patch aggressive inline for v0.17 (#513)

Fix #508 (this is fixed in main but needs a patch for release-0-17). Add
logic in the `Inline` rewrite rule to permute kwargs.

---------

Co-authored-by: Neelay Fruitwala <[email protected]>
Co-authored-by: Neelay Fruitwala <[email protected]>

Author: 3 people

src/kirin/rewrite/inline.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass
 
 from kirin import ir
+from kirin.interp import BaseInterpreter
 from kirin.dialects import cf, func
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
@@ -30,7 +31,13 @@ def rewrite_func_Call(self, node: func.Call) -> RewriteResult:
             return RewriteResult()
 
         # NOTE: a lambda statement is defined and used in the same scope
-        self.inline_call_like(node, tuple(node.args), lambda_stmt.body)
+        arg_names = [arg.name for arg in node.callee.owner.body.blocks[0].args]
+        args = BaseInterpreter.permute_values(
+            arg_names=arg_names,
+            values=tuple(node.args[1:]),
+            kwarg_names=node.kwargs,
+        )
+        self.inline_call_like(node, (node.args[0],) + args, lambda_stmt.body)
         return RewriteResult(has_done_something=True)
 
     def rewrite_func_Invoke(self, node: func.Invoke) -> RewriteResult:
@@ -47,9 +54,12 @@ def rewrite_func_Invoke(self, node: func.Invoke) -> RewriteResult:
             func_self = Constant(node.callee)
             func_self.result.name = node.callee.sym_name
             func_self.insert_before(node)
-            self.inline_call_like(
-                node, (func_self.result,) + tuple(arg for arg in node.args), region
+            args = BaseInterpreter.permute_values(
+                arg_names=node.callee.arg_names,
+                values=tuple(node.args),
+                kwarg_names=node.kwargs,
             )
+            self.inline_call_like(node, (func_self.result,) + tuple(args), region)
             has_done_something = True
 
         return RewriteResult(has_done_something=has_done_something)

test/dialects/test_aggressive_inline.py

@@ -0,0 +1,89 @@
+from kirin.passes import aggressive
+from kirin.prelude import basic
+
+
+def test_aggressive_inline():
+
+    @basic(aggressive=False)
+    def foo0(arg0, arg1):
+        return arg0 - arg1
+
+    @basic(aggressive=False)
+    def main_aggressive(arg0):
+        return foo0(arg1=2, arg0=arg0)
+
+    main_aggressive = main_aggressive.similar()
+    aggressive.Fold(main_aggressive.dialects).fixpoint(main_aggressive)
+
+    assert main_aggressive(1) == -1
+
+
+def test_aggressive_inline_noargs():
+
+    @basic(aggressive=False)
+    def foo1(arg0, arg1):
+        return arg0 - arg1
+
+    @basic(aggressive=True)
+    def main_aggressive():
+        return foo1(arg1=2, arg0=1)
+
+    assert main_aggressive() == -1
+
+
+def test_aggressive_inline_pos_args():
+
+    @basic(aggressive=False)
+    def foo2(arg0, arg1):
+        return arg0 - arg1
+
+    @basic(aggressive=True)
+    def main_aggressive(arg0):
+        return foo2(arg0, 2)
+
+    assert main_aggressive(1) == -1
+
+
+def test_aggressive_inline_closure():
+
+    # @basic(aggressive=False, fold=False, typeinfer=True)
+    @basic
+    def main_aggressive(param: int):
+        def foo3(arg0: int, arg1: int):
+            return arg0 - arg1 + param
+
+        return foo3(arg1=2, arg0=1)
+
+    main_aggressive = main_aggressive.similar()
+    aggressive.Fold(main_aggressive.dialects).fixpoint(main_aggressive)
+
+    assert main_aggressive(1) == 0
+
+
+def test_aggressive_inline_closure_pos_args():
+
+    # @basic(aggressive=False, fold=False, typeinfer=True)
+    @basic
+    def main_aggressive(param: int):
+        def foo3(arg0: int, arg1: int):
+            return arg0 - arg1 + param
+
+        return foo3(1, arg1=2)
+
+    main_aggressive = main_aggressive.similar()
+    aggressive.Fold(main_aggressive.dialects).fixpoint(main_aggressive)
+
+    assert main_aggressive(1) == 0
+
+
+def test_aggressive_inline_closure_alias():
+    @basic(aggressive=True)
+    def main_aggressive2(param: int):
+        def foo4(arg0: int, arg1: int):
+            return arg0 - arg1 + param
+
+        alias_foo4 = foo4
+
+        return alias_foo4(arg1=2, arg0=1)
+
+    assert main_aggressive2(1) == 0