From 25db4a50102e2629b39e350eb2fdad2460385bee Mon Sep 17 00:00:00 2001
From: Aaron Councilman <aaronjc4@illinois.edu>
Date: Wed, 5 Mar 2025 14:49:33 -0600
Subject: [PATCH] Fix backprop and matmul cpu schedules

---
 juno_samples/matmul/src/matmul.sch        | 8 ++++++--
 juno_samples/rodinia/backprop/src/cpu.sch | 8 ++++----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/juno_samples/matmul/src/matmul.sch b/juno_samples/matmul/src/matmul.sch
index 6867576e..e6fe7bc8 100644
--- a/juno_samples/matmul/src/matmul.sch
+++ b/juno_samples/matmul/src/matmul.sch
@@ -25,6 +25,10 @@ macro forkify!(X) {
   }
 }
 
+macro fork-chunk![n](X) {
+  fork-tile[n, 0, false, false](X);
+}
+
 macro fork-tile![n](X) {
   fork-tile[n, 0, false, true](X);
 }
@@ -66,8 +70,8 @@ if feature("cuda") {
 
   // Parallelize by computing output array as 16 chunks
   let par = matmul@outer \ matmul@inner;
-  fork-tile![4](par);
-  let (outer, inner, _) = fork-reshape[[1, 3], [0], [2]](par);
+  fork-chunk![4](par);
+  let (outer, inner, _) = fork-reshape[[0, 2], [1], [3]](par);
   parallelize!(outer \ inner);
 
   let body = outline(inner);
diff --git a/juno_samples/rodinia/backprop/src/cpu.sch b/juno_samples/rodinia/backprop/src/cpu.sch
index 4796f427..cbdd90fd 100644
--- a/juno_samples/rodinia/backprop/src/cpu.sch
+++ b/juno_samples/rodinia/backprop/src/cpu.sch
@@ -38,8 +38,8 @@ let forward_input = outline(backprop@forward_input);
 let forward_hidden = outline(backprop@forward_hidden);
 
 if !feature("seq") {
-  fork-tile[16, 0, false, true](forward_input@outer_loop \ forward_input@inner_loop);
-  let (outer, inner) = fork-reshape[[1], [0]](forward_input@outer_loop \ forward_input@inner_loop);
+  fork-tile[16, 0, false, false](forward_input@outer_loop \ forward_input@inner_loop);
+  let (outer, inner) = fork-reshape[[0], [1]](forward_input@outer_loop \ forward_input@inner_loop);
   forward_input = outline(inner);
   inline(backprop@forward_input);
 }
@@ -53,8 +53,8 @@ let adjust_hidden = outline(backprop@adjust_hidden);
 let adjust_input = outline(backprop@adjust_input);
 
 if !feature("seq") {
-  fork-tile[16, 0, false, true](adjust_input);
-  let (outer, inner) = fork-reshape[[1], [0, 2]](adjust_input);
+  fork-tile[16, 0, false, false](adjust_input);
+  let (outer, inner) = fork-reshape[[0], [1, 2]](adjust_input);
   adjust_input = outline(inner);
   inline(backprop@adjust_input);
 }
-- 
GitLab