support lowering loop to cf (#187)

this PR supports lowering loop to `cf`, e.g

```python
@basic_no_opt
def main(x: int):
    for i in range(5):
        for j in range(5):
            x = x + 1
    return x
```

<details>
<summary>compiles to</summary>

```mlir
func.func main(!py.int) -> !Any {
  ^0(%main_self, %x_1):
  │  %0 = py.constant.constant 0 : !py.int
  │  %1 = py.constant.constant 5 : !py.int
  │  %2 = py.constant.constant 1 : !py.int
  │  %3 = py.range.range(start=%0, stop=%1, step=%2) : !py.range
  │  %4 = py.iterable.iter(value=%3) : !Any
  │  %5 = py.constant.constant None : !py.NoneType
  │  %6 = py.iterable.next(iter=%4) : !Any
  │  %7 = py.cmp.is(lhs=%6, rhs=%5) : !py.bool
  │       cf.cond_br %7 goto ^4(%x_1 : !py.int) else ^1(%6, %x_1 : !py.int)
  ^1(%i, %x_2):
  │  %8 = py.constant.constant 0 : !py.int
  │  %9 = py.constant.constant 5 : !py.int
  │ %10 = py.constant.constant 1 : !py.int
  │ %11 = py.range.range(start=%8, stop=%9, step=%10) : !py.range
  │ %12 = py.iterable.iter(value=%11) : !Any
  │ %13 = py.constant.constant None : !py.NoneType
  │ %14 = py.iterable.next(iter=%12) : !Any
  │ %15 = py.cmp.is(lhs=%14, rhs=%13) : !py.bool
  │       cf.cond_br %15 goto ^3(%x_2 : !py.int) else ^2(%14, %x_2 : !py.int)
  ^2(%j, %x_3):
  │ %16 = py.constant.constant 1 : !py.int
  │  %x = py.binop.add(%x_3 : !py.int, %16) : ~T
  │ %17 = py.iterable.next(iter=%12) : !Any
  │ %18 = py.cmp.is(lhs=%17, rhs=%13) : !py.bool
  │       cf.cond_br %18 goto ^3(%x) else ^2(%17, %x)
  ^3(%x_4):
  │ %19 = py.iterable.next(iter=%4) : !Any
  │ %20 = py.cmp.is(lhs=%19, rhs=%5) : !py.bool
  │       cf.cond_br %20 goto ^4(%x_4 : ~T) else ^1(%19, %x_4 : ~T)
  ^4(%x_5):
  │       func.return %x_5 : ~T
  ^5():
  │ %21 = func.const.none() : !py.NoneType
  │       func.return %21
} // func.func main
```
</details>

Author: Roger-luo

src/kirin/analysis/cfg.py

@@ -44,6 +44,7 @@ def predecessors(self):
     def successors(self):
         """CFG data, mapping a block to its neighbors."""
         graph: dict[ir.Block, set[ir.Block]] = {}
+        visited: set[ir.Block] = set()
         worklist: WorkList[ir.Block] = WorkList()
         if self.parent.blocks.isempty():
             return graph
@@ -54,8 +55,11 @@ def successors(self):
             if block.last_stmt is not None:
                 neighbors.update(block.last_stmt.successors)
                 worklist.extend(block.last_stmt.successors)
+            visited.add(block)
 
             block = worklist.pop()
+            while block is not None and block in visited:
+                block = worklist.pop()
         return graph
 
     # graph interface

src/kirin/dialects/cf/stmts.py

@@ -13,6 +13,9 @@ class Branch(Statement):
     arguments: tuple[SSAValue, ...]
     successor: Block = info.block()
 
+    def verify(self) -> None:
+        return
+
     def print_impl(self, printer: Printer) -> None:
         with printer.rich(style="keyword"):
             printer.print_name(self)
@@ -61,3 +64,6 @@ def print_impl(self, printer: Printer) -> None:
         printer.plain_print("(")
         printer.print_seq(self.else_arguments, delim=", ")
         printer.plain_print(")")
+
+    def verify(self) -> None:
+        return

src/kirin/dialects/lowering/cf.py

@@ -4,9 +4,9 @@
 import ast
 
 from kirin import ir
+from kirin.dialects import cf, py
 from kirin.lowering import Frame, Result, FromPythonAST, LoweringState
 from kirin.exceptions import DialectLoweringError
-from kirin.dialects.cf import stmts as cf
 
 dialect = ir.Dialect("lowering.cf")
 
@@ -21,6 +21,83 @@ def lower_Pass(self, state: LoweringState, node: ast.Pass) -> Result:
         state.append_stmt(cf.Branch(arguments=(), successor=next))
         return Result()
 
+    def lower_For(self, state: LoweringState, node: ast.For) -> Result:
+        iterable = state.visit(node.iter).expect_one()
+        iter_stmt = state.append_stmt(py.iterable.Iter(iterable))
+        none_stmt = state.append_stmt(py.Constant(None))
+        yields: list[str] = []
+
+        def new_block_arg_if_inside_loop(frame: Frame, capture: ir.SSAValue):
+            if not capture.name:
+                raise DialectLoweringError("unexpected loop variable captured")
+            yields.append(capture.name)
+            return frame.entry_block.args.append_from(capture.type, capture.name)
+
+        frame = state.current_frame
+        next_block = ir.Block()
+
+        body_frame = state.push_frame(
+            Frame.from_stmts(
+                node.body,
+                state,
+                region=frame.current_region,
+                globals=state.current_frame.globals,
+                capture_callback=new_block_arg_if_inside_loop,
+            )
+        )
+        body_frame.next_block = next_block
+        next_value = body_frame.entry_block.args.append_from(ir.types.Any, "next_value")
+        py.unpack.unpackable(state, node.target, next_value)
+        state.exhaust(body_frame)
+        next_stmt = body_frame.append_stmt(py.iterable.Next(iter_stmt.iter))
+        cond_stmt = body_frame.append_stmt(py.cmp.Is(next_stmt.value, none_stmt.result))
+        yield_args = tuple(body_frame.get_scope(name) for name in yields)
+        state.pop_frame()
+
+        next_frame = state.push_frame(
+            Frame.from_stmts(
+                frame.stream.split(),
+                state,
+                region=frame.current_region,
+                block=next_block,
+                globals=frame.globals,
+            )
+        )
+        next_frame.next_block = frame.next_block
+        for name, arg in zip(yields, yield_args):
+            input = next_frame.current_block.args.append_from(arg.type, name)
+            next_frame.defs[name] = input
+        state.exhaust()
+        state.pop_frame()
+
+        yield_args = tuple(body_frame.get_scope(name) for name in yields)
+        body_frame.append_stmt(
+            cf.ConditionalBranch(
+                cond_stmt.result,
+                yield_args,
+                (next_stmt.value,) + yield_args,
+                then_successor=next_frame.entry_block,
+                else_successor=body_frame.entry_block,
+            )
+        )
+
+        next_stmt = frame.append_stmt(py.iterable.Next(iter_stmt.iter))
+        cond_stmt = frame.append_stmt(py.cmp.Is(next_stmt.value, none_stmt.result))
+        yield_args = tuple(frame.get_scope(name) for name in yields)
+        frame.append_stmt(
+            cf.ConditionalBranch(
+                cond_stmt.result,
+                yield_args,
+                (next_stmt.value,) + yield_args,
+                then_successor=next_frame.entry_block,
+                else_successor=body_frame.entry_block,
+            )
+        )
+        frame.current_block = next_frame.current_block
+        frame.next_block = next_frame.next_block
+        frame.defs.update(next_frame.defs)
+        return Result()
+
     def lower_If(self, state: LoweringState, node: ast.If) -> Result:
         cond = state.visit(node.test).expect_one()
 

src/kirin/dialects/py/__init__.py

@@ -15,6 +15,7 @@
     builtin as builtin,
     constant as constant,
     indexing as indexing,
+    iterable as iterable,
 )
 from .len import Len as Len
 from .attr import GetAttr as GetAttr

