Refactor Scf to Cf rewrite rules. (#467)

The first pass on this rule was a bit clunky and was trying to optimize
the rewrite. I have refactored the rule now to be less optimal but more
readable.

Author: weinbe58

src/kirin/dialects/scf/scf2cf.py

@@ -1,190 +1,181 @@
-from ... import ir
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from ..cf import Branch, ConditionalBranch
+from ..func import ConstantNone
 from .stmts import For, Yield, IfElse
-from ...rewrite.abc import RewriteRule, RewriteResult
+from ..py.cmp import Is
+from ..py.iterable import Iter, Next
 
 
-class ScfToCfRule(RewriteRule):
+class ScfRule(RewriteRule):
 
-    def rewrite_ifelse(
-        self, node: ir.Region, block_idx: int, curr_block: ir.Block, stmt: IfElse
-    ):
-        from kirin.dialects import cf
+    def get_entr_and_exit_blks(self, node: For | IfElse):
+        """Get the enter and exit blocks for the given SCF node.
 
-        # create a new block for entering the if statement
-        entry_block = ir.Block()
-        for arg in curr_block.args:
-            arg.replace_by(entry_block.args.append_from(arg.type, arg.name))
+        The exit block is a new block that will be created to hold the
+        statements that follow the SCF node in the current block and the
+        enter block is a new block that will be created to hold the
+        any logic required to enter the SCF node.
 
-        # delete the args of the old block and replace with the result of the # if statement
-        for arg in curr_block.args:
-            curr_block.args.delete(arg)
+        """
+        # split the current block into two parts
+        exit_block = ir.Block()
+        stmt = node.next_stmt
+        while stmt is not None:
+            next_stmt = stmt.next_stmt
+            stmt.detach()
+            exit_block.stmts.append(stmt)
+            stmt = next_stmt
 
-        for arg in stmt.results:
-            arg.replace_by(curr_block.args.append_from(arg.type, arg.name))
+        for result in node.results:
+            result.replace_by(exit_block.args.append_from(result.type, result.name))
 
-        (then_block := stmt.then_body.blocks[0]).detach()
-        (else_block := stmt.else_body.blocks[0]).detach()
+        curr_block = node.parent_block
+        assert isinstance(curr_block, ir.Block), "Node must be inside a block"
 
-        entry_block.stmts.append(
-            cf.ConditionalBranch(
-                cond=stmt.cond,
-                then_arguments=tuple(stmt.args),
-                then_successor=then_block,
-                else_arguments=tuple(stmt.args),
-                else_successor=else_block,
-            )
+        curr_block.stmts.append(
+            Branch(arguments=(), successor=(entr_block := ir.Block()))
         )
 
-        # insert the then/else blocks and add branch to the current block
-        # if the last statement of the then block is a yield
-        if isinstance(last_stmt := else_block.last_stmt, Yield):
-            last_stmt.replace_by(
-                cf.Branch(
-                    arguments=tuple(last_stmt.args),
-                    successor=curr_block,
-                )
-            )
+        return exit_block, entr_block
 
-        if isinstance(last_stmt := then_block.last_stmt, Yield):
-            last_stmt.replace_by(
-                cf.Branch(
-                    arguments=tuple(last_stmt.args),
-                    successor=curr_block,
-                )
-            )
+    def get_curr_blk_info(self, node: For | IfElse) -> tuple[ir.Region, int]:
+        """Get the current region and the block index of the node in the region."""
+        curr_block = node.parent_block
+        region = node.parent_region
 
-        node.blocks.insert(block_idx, curr_block)
-        node.blocks.insert(block_idx, else_block)
-        node.blocks.insert(block_idx, then_block)
+        assert isinstance(region, ir.Region), "Node must be inside a region"
+        assert isinstance(curr_block, ir.Block), "Node must be inside a block"
 
-        curr_stmt = stmt
-        next_stmt = stmt.prev_stmt
-        curr_stmt.delete()
+        block_idx = region._block_idx[curr_block]
+        return region, block_idx
 
-        return next_stmt, entry_block
 
-    def rewrite_for(
-        self, node: ir.Region, block_idx: int, curr_block: ir.Block, stmt: For
-    ):
-        from kirin.dialects import cf, py, func
+class ForRule(ScfRule):
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if (
+            not isinstance(node, For)
+            # must be inside a callable statement
+            or not isinstance(parent_stmt := node.parent_stmt, ir.Statement)
+            or not parent_stmt.has_trait(ir.CallableStmtInterface)
+        ):
+            return RewriteResult()
 
-        (body_block := stmt.body.blocks[0]).detach()
+        region, block_idx = self.get_curr_blk_info(node)
+        exit_block, entr_block = self.get_entr_and_exit_blks(node)
 
-        entry_block = ir.Block()
-        for arg in curr_block.args:
-            arg.replace_by(entry_block.args.append_from(arg.type, arg.name))
+        (body_block := node.body.blocks[0]).detach()
 
         # Get iterator from iterable object
-        entry_block.stmts.append(iterable_stmt := py.iterable.Iter(stmt.iterable))
-        entry_block.stmts.append(const_none := func.ConstantNone())
-        last_stmt = entry_block.last_stmt
-        entry_block.stmts.append(
-            next_stmt := py.iterable.Next(iterable_stmt.expect_one_result())
-        )
-        entry_block.stmts.append(
-            loop_cmp := py.cmp.Is(next_stmt.expect_one_result(), const_none.result)
+        entr_block.stmts.append(iterable_stmt := Iter(node.iterable))
+        entr_block.stmts.append(next_stmt := Next(iterable_stmt.expect_one_result()))
+        entr_block.stmts.append(const_none := ConstantNone())
+        entr_block.stmts.append(
+            loop_cmp := Is(next_stmt.expect_one_result(), const_none.result)
         )
-        entry_block.stmts.append(
-            cf.ConditionalBranch(
+        entr_block.stmts.append(
+            ConditionalBranch(
                 cond=loop_cmp.result,
-                then_arguments=tuple(stmt.initializers),
-                then_successor=curr_block,
+                then_arguments=tuple(node.initializers),
+                then_successor=exit_block,
                 else_arguments=(next_stmt.expect_one_result(),)
-                + tuple(stmt.initializers),
+                + tuple(node.initializers),
                 else_successor=body_block,
             )
         )
 
-        for arg in curr_block.args:
-            curr_block.args.delete(arg)
-
-        for arg in stmt.results:
-            arg.replace_by(curr_block.args.append_from(arg.type, arg.name))
-
         if isinstance(last_stmt := body_block.last_stmt, Yield):
+            (next_stmt := Next(iterable_stmt.expect_one_result())).insert_before(
+                last_stmt
+            )
+            (const_none := ConstantNone()).insert_before(last_stmt)
             (
-                next_stmt := py.iterable.Next(iterable_stmt.expect_one_result())
-            ).insert_before(last_stmt)
-            (
-                loop_cmp := py.cmp.Is(next_stmt.expect_one_result(), const_none.result)
+                loop_cmp := Is(next_stmt.expect_one_result(), const_none.result)
             ).insert_before(last_stmt)
             last_stmt.replace_by(
-                cf.ConditionalBranch(
+                ConditionalBranch(
                     cond=loop_cmp.result,
                     else_arguments=(next_stmt.expect_one_result(),)
                     + tuple(last_stmt.args),
                     else_successor=body_block,
                     then_arguments=tuple(last_stmt.args),
-                    then_successor=curr_block,
+                    then_successor=exit_block,
                 )
             )
 
         # insert the body block and add branch to the current block
-        node.blocks.insert(block_idx, curr_block)
-        node.blocks.insert(block_idx, body_block)
-
-        curr_stmt = stmt
-        next_stmt = stmt.prev_stmt
-        curr_stmt.delete()
+        region.blocks.insert(block_idx + 1, exit_block)
+        region.blocks.insert(block_idx + 1, body_block)
+        region.blocks.insert(block_idx + 1, entr_block)
 
-        return next_stmt, entry_block
+        node.delete()
 
-    def rewrite_ssacfg(self, node: ir.Region):
+        return RewriteResult(has_done_something=True)
 
-        has_done_something = False
 
-        for block_idx in range(len(node.blocks)):
+class IfElseRule(ScfRule):
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if (
+            not isinstance(node, IfElse)
+            or not isinstance(parent_stmt := node.parent_stmt, ir.Statement)
+            or not parent_stmt.has_trait(ir.CallableStmtInterface)
+        ):
+            return RewriteResult()
 
-            block = node.blocks.pop(block_idx)
+        region, block_idx = self.get_curr_blk_info(node)
+        exit_block, entr_block = self.get_entr_and_exit_blks(node)
 
-            stmt = block.last_stmt
-            if stmt is None:
-                continue
+        (then_block := node.then_body.blocks[0]).detach()
+        (else_block := node.else_body.blocks[0]).detach()
+        entr_block.stmts.append(
+            ConditionalBranch(
+                node.cond,
+                then_arguments=tuple(node.args),
+                then_successor=then_block,
+                else_arguments=tuple(node.args),
+                else_successor=else_block,
+            )
+        )
 
-            curr_block = ir.Block()
+        if isinstance(last_stmt := then_block.last_stmt, Yield):
+            last_stmt.replace_by(
+                Branch(
+                    arguments=tuple(last_stmt.args),
+                    successor=exit_block,
+                )
+            )
 
-            for arg in block.args:
-                arg.replace_by(curr_block.args.append_from(arg.type, arg.name))
+        if isinstance(last_stmt := else_block.last_stmt, Yield):
+            last_stmt.replace_by(
+                Branch(
+                    arguments=tuple(last_stmt.args),
+                    successor=exit_block,
+                )
+            )
 
-            while stmt is not None:
-                if isinstance(stmt, For):
-                    has_done_something = True
-                    stmt, curr_block = self.rewrite_for(
-                        node, block_idx, curr_block, stmt
-                    )
+        # insert the new blocks
+        region.blocks.insert(block_idx + 1, exit_block)
+        region.blocks.insert(block_idx + 1, else_block)
+        region.blocks.insert(block_idx + 1, then_block)
+        region.blocks.insert(block_idx + 1, entr_block)
 
-                elif isinstance(stmt, IfElse):
-                    has_done_something = True
-                    stmt, curr_block = self.rewrite_ifelse(
-                        node, block_idx, curr_block, stmt
-                    )
-                else:
-                    curr_stmt = stmt
-                    stmt = stmt.prev_stmt
-                    curr_stmt.detach()
+        node.delete()
+        return RewriteResult(has_done_something=True)
 
-                    if curr_block.first_stmt is None:
-                        curr_block.stmts.append(curr_stmt)
-                    else:
-                        curr_stmt.insert_before(curr_block.first_stmt)
 
-            # if the last block is empty, remove it
-            if curr_block.parent is None and curr_block.first_stmt is not None:
-                node.blocks.insert(block_idx, curr_block)
+@dataclass
+class ScfToCfRule(RewriteRule):
 
-        return RewriteResult(has_done_something=has_done_something)
+    for_rule: ForRule = field(default_factory=ForRule, init=False)
+    if_else_rule: IfElseRule = field(default_factory=IfElseRule, init=False)
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        if (
-            isinstance(node, (For, IfElse))
-            or not node.has_trait(ir.HasCFG)
-            and not node.has_trait(ir.SSACFG)
-        ):
-            # do not do rewrite in scf regions
+        if isinstance(node, For):
+            return self.for_rule.rewrite_Statement(node)
+        elif isinstance(node, IfElse):
+            return self.if_else_rule.rewrite_Statement(node)
+        else:
             return RewriteResult()
-
-        result = RewriteResult()
-        for region in node.regions:
-            result = result.join(self.rewrite_ssacfg(region))
-
-        return result

test/dialects/scf/test_scf2cf.py

@@ -99,31 +99,39 @@ def test():
 
     expected_callable_region = ir.Region(
         [
+            curr_block := ir.Block(),
             entry_block := ir.Block(),
             body_block := ir.Block(),
             exit_block := ir.Block(),
         ]
     )
 
-    entry_block.args.append_from(types.MethodType, "self")
-    entry_block.stmts.append(j_start := py.Constant(value=0))
+    curr_block.args.append_from(types.MethodType, "self")
+    curr_block.stmts.append(j_start := py.Constant(value=0))
+
     j_start.result.name = "j"
-    entry_block.stmts.append(iter_start := py.Constant(value=0))
-    entry_block.stmts.append(iter_end := py.Constant(value=10))
-    entry_block.stmts.append(iter_step := py.Constant(value=1))
-    entry_block.stmts.append(
+    curr_block.stmts.append(iter_start := py.Constant(value=0))
+    curr_block.stmts.append(iter_end := py.Constant(value=10))
+    curr_block.stmts.append(iter_step := py.Constant(value=1))
+    curr_block.stmts.append(
         range_stmt := ilist.stmts.Range(
             start=iter_start.result,
             stop=iter_end.result,
             step=iter_step.result,
         )
     )
     range_stmt.result.type = ilist.IListType[types.Int, types.Literal(10)]
+    curr_block.stmts.append(
+        cf.Branch(
+            arguments=(),
+            successor=entry_block,
+        )
+    )
     entry_block.stmts.append(iterable_stmt := py.iterable.Iter(range_stmt.result))
-    entry_block.stmts.append(none_stmt := func.ConstantNone())
     entry_block.stmts.append(
         first_iter := py.iterable.Next(iterable_stmt.expect_one_result())
     )
+    entry_block.stmts.append(none_stmt := func.ConstantNone())
     entry_block.stmts.append(
         loop_cmp := py.cmp.Is(first_iter.expect_one_result(), none_stmt.result)
     )
@@ -157,6 +165,7 @@ def test():
     body_block.stmts.append(
         next_iter := py.iterable.Next(iterable_stmt.expect_one_result())
     )
+    body_block.stmts.append(none_stmt := func.ConstantNone())
     body_block.stmts.append(
         loop_cmp := py.cmp.Is(next_iter.expect_one_result(), none_stmt.result)
     )