Moved ClosureField into dialects/func. Implemented LambdaLifting. (#543)

Renamed `ClosureFieldLowering` to `ClosureField` and moved it into the
`dialects/func` directory because it is specific to `func.Invoke`.
Register the rewrite as canonicalize rule.

Also, implemented `LambdaLifting` that rewrites `py.Constant(ir.Method)`
with method body being lambda into a regular function.

e.g, this function:
```python
@basic
def outer():
    def inner(x: int):
        return x + 1

    return inner
```
Will rewrite from:
```
func.func @outer() -> !Any {
  ^0(%outer_self):
  │ %inner = py.constant.constant Method("inner") : !py.Method[!py.tuple[!py.int], !Any]
  │          func.return %inner
} // func.func outer
func.lambda inner() -> !Any {
  ^0(%inner_self, %x):
  │ %0 = py.constant.constant 1 : !py.int
  │ %1 = py.binop.add(%x : !py.int, %0) : ~T
  │      func.return %1
} // func.lambda inner
```

to:
```
func.func @outer() -> !Any {
  ^0(%outer_self):
  │ %inner = py.constant.constant Method("inner") : !py.Method[!py.tuple[!py.int], !Any]
  │          func.return %inner
} // func.func outer
func.func @inner(x : !py.int) -> !Any {
  ^0(%inner_self, %x):
  │ %0 = py.constant.constant 1 : !py.int
  │ %1 = py.binop.add(%x : !py.int, %0) : ~T
  │      func.return %1
} // func.func inner
```

Added test files for these 2 rewrites.

Author: zhenrongliew

src/kirin/dialects/func/__init__.py

@@ -18,4 +18,8 @@
 )
 from kirin.dialects.func._dialect import dialect as dialect
 
-from . import _julia as _julia
+from . import (
+    _julia as _julia,
+    closurefield as closurefield,
+    lambdalifting as lambdalifting,
+)

src/kirin/rewrite/closurefieldlowering.py src/kirin/dialects/func/closurefield.pysrc/kirin/rewrite/closurefieldlowering.py renamed to src/kirin/dialects/func/closurefield.py

@@ -1,12 +1,12 @@
-from dataclasses import dataclass
-
 from kirin import ir
 from kirin.dialects import py, func
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
+from ._dialect import dialect
+
 
-@dataclass
-class ClosureFieldLowering(RewriteRule):
+@dialect.canonicalize
+class ClosureField(RewriteRule):
     """Lowers captured closure fields into py.Constants.
     - Trigger on func.Invoke
     - If the callee Method has non-empty .fields, lower its func.GetField to py.Constant
@@ -15,15 +15,17 @@ class ClosureFieldLowering(RewriteRule):
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         if not isinstance(node, func.Invoke):
             return RewriteResult(has_done_something=False)
-
         method = node.callee
         if not method.fields:
             return RewriteResult(has_done_something=False)
         # Replace func.GetField with py.Constant.
         changed = self._lower_captured_fields(method)
         if changed:
             method.fields = ()
-        return RewriteResult(has_done_something=changed)
+        from kirin.passes import TypeInfer
+
+        rewrite_result = TypeInfer(dialects=method.dialects).unsafe_run(method)
+        return RewriteResult(has_done_something=changed).join(rewrite_result)
 
     def _get_field_index(self, getfield_stmt: func.GetField) -> int | None:
         fld = getfield_stmt.attributes.get("field")

src/kirin/dialects/func/lambdalifting.py

@@ -0,0 +1,72 @@
+from kirin import ir
+from kirin.dialects import py, func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from ._dialect import dialect
+
+
+@dialect.canonicalize
+class LambdaLifting(RewriteRule):
+    """Lifts func.Lambda methods embedded in py.Constant into func.Function.
+    - Trigger on py.Constant
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, py.Constant):
+            return RewriteResult(has_done_something=False)
+        method = self._get_method_from_constant(node)
+        if method is None:
+            return RewriteResult(has_done_something=False)
+        if not isinstance(method.code, func.Lambda):
+            return RewriteResult(has_done_something=False)
+        self._promote_lambda(method)
+
+        from kirin.passes import TypeInfer
+
+        rewrite_result = TypeInfer(dialects=method.dialects).unsafe_run(method)
+        return RewriteResult(has_done_something=True).join(rewrite_result)
+
+    def _get_method_from_constant(self, const_stmt: py.Constant) -> ir.Method | None:
+        pyattr_data = const_stmt.value
+        if isinstance(pyattr_data, ir.PyAttr) and isinstance(
+            pyattr_data.data, ir.Method
+        ):
+            return pyattr_data.data
+        return None
+
+    def _get_field_index(self, getfield_stmt: func.GetField) -> int | None:
+        fld = getfield_stmt.attributes.get("field")
+        if fld:
+            return getfield_stmt.field
+        else:
+            return None
+
+    def _promote_lambda(self, method: ir.Method) -> None:
+        new_method = method.similar()
+        assert isinstance(
+            new_method.code, func.Lambda
+        ), "expected method.code to be func.Lambda before promotion"
+
+        captured_fields = method.fields
+        if captured_fields:
+            for stmt in new_method.code.body.blocks[0].stmts:
+                if not isinstance(stmt, func.GetField):
+                    continue
+                idx = self._get_field_index(stmt)
+                if idx is None:
+                    continue
+                captured = new_method.fields[idx]
+                const_stmt = py.Constant(captured)
+                const_stmt.insert_before(stmt)
+                if stmt.results and const_stmt.results:
+                    stmt.results[0].replace_by(const_stmt.results[0])
+                stmt.delete()
+                new_method.code
+
+        fn = func.Function(
+            sym_name=new_method.code.sym_name,
+            slots=new_method.code.slots,
+            signature=new_method.code.signature,
+            body=new_method.code.body,
+        )
+        method.code = fn

