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

juno_samples/matmul/src/cpu.sch

@@ -10,27 +10,45 @@ macro optimize!(X) {
   dce(X);
 }
 
-macro codegen!(X) {
-  gcm(*);
-  float-collections(*);
-  dce(*);
-  gcm(*);
+macro codegen-prep!(X) {
+  optimize!(X);
+  gcm(X);
+  float-collections(X);
+  dce(X);
+  gcm(X);
 }
 
-optimize!(*);
+macro forkify!(X) {
+  fixpoint {
+    forkify(X);
+    fork-guard-elim(X);
+  }
+}
+
+macro fork-tile![n](X) {
+  fork-tile[n, 0, false, true](X);
+}
+
+macro parallelize!(X) {
+  parallel-fork(X);
+  parallel-reduce(X);
+}
 
-fixpoint panic after 20 {
-  forkify(matmul);
-  fork-guard-elim(matmul);
+macro unforkify!(X) {
+  fork-split(X);
+  unforkify(X);
 }
 
+optimize!(*);
+forkify!(*);
+
 // Mark the whole loop nest as associative, any order of iterations is equivalent
 associative(matmul@outer);
 
 // Tile the outer 2 loops to create 16 parallel threads (each responsible for
 // computing one block of the output
 let par = matmul@outer \ matmul@inner;
-fork-tile[4, 0, false, true](par);
+fork-tile![4](par);
 fork-coalesce(par);
 fork-interchange[0, 1](par);
 fork-interchange[2, 3](par);
@@ -38,29 +56,25 @@ fork-interchange[1, 2](par);
 
 let split = fork-split(*);
 fork-coalesce(split.matmul.fj0 \ split.matmul.fj2);
-parallel-fork(split.matmul.fj0 \ split.matmul.fj2);
+
+parallelize!(split.matmul.fj0 \ split.matmul.fj2);
 
 // Pull the body of the parallel loop out into its own device function
 let body = outline(split.matmul.fj2);
 cpu(body);
 
 // Tile the loop nest for cache performance; 16x16x16 tile
-fork-tile[16, 0, false, true](body);
+fork-tile![16](body);
 fixpoint { fork-coalesce(body); }
-
 fork-interchange[1, 2](body);
 fork-interchange[3, 4](body);
 fork-interchange[2, 3](body);
 
-optimize!(*);
-
 fork-split(body);
 reduce-slf(*);
-unforkify(body);
-
-optimize!(*);
 
-parallel-reduce(split.matmul.fj0);
 xdot[true](*);
 
-codegen!(*);
+unforkify!(body);
+
+codegen-prep!(*);