Fix nested for-loop handling in squin -> stim pass (#394)

Previously, the following would cause the squin to stim pass to fail:

```python
for i in range(...):
  for j in range(...):
    qubit.apply(op.cx(), [q[i], q[j])
```

This was because the unroll rule was not successful in fully unrolling
the loops, getting (at most) one unroll before silently giving up. This
behavior was first identified by @liupengy19 and encountered more
recently by @ehua7365 as well.

@kaihsin pointed out to me that the problem might be because of a lack
of constprop occuring between each unroll. This turns out to be the case
and I've implemented a mini-pass that will do the constprop after each
unroll, with the pass itself applied via fixpoint instead of the more
traditional unsafe_run, single application style.

I've also taken the liberty to rename some of the tests to be a bit more
accurate and removed some outdated comments.

---------

Co-authored-by: David Plankensteiner <[email protected]>

Author: johnzl-777

src/bloqade/stim/passes/squin_to_stim.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 
-from kirin.passes import Fold
+from kirin.passes import Fold, HintConst
 from kirin.rewrite import (
     Walk,
     Chain,
@@ -38,6 +38,37 @@
 from ..rewrite.ifs_to_stim import IfToStim
 
 
+@dataclass
+class AggressiveForLoopUnroll(Pass):
+    """
+    Aggressive unrolling of for loops, addresses cases where unroll
+    does not successfully handle nested loops because of a lack of constprop.
+
+    This should be invoked via fixpoint to let this be repeatedly applied until
+    no further rewrites are possible.
+    """
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        rule = Chain(
+            InlineGetField(),
+            InlineGetItem(),
+            scf.unroll.ForLoop(),
+            scf.trim.UnusedYield(),
+        )
+
+        # Intentionally only walk ONCE, let fixpoint happen with the WHOLE pass
+        # so that HintConst gets run right after, allowing subsequent unrolls to happen
+        rewrite_result = Walk(rule).rewrite(mt.code)
+
+        rewrite_result = (
+            HintConst(dialects=mt.dialects, no_raise=self.no_raise)
+            .unsafe_run(mt)
+            .join(rewrite_result)
+        )
+
+        return rewrite_result
+
+
 @dataclass
 class SquinToStimPass(Pass):
 
@@ -48,15 +79,10 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             dialects=mt.dialects, no_raise=self.no_raise
         ).unsafe_run(mt)
 
-        rule = Chain(
-            InlineGetField(),
-            InlineGetItem(),
-            scf.unroll.ForLoop(),
-            scf.trim.UnusedYield(),
-        )
-        rewrite_result = Fixpoint(Walk(rule)).rewrite(mt.code).join(rewrite_result)
-        # fold_pass = Fold(mt.dialects, no_raise=self.no_raise)
-        # rewrite_result = fold_pass(mt)
+        rewrite_result = AggressiveForLoopUnroll(
+            dialects=mt.dialects, no_raise=self.no_raise
+        ).fixpoint(mt)
+
         rewrite_result = (
             Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(rewrite_result)
         )
@@ -66,9 +92,6 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             .join(rewrite_result)
         )
 
-        # run typeinfer again after unroll etc. because we now insert
-        # a lot of new nodes, which might have more precise types
-        # self.typeinfer.unsafe_run(mt)
         rewrite_result = (
             Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
             .rewrite(mt.code)

test/stim/passes/stim_reference_programs/qubit/for_loop.stim renamed to test/stim/passes/stim_reference_programs/qubit/for_loop_nontrivial_index.stim

@@ -0,0 +1,4 @@
+
+H 0
+CX 0 2
+CX 1 2

test/stim/passes/stim_reference_programs/qubit/nested_for_loop.stim

@@ -153,7 +153,7 @@ def test():
     assert codegen(test).strip() == "SQRT_Y 0"
 
 
-def test_for_loop_rewrite():
+def test_for_loop_nontrivial_index_rewrite():
 
     @squin.kernel
     def main():
@@ -164,6 +164,23 @@ def main():
             squin.qubit.apply(cx, [q[i], q[i + 1]])
 
     SquinToStimPass(main.dialects)(main)
-    base_stim_prog = load_reference_program("for_loop.stim")
+    base_stim_prog = load_reference_program("for_loop_nontrivial_index.stim")
+
+    assert codegen(main) == base_stim_prog.rstrip()
+
+
+def test_nested_for_loop_rewrite():
+
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(5)
+        squin.qubit.apply(squin.op.h(), q[0])
+        cx = squin.op.cx()
+        for i in range(2):
+            for j in range(2, 3):
+                squin.qubit.apply(cx, [q[i], q[j]])
+
+    SquinToStimPass(main.dialects)(main)
+    base_stim_prog = load_reference_program("nested_for_loop.stim")
 
     assert codegen(main) == base_stim_prog.rstrip()