From d462a7e982ab58eb865736559eebd587f75baaa2 Mon Sep 17 00:00:00 2001
From: Ryan Ziegler <rzig408@gmail.com>
Date: Sun, 27 Apr 2025 17:24:07 -0400
Subject: [PATCH] create reassociate pass that does nothing

---
 hercules_opt/src/lib.rs                    |  2 ++
 hercules_opt/src/reassociate.rs            | 37 ++++++++++++++++++++++
 juno_samples/grape_reduction/src/grape.sch |  9 ++++--
 juno_scheduler/src/compile.rs              |  1 +
 juno_scheduler/src/ir.rs                   |  1 +
 juno_scheduler/src/pm.rs                   | 12 +++++++
 6 files changed, 59 insertions(+), 3 deletions(-)
 create mode 100644 hercules_opt/src/reassociate.rs

diff --git a/hercules_opt/src/lib.rs b/hercules_opt/src/lib.rs
index b25449e7..e9550cd4 100644
--- a/hercules_opt/src/lib.rs
+++ b/hercules_opt/src/lib.rs
@@ -20,6 +20,7 @@ pub mod loop_bound_canon;
 pub mod outline;
 pub mod phi_elim;
 pub mod pred;
+pub mod reassociate;
 pub mod reuse_products;
 pub mod rewrite_math_expressions;
 pub mod schedule;
@@ -49,6 +50,7 @@ pub use crate::loop_bound_canon::*;
 pub use crate::outline::*;
 pub use crate::phi_elim::*;
 pub use crate::pred::*;
+pub use crate::reassociate::*;
 pub use crate::reuse_products::*;
 pub use crate::rewrite_math_expressions::*;
 pub use crate::schedule::*;
diff --git a/hercules_opt/src/reassociate.rs b/hercules_opt/src/reassociate.rs
new file mode 100644
index 00000000..f876c55a
--- /dev/null
+++ b/hercules_opt/src/reassociate.rs
@@ -0,0 +1,37 @@
+use hercules_ir::def_use::*;
+use hercules_ir::ir::*;
+
+use crate::*;
+
+/*
+ * Top level function to run dead code elimination.
+ */
+pub fn reassociate(editor: &mut FunctionEditor) {
+    panic!("Reassociate was called");
+    // Create worklist (starts as all nodes).
+    let mut worklist: Vec<NodeID> = editor.node_ids().collect();
+
+    while let Some(work) = worklist.pop() {
+        // If a node on the worklist is a start node, it is either *the* start
+        // node (which we shouldn't delete), or is a gravestone for an already
+        // deleted node earlier in the worklist. If a node is a return node, it
+        // shouldn't be removed.
+        if editor.func().nodes[work.idx()].is_start() || editor.func().nodes[work.idx()].is_return()
+        {
+            continue;
+        }
+
+        // If a node on the worklist has 0 users, delete it. Add its uses onto
+        // the worklist.
+        if editor.get_users(work).len() == 0 {
+            let uses = get_uses(&editor.func().nodes[work.idx()]);
+            let success = editor.edit(|edit| edit.delete_node(work));
+            if success {
+                // If the edit was performed, then its uses may now be dead.
+                for u in uses.as_ref() {
+                    worklist.push(*u);
+                }
+            }
+        }
+    }
+}
diff --git a/juno_samples/grape_reduction/src/grape.sch b/juno_samples/grape_reduction/src/grape.sch
index 2f4aba67..3be75117 100644
--- a/juno_samples/grape_reduction/src/grape.sch
+++ b/juno_samples/grape_reduction/src/grape.sch
@@ -10,10 +10,10 @@ delete-uncalled(*);
 forkify(*);
 fork-guard-elim(*);
 dce(*);
-xdot[true](*);
+// xdot[true](*);
 // fork-tile[8, 0, false, true](*);
 let a = fork-split(*);
-xdot[true](*);
+// xdot[true](*);
 
 
 fixpoint stop after 10 {
@@ -28,12 +28,15 @@ fixpoint stop after 10 {
   dce(entry@inner);
   lift-dc-math(entry@inner);
 }
-xdot[true](*);
+// xdot[true](*);
 
 a2p(*);
 sroa(*);
 xdot[true](*);
 
+// reassociate go brr
+reassociate(*);
+
 gvn(*);
 phi-elim(*);
 ccp(*);
diff --git a/juno_scheduler/src/compile.rs b/juno_scheduler/src/compile.rs
index 361f08a4..2c941b79 100644
--- a/juno_scheduler/src/compile.rs
+++ b/juno_scheduler/src/compile.rs
@@ -152,6 +152,7 @@ impl FromStr for Appliable {
             "outline" => Ok(Appliable::Pass(ir::Pass::Outline)),
             "phi-elim" => Ok(Appliable::Pass(ir::Pass::PhiElim)),
             "predication" => Ok(Appliable::Pass(ir::Pass::Predication)),
+            "reassociate" => Ok(Appliable::Pass(ir::Pass::Reassociate)),
             "reduce-slf" => Ok(Appliable::Pass(ir::Pass::ReduceSLF)),
             "rename" => Ok(Appliable::Pass(ir::Pass::Rename)),
             "reuse-products" => Ok(Appliable::Pass(ir::Pass::ReuseProducts)),
diff --git a/juno_scheduler/src/ir.rs b/juno_scheduler/src/ir.rs
index aa5926cb..bee0cf06 100644
--- a/juno_scheduler/src/ir.rs
+++ b/juno_scheduler/src/ir.rs
@@ -37,6 +37,7 @@ pub enum Pass {
     PhiElim,
     Predication,
     Print,
+    Reassociate,
     ReduceSLF,
     Rename,
     ReuseProducts,
diff --git a/juno_scheduler/src/pm.rs b/juno_scheduler/src/pm.rs
index 20413057..3627c0a0 100644
--- a/juno_scheduler/src/pm.rs
+++ b/juno_scheduler/src/pm.rs
@@ -2572,6 +2572,18 @@ fn run_pass(
             pm.delete_gravestones();
             pm.clear_analyses();
         }
+        Pass::Reassociate => {
+            assert!(args.is_empty());
+            for func in build_selection(pm, selection, false) {
+                let Some(mut func) = func else {
+                    continue;
+                };
+                reassociate(&mut func);
+                changed |= func.modified();
+            }
+            pm.delete_gravestones();
+            pm.clear_analyses();
+        }
         Pass::ReduceSLF => {
             assert!(args.is_empty());
             pm.make_fork_join_maps();
-- 
GitLab