backport 493 (#494)

Author: weinbe58

src/kirin/dialects/scf/unroll.py

@@ -88,4 +88,4 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         for result, output in zip(node.results, loop_vars):
             result.replace_by(output)
         node.delete()
-        return RewriteResult(has_done_something=True, terminated=True)
+        return RewriteResult(has_done_something=True)

src/kirin/passes/aggressive/__init__.py

@@ -1 +1,2 @@
 from .fold import Fold as Fold
+from .unroll import UnrollScf as UnrollScf

src/kirin/passes/aggressive/unroll.py

@@ -0,0 +1,35 @@
+from dataclasses import field, dataclass
+
+from kirin.ir import Method
+from kirin.passes import Fold, Pass, TypeInfer
+from kirin.rewrite import Walk
+from kirin.rewrite.abc import RewriteResult
+from kirin.dialects.scf.unroll import ForLoop, PickIfElse
+
+
+@dataclass
+class UnrollScf(Pass):
+    """This pass can be used to unroll scf.For loops and inline/expand scf.IfElse when
+    the input are known at compile time.
+
+    usage:
+        UnrollScf(dialects).fixpoint(method)
+
+    Note: This pass should be used in a fixpoint manner, to unroll nested scf nodes.
+
+    """
+
+    typeinfer: TypeInfer = field(init=False)
+    fold: Fold = field(init=False)
+
+    def __post_init__(self):
+        self.typeinfer = TypeInfer(self.dialects, no_raise=self.no_raise)
+        self.fold = Fold(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: Method):
+        result = RewriteResult()
+        result = Walk(PickIfElse()).rewrite(mt.code).join(result)
+        result = Walk(ForLoop()).rewrite(mt.code).join(result)
+        result = self.typeinfer(mt).join(result)
+        result = self.fold(mt).join(result)
+        return result

test/passes/test_unroll_scf.py

@@ -0,0 +1,39 @@
+from kirin.prelude import structural
+from kirin.dialects import py, func
+from kirin.passes.aggressive import UnrollScf
+
+
+def test_unroll_scf():
+    @structural
+    def main(r: list[int], cond: bool):
+        if cond:
+            for i in range(4):
+                tmp = r[-1]
+                if i < 2:
+                    tmp += i * 2
+                else:
+                    for j in range(4):
+                        if i > j:
+                            tmp += i + j
+                        else:
+                            tmp += i - j
+
+                r.append(tmp)
+        else:
+            for i in range(4):
+                r.append(i)
+        return r
+
+    UnrollScf(structural).fixpoint(main)
+
+    num_adds = 0
+    num_calls = 0
+
+    for op in main.callable_region.walk():
+        if isinstance(op, py.Add):
+            num_adds += 1
+        elif isinstance(op, func.Call):
+            num_calls += 1
+
+    assert num_adds == 10
+    assert num_calls == 8