pytorch
859ac58c - [Inductor] Support loop split at given depth in CPP codegen (#91397)

Commit
1 year ago
[Inductor] Support loop split at given depth in CPP codegen (#91397) This PR refactors the loop related data structure to support the loop split at a given depth. Before this PR, the loop split is always supported at the inner-most level. With this PR, it is possible to support tiling at outer levels and at more than one levels. The `LoopNest` data structure is extended to support loop splits at various levels and renamed to `LoopNestWithSplit`. The `codegen_loops` function is also rewritten to be general to support arbitrary kernels set at the leaves of the loop structure. This PR also improves the handling of reduction loops with split. The main loop and tail loop now work on their own reduction variables in parallel without data dependency as previous do. With this, two workarounds can be removed in the `CppVecKernel`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/91397 Approved by: https://github.com/EikanWang, https://github.com/jansel
Author
Committer
Parents
Loading