rewrite rule for unrolling scf.ifelse (#244)

also adjust the aggressive fold rule order.

Author: Roger-luo

src/kirin/dialects/scf/constprop.py

@@ -41,6 +41,10 @@ def if_else(
                 interp_, frame, stmt, cond, body
             )
             frame.entries.update(body_frame.entries)
+            if not body_frame.frame_is_not_pure and not isinstance(
+                body.blocks[0].last_stmt, func.Return
+            ):
+                frame.should_be_pure.add(stmt)
             return ret
         else:
             then_frame, then_results = self._prop_const_cond_ifelse(
@@ -80,11 +84,6 @@ def _prop_const_cond_ifelse(
             body_frame.entries.update(frame.entries)
             body_frame.set(body.blocks[0].args[0], cond)
             results = interp_.run_ssacfg_region(body_frame, body)
-
-        if not body_frame.frame_is_not_pure and not isinstance(
-            body.blocks[0].last_stmt, func.Return
-        ):
-            frame.should_be_pure.add(stmt)
         return body_frame, results
 
     @interp.impl(For)

src/kirin/dialects/scf/stmts.py

@@ -87,6 +87,9 @@ def print_impl(self, printer: Printer) -> None:
             printer.plain_print(" else ", style="keyword")
             printer.print(self.else_body)
 
+        with printer.rich(style="comment"):
+            printer.plain_print(f" -> purity={self.purity}")
+
     def verify(self) -> None:
         from kirin.dialects.func import Return
 
@@ -225,6 +228,8 @@ def print_impl(self, printer: Printer) -> None:
                     printer.print_stmt(stmt)
         printer.print_newline()
         printer.plain_print("}")
+        with printer.rich(style="comment"):
+            printer.plain_print(f" -> purity={self.purity}")
 
 
 @statement(dialect=dialect)

src/kirin/dialects/scf/unroll.py

@@ -1,10 +1,56 @@
 from kirin import ir
 from kirin.analysis import const
+from kirin.dialects import func
 from kirin.rewrite.abc import RewriteRule
 from kirin.rewrite.result import RewriteResult
 from kirin.dialects.py.constant import Constant
 
-from .stmts import For, Yield
+from .stmts import For, Yield, IfElse
+
+
+class PickIfElse(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, IfElse):
+            return RewriteResult()
+
+        if not isinstance(hint := node.cond.hints.get("const"), const.Value):
+            return RewriteResult()
+
+        if hint.data:
+            return self.insert_body(node, node.then_body)
+        else:
+            return self.insert_body(node, node.else_body)
+
+    def insert_body(self, node: IfElse, body: ir.Region):
+        body_block = body.clone().blocks[0]
+        body_block.args[0].replace_by(node.cond)
+        block_stmt = body_block.first_stmt
+        while block_stmt and not block_stmt.has_trait(ir.IsTerminator):
+            block_stmt.detach()
+            block_stmt.insert_before(node)
+            block_stmt = body_block.first_stmt
+
+        terminator = body_block.last_stmt
+        if isinstance(terminator, Yield):
+            for result, output in zip(node.results, terminator.values):
+                result.replace_by(output)
+            node.delete()
+            return RewriteResult(has_done_something=True)
+        elif isinstance(terminator, func.Return):
+            block = node.parent
+            assert block is not None
+            stmt = block.last_stmt
+            while stmt is not None and stmt is not node:  # remove the rest of the block
+                delete_stmt = stmt
+                stmt = stmt.prev_stmt
+                delete_stmt.delete()
+
+            terminator.detach()
+            node.replace_by(terminator)
+            return RewriteResult(has_done_something=True)
+        else:
+            return RewriteResult()
 
 
 class ForLoop(RewriteRule):

src/kirin/ir/nodes/region.py

@@ -152,7 +152,7 @@ def clone(self, ssamap: dict[SSAValue, SSAValue] | None = None) -> Region:
             for stmt in block.stmts:
                 new_stmt = stmt.from_stmt(
                     stmt,
-                    args=[_ssamap[arg] for arg in stmt.args],
+                    args=[_ssamap.get(arg, arg) for arg in stmt.args],
                     regions=[region.clone(_ssamap) for region in stmt.regions],
                     successors=[
                         successor_map[successor] for successor in stmt.successors

src/kirin/passes/aggressive/fold.py

@@ -1,16 +1,43 @@
-from dataclasses import dataclass
+from dataclasses import field, dataclass
 
 from kirin.passes import Pass
-from kirin.rewrite import aggressive
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Inline,
+    Fixpoint,
+    WrapConst,
+    Call2Invoke,
+    ConstantFold,
+    CFGCompactify,
+    InlineGetItem,
+    InlineGetField,
+    DeadCodeElimination,
+)
 from kirin.analysis import const
 from kirin.ir.method import Method
 from kirin.rewrite.abc import RewriteResult
 
 
 @dataclass
 class Fold(Pass):
+    constprop: const.Propagate = field(init=False)
+
+    def __post_init__(self):
+        self.constprop = const.Propagate(self.dialects)
 
     def unsafe_run(self, mt: Method) -> RewriteResult:
-        constprop = const.Propagate(self.dialects)
-        constprop_results, _ = constprop.run_analysis(mt)
-        return aggressive.Fold(constprop_results).rewrite(mt.code)
+        result = RewriteResult()
+        frame, _ = self.constprop.run_analysis(mt)
+        result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
+        rule = Chain(
+            ConstantFold(),
+            Call2Invoke(),
+            InlineGetField(),
+            InlineGetItem(),
+            DeadCodeElimination(),
+        )
+        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+        result = Walk(Inline(lambda _: True)).rewrite(mt.code).join(result)
+        result = Fixpoint(CFGCompactify()).rewrite(mt.code).join(result)
+        return result

src/kirin/print/printer.py

@@ -145,10 +145,10 @@ def print_stmt(self, node: "ir.Statement"):
     def print_hint(
         self,
         *values: "ir.SSAValue",
-        prefix: str = "//hint<",
+        prefix: str = " // hint<",
         suffix: str = ">",
     ):
-        if not self.hint:
+        if not self.hint or not values:
             return
 
         self.plain_print(prefix)
@@ -157,7 +157,8 @@ def print_hint(
             if idx > 0:
                 self.plain_print(", ")
 
-            if item.hints.get(self.hint):
+            self.plain_print("=")
+            if item.hints.get(self.hint) is not None:
                 self.plain_print(repr(item.hints.get(self.hint)))
             else:
                 self.plain_print("missing")