src/kirin/dialects/py/iterable.py

@@ -0,0 +1,81 @@
+"""This module provides access to Python iterables.
+
+This is used to lower Python loops into `cf` dialect.
+"""
+
+from ast import Call
+
+from kirin import ir, interp, lowering
+from kirin.decl import info, statement
+from kirin.exceptions import DialectLoweringError
+
+dialect = ir.Dialect("py.iterable")
+
+PyRangeIterType = ir.types.PyClass(type(iter(range(0))))
+
+
+@statement(dialect=dialect)
+class Iter(ir.Statement):
+    """This is equivalent to `iter(value)` in Python."""
+
+    traits = frozenset({ir.Pure()})
+    value: ir.SSAValue = info.argument(ir.types.Any)
+    iter: ir.ResultValue = info.result(ir.types.Any)
+
+
+@statement(dialect=dialect)
+class Next(ir.Statement):
+    """This is equivalent to `next(iterable, None)` in Python."""
+
+    iter: ir.SSAValue = info.argument(ir.types.Any)
+    value: ir.ResultValue = info.result(ir.types.Any)
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Iter)
+    def iter_(self, interp, frame: interp.Frame, stmt: Iter):
+        return (iter(frame.get(stmt.value)),)
+
+    @interp.impl(Next)
+    def next_(self, interp, frame: interp.Frame, stmt: Next):
+        return (next(frame.get(stmt.iter), None),)
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(Iter, ir.types.PyClass(range))
+    def iter_(self, interp, frame: interp.Frame, stmt: Iter):
+        return (PyRangeIterType,)
+
+    @interp.impl(Next, PyRangeIterType)
+    def next_(self, interp, frame: interp.Frame, stmt: Next):
+        return (ir.types.Int,)
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_Call_iter(
+        self, state: lowering.LoweringState, node: Call
+    ) -> lowering.Result:
+        if len(node.args) != 1:
+            raise DialectLoweringError("iter() takes exactly 1 argument")
+        return lowering.Result(
+            state.append_stmt(Iter(state.visit(node.args[0]).expect_one()))
+        )
+
+    def lower_Call_next(
+        self, state: lowering.LoweringState, node: Call
+    ) -> lowering.Result:
+        if len(node.args) == 2:
+            raise DialectLoweringError(
+                "next() does not throw StopIteration inside kernel"
+            )
+        if len(node.args) != 1:
+            raise DialectLoweringError("next() takes exactly 1 argument")
+        return lowering.Result(
+            state.append_stmt(Next(state.visit(node.args[0]).expect_one()))
+        )

