scf.for loop unroll (#237)

this PR implements loop unroll for `scf.For`, note that the loop body
here asssumes no branching inside thus the rewrite is just copy and
replacing SSAValues from the previous expanded body.

Author: Roger-luo

src/kirin/dialects/scf/__init__.py

@@ -12,6 +12,7 @@
 
 from . import (
     interp as interp,
+    unroll as unroll,
     lowering as lowering,
     constprop as constprop,
     typeinfer as typeinfer,

src/kirin/dialects/scf/unroll.py

@@ -0,0 +1,45 @@
+from kirin import ir
+from kirin.analysis import const
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+from kirin.dialects.py.constant import Constant
+
+from .stmts import For, Yield
+
+
+class ForLoop(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, For):
+            return RewriteResult()
+
+        # TODO: support for PartialTuple and IList with known length
+        if not isinstance(hint := node.iterable.hints.get("const"), const.Value):
+            return RewriteResult()
+
+        loop_vars = node.initializers
+        for item in hint.data:
+            body = node.body.clone()
+            block = body.blocks[0]
+            item_stmt = Constant(item)
+            item_stmt.insert_before(node)
+            block.args[0].replace_by(item_stmt.result)
+            for var, input in zip(block.args[1:], loop_vars):
+                var.replace_by(input)
+
+            block_stmt = block.first_stmt
+            while block_stmt and not block_stmt.has_trait(ir.IsTerminator):
+                block_stmt.detach()
+                block_stmt.insert_before(node)
+                block_stmt = block.first_stmt
+
+            terminator = block.last_stmt
+            # we assume Yield has the same # of values as initializers
+            # TODO: check this in validation
+            if isinstance(terminator, Yield):
+                loop_vars = terminator.values
+
+        for result, output in zip(node.results, loop_vars):
+            result.replace_by(output)
+        node.delete()
+        return RewriteResult(has_done_something=True)

test/dialects/scf/test_unroll.py

@@ -0,0 +1,27 @@
+from kirin.passes import Fold
+from kirin.prelude import structural_no_opt
+from kirin.rewrite import Walk
+from kirin.dialects import py, scf, func
+
+
+def test_simple_loop_unroll():
+    @structural_no_opt
+    def simple_loop(x):
+        for i in range(3):
+            x = x + i
+        return x
+
+    fold = Fold(structural_no_opt)
+    fold(simple_loop)
+    Walk(scf.unroll.ForLoop()).rewrite(simple_loop.code)
+    assert len(simple_loop.callable_region.blocks) == 1
+    stmts = simple_loop.callable_region.blocks[0].stmts
+    assert isinstance(stmts.at(0), py.Constant)
+    assert isinstance(stmts.at(1), py.Constant)
+    assert isinstance(stmts.at(2), py.Add)
+    assert isinstance(stmts.at(3), py.Constant)
+    assert isinstance(stmts.at(4), py.Add)
+    assert isinstance(stmts.at(5), py.Constant)
+    assert isinstance(stmts.at(6), py.Add)
+    assert isinstance(stmts.at(7), func.Return)
+    assert simple_loop(1) == 4