Skip to content
Snippets Groups Projects

Fork reshape

Merged Aaron Councilman requested to merge fork-reshape into main

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

Pipeline #201934 passed

Pipeline passed for 7370cb51 on fork-reshape

Merged by rarbore2rarbore2 2 weeks ago (Feb 28, 2025 6:50pm UTC)

Merge details

  • Changes merged into main with cc694f57 (commits were squashed).
  • Deleted the source branch.

Pipeline #201935 passed

Pipeline passed for cc694f57 on main

Activity

Filter activity
  • Approvals
  • Assignees & reviewers
  • Comments (from bots)
  • Comments (from users)
  • Commits & branches
  • Edits
  • Labels
  • Lock status
  • Mentions
  • Merge request status
  • Tracking
Please register or sign in to reply
Loading