src/kirin/interp/abstract.py

@@ -2,7 +2,7 @@
 from typing import TypeVar, Iterable, TypeAlias, overload
 from dataclasses import field, dataclass
 
-from kirin.ir import Region, SSAValue, Statement
+from kirin.ir import Block, Region, SSAValue, Statement
 from kirin.lattice import BoundedLattice
 from kirin.worklist import WorkList
 from kirin.interp.base import BaseInterpreter, InterpreterMeta
@@ -27,6 +27,7 @@ class AbstractFrame(Frame[ResultType]):
     """
 
     worklist: WorkList[Successor[ResultType]] = field(default_factory=WorkList)
+    visited: dict[Block, set[Successor[ResultType]]] = field(default_factory=dict)
 
 
 AbstractFrameType = TypeVar("AbstractFrameType", bound=AbstractFrame)
@@ -151,8 +152,18 @@ def run_ssacfg_region(
             Successor(region.blocks[0], *frame.get_values(region.blocks[0].args))
         )
         while (succ := frame.worklist.pop()) is not None:
+            if succ.block in frame.visited:
+                if succ in frame.visited[succ.block]:
+                    continue
+            else:
+                frame.visited[succ.block] = set()
             self.prehook_succ(frame, succ)
             block_result = self.run_block(frame, succ)
+            if len(frame.visited[succ.block]) < 128:
+                frame.visited[succ.block].add(succ)
+            else:
+                continue
+
             if isinstance(block_result, Successor):
                 raise InterpreterError(
                     "unexpected successor, successors should be in worklist"

src/kirin/lowering/dialect.pyi

@@ -75,6 +75,8 @@ class FromPythonAST:
     def lower_Call_slice(self, state: LoweringState, node: ast.Call) -> Result: ...
     def lower_Call_range(self, state: LoweringState, node: ast.Call) -> Result: ...
     def lower_Call_len(self, state: LoweringState, node: ast.Call) -> Result: ...
+    def lower_Call_iter(self, state: LoweringState, node: ast.Call) -> Result: ...
+    def lower_Call_next(self, state: LoweringState, node: ast.Call) -> Result: ...
     def lower_Call_local(
         self, state: LoweringState, callee: SSAValue, node: ast.Call
     ) -> Result: ...

src/kirin/prelude.py

@@ -24,6 +24,7 @@
     builtin,
     constant,
     indexing,
+    iterable,
     assertion,
 )
 from kirin.passes.fold import Fold
@@ -45,6 +46,7 @@
         len,
         tuple,
         assertion,
+        iterable,
     ]
 )
 def python_basic(self):

test/program/py/test_loop.py

@@ -0,0 +1,38 @@
+from kirin.prelude import basic_no_opt
+
+
+def test_simple():
+    @basic_no_opt
+    def main(x: int):
+        for i in range(5):
+            x = x + 1
+        return x
+
+    assert main.py_func is not None
+    assert main.py_func(1) == main(1)
+
+
+# generate some more complicated loop
+def test_nested():
+    @basic_no_opt
+    def main(x: int):
+        for i in range(5):
+            for j in range(5):
+                x = x + 1
+        return x
+
+    assert main.py_func is not None
+    assert main.py_func(1) == main(1)
+
+
+def test_nested2():
+    @basic_no_opt
+    def main(x: int):
+        for i in range(5):
+            for j in range(5):
+                for k in range(5):
+                    x = x + 1
+        return x
+
+    assert main.py_func is not None
+    assert main.py_func(1) == main(1)