Implement User Scheduling in Python Frontend

[rdspring1]

Why are we doing this?

Facilitate rapid prototyping of new schedulers.

• The development of new schedulers is relatively bug-prone, so the cpp compile time latency is cumbersome.
• pytestframework allows creating parameter space for a single test. A developer can iterate and benchmark several configurations while developing heuristics for a scheduler.

Open Questions:

How should we refactor the CPP scheduling API to improve usability?

• @kevinstephano has suggested splitting the CPP ParallelType enum into Parallelization and LoopTransform enums in python frontend.

Error-checking?

• Error checking is already handled in CPP, but reimplementing simple checks could improve developer quality of life.

Experimental Scheduling API?

Some aspects of NvFuser’s scheduling API are experimental. Backward and forward compatibility for experiment functions would not be guaranteed. Currently, an experimental API is denoted with an underscore appended to the front of the function's name.

Which scheduling operations to implement?

TensorView

• [x] split
• [x] merge
• [x] reorder
• [x] rFactor
• [x] cacheAfter
• [x] cacheBefore
• [x] setMemoryType

Scheduling Utilities

• [x] propagateTransformation
• [x] parallelizeAllLike
• [x] inlineSelectedAt
• [x] inlineMost

Filtering Utilities

• [x] getReductionTVs
• [x] allTvsExcept

Enum

• [x] LoadStoreOpType
• [x] CacheOp
• [x] MemoryType
• [x] ParallelType --- DIDx, BIDx, BIDy, BIDz, TIDx, TIDy, TIDz
• [x] LoopTransform --- Vectorize, Unroll, Unswitch, Bulk, Serial

Debug

• [x] Print TensorView
• [x] Print fusion ir for user-scheduled fusions

[rdspring1]

How to print transformed state of TensorView or Fusion in python frontend for debugging?

[rdspring1]

Which functions to primitives implement?

TensorView

• [ ] cacheFork
• [ ] doubleBuffer
• [ ] swizzle

Filtering Utilities

• [ ] ir_utils::getViewOps
• [ ] DomainMap::findReferenceFor
• [ ] DomainMap::groupInputsOutputsByInnerDim
• [ ] Functions for DependencyCheck and StmtSort

Scheduling Utilities

• [ ] scheduler_utils::propagateReshapeTransforms