src/kirin/passes/fold.py

@@ -10,7 +10,6 @@
     CFGCompactify,
     InlineGetItem,
     DeadCodeElimination,
-    ClosureFieldLowering,
 )
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
@@ -29,7 +28,6 @@ class Fold(Pass):
     - `InlineGetItem`
     - `Call2Invoke`
     - `DeadCodeElimination`
-    - `ClosureFieldLowering`
     """
 
     hint_const: HintConst = field(init=False)
@@ -48,7 +46,6 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
                         InlineGetItem(),
                         Call2Invoke(),
                         DeadCodeElimination(),
-                        ClosureFieldLowering(),
                     )
                 )
             )

src/kirin/rewrite/__init__.py

@@ -13,4 +13,3 @@
 from .wrap_const import WrapConst as WrapConst
 from .call2invoke import Call2Invoke as Call2Invoke
 from .type_assert import InlineTypeAssert as InlineTypeAssert
-from .closurefieldlowering import ClosureFieldLowering as ClosureFieldLowering

test/dialects/func/__init__.py

@@ -0,0 +1,52 @@
+from typing import cast
+
+from kirin import rewrite
+from kirin.prelude import basic
+from kirin.dialects import py, func
+
+
+def test_rewrite_closure_inner_lambda():
+    @basic
+    def outer(y: int):
+        def inner(x: int):
+            return x * y + 1
+
+        return inner
+
+    inner_ker = outer(y=10)
+
+    @basic
+    def main_lambda(z: int):
+        return inner_ker(z)
+
+    main_invoke = main_lambda.code.regions[0].blocks[0].stmts.at(0)
+    inner_lambda = cast(func.Invoke, main_invoke).callee.code
+    inner_getfield_stmt = inner_lambda.regions[0].blocks[0].stmts.at(0)
+    assert isinstance(
+        inner_getfield_stmt, func.GetField
+    ), "expected GetField before rewrite"
+
+    rewrite.Walk(func.closurefield.ClosureField()).rewrite(main_lambda.code)
+
+    inner_getfield_stmt = inner_lambda.regions[0].blocks[0].stmts.at(0)
+    assert isinstance(
+        inner_getfield_stmt, py.Constant
+    ), "GetField should be lowered to Constant"
+
+
+def test_rewrite_closure_no_fields():
+    @basic
+    def bar():
+        def goo(x: int):
+            a = (3, 4)
+            return a[0]
+
+        def boo(y):
+            return goo(y) + 1
+
+        return boo(4)
+
+    before = bar.code.regions[0].blocks[0].stmts.at(0)
+    rewrite.Walk(func.closurefield.ClosureField()).rewrite(bar.code)
+    after = bar.code.regions[0].blocks[0].stmts.at(0)
+    assert before is after

test/dialects/func/test_closurefield.py

@@ -0,0 +1,51 @@
+from kirin import ir, rewrite
+from kirin.prelude import basic
+from kirin.dialects import py, func
+
+
+def test_rewrite_inner_lambda():
+    @basic
+    def outer():
+        def inner(x: int):
+            return x + 1
+
+        return inner
+
+    pyconstant_stmt = outer.code.regions[0].blocks[0].stmts.at(0)
+    assert isinstance(pyconstant_stmt, py.Constant), "expected a Constant in outer body"
+    assert isinstance(
+        pyconstant_stmt.value, ir.PyAttr
+    ), "expected a PyAttr in outer body"
+    assert isinstance(
+        pyconstant_stmt.value.data.code, func.Lambda
+    ), "expected a lambda Method in outer body"
+
+    rewrite.Walk(func.lambdalifting.LambdaLifting()).rewrite(outer.code)
+    assert isinstance(
+        pyconstant_stmt.value.data.code, func.Function
+    ), "expected a Function in outer body"
+
+
+def test_rewrite_inner_lambda_with_captured_vars():
+    @basic
+    def outer2():
+        z = 10
+        y = 3 + z
+
+        def inner2(x: int):
+            return x + y + 5
+
+        return inner2
+
+    pyconstant_stmt = outer2.code.regions[0].blocks[0].stmts.at(0)
+    assert isinstance(pyconstant_stmt, py.Constant), "expected a Constant in outer body"
+    assert isinstance(
+        pyconstant_stmt.value, ir.PyAttr
+    ), "expected a PyAttr in outer body"
+    assert isinstance(
+        pyconstant_stmt.value.data.code, func.Lambda
+    ), "expected a lambda Method in outer body"
+    rewrite.Walk(func.lambdalifting.LambdaLifting()).rewrite(outer2.code)
+    assert isinstance(
+        pyconstant_stmt.value.data.code, func.Function
+    ), "expected a Function in outer body"