Fork reshape
Implements a fork-reshape
pass in the scheduling language which can be used to coalesce a selection of fork-joins, interchange them into a specified order, and then split them into a specified series of chunks.
Adds a parallel and tiled schedule for matmul that uses fork-reshape
to keep the schedule quite clean.
Also adds tuples to the scheduling language and a new form of let-statements that allow destructing tuples, it's not the nicest (it's not a general pattern handler and will only cause an error if there are too many fields on the left-hand side, not too few, but it's enough to make the code that uses fork-reshape
cleaner).
Includes a few other tweaks I ran into
- When adding a node in the editor, only insert an empty def_use set if the node doesn't have one yet (this is important due to cycles between nodes, and was causing issues in fork reshape since it invokes a series of passes using the same editor)
- Made fork-coalesce return the fork and join nodes when it succeeds so that they can be used in further operations.
- Modified fork coalesce's check for control flow between the forks and joins to check via uses rather than users because I was having weird issues with checking users
- Modified fork coalesce to perform everything as a single edit because I was having issues when manually scheduling matmul when using selections with coalesce.
- Modified fork split to return the ID of the new fork on a success
- Modified both fork split and coalesce to propagate parallel fork schedules when appropriate to avoid having to re-assign/re-infer schedules
Merge request reports
Activity
requested review from @rarbore2
assigned to @aaronjc4
added 7 commits
-
e95929fd...b2231f2d - 6 commits from branch
main
- 43c5f06c - Merge branch 'main' into fork-reshape
-
e95929fd...b2231f2d - 6 commits from branch
mentioned in commit cc694f57