Proper state edge passthrough

egraphs-good · Feb 7, 2025 · c0c4951 · c0c4951
1 parent 83a6209
commit c0c4951
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/dag_in_context/src/optimizations/hackers_delight.egg b/dag_in_context/src/optimizations/hackers_delight.egg
@@ -58,14 +58,16 @@
     (union (Get outerif i) (Bop (Div) n lowbitn))
 ) :ruleset hacker)
 
-;; Try to get rid of the state_edge hanging there
-;; TODO: This is obviously cheating
+;; Try to do a state-edge-passthrough for loops
+;; PopcountIterations guarantees termination for non-zero values
+;; lowbit(0) is undefined behavior
 
 (rule (
     (PopcountIterations anyif n i)
     (= anyif (If cond inputs thenbr elsebr))
-    (HasType (Get anyif j) (Base (StateT)))
-    (HasType (Get inputs k) (Base (StateT)))
+    (= thenbr (DoWhile lpinputs pred_outputs))
+    (= (Get pred_outputs (+ j 1)) (Get (Arg arg_ty then_ctx) j))
+    (HasType (Get pred_outputs (+ j 1)) (Base (StateT)))
 ) (
-    (union (Get anyif j) (Get inputs k))
+    (union (Get thenbr j) (Get lpinputs j))
 ) :ruleset hacker)
diff --git a/dag_in_context/src/schedule.rs b/dag_in_context/src/schedule.rs
@@ -196,7 +196,8 @@ pub fn parallel_schedule() -> Vec<CompilerPass> {
 (run-schedule
     (saturate
       {helpers}
-      passthrough)
+      passthrough
+      state-edge-passthrough)
     (repeat 2
         {helpers}
         all-optimizations
@@ -207,7 +208,10 @@ pub fn parallel_schedule() -> Vec<CompilerPass> {
         cheap-optimizations
     )
 
-    {helpers}
+    (saturate
+      {helpers}
+      passthrough
+      state-edge-passthrough)
 )
 "
         )),