fcf.Map Inline (#130)

Co-authored-by: Xiu-zhe (Roger) Luo <[email protected]>

Author: kaihsin

src/kirin/dialects/fcf/interp.py

@@ -36,7 +36,7 @@ def foldr(self, interp: Interpreter, frame: Frame, stmt: Foldr):
         return (acc,)
 
     @impl(Map)
-    def map_list(self, interp: Interpreter, frame: Frame, stmt: Map):
+    def map_tuple(self, interp: Interpreter, frame: Frame, stmt: Map):
         fn: ir.Method = frame.get(stmt.fn)
         coll = frame.get(stmt.coll)
         ret = []
@@ -47,7 +47,7 @@ def map_list(self, interp: Interpreter, frame: Frame, stmt: Map):
                 return _ret
             else:
                 ret.append(_ret)
-        return (ret,)
+        return (tuple(ret),)
 
     @impl(Scan)
     def scan(self, interp: Interpreter, frame: Frame, stmt: Scan):

src/kirin/dialects/fcf/rewrite/__init__.py

@@ -0,0 +1,50 @@
+from typing import Dict
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.analysis import const
+from kirin.dialects import py, fcf, func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.ir.nodes.stmt import Statement
+
+
+@dataclass
+class InlineFcfMap(RewriteRule):
+    cp_results: Dict[ir.SSAValue, const.JointResult]
+
+    def rewrite_Statement(self, node: Statement) -> RewriteResult:
+        match node:
+            case fcf.Map():
+                return self.rewrite_fcf_map(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_fcf_map(self, node: fcf.Map) -> RewriteResult:
+        # TODO make this more generic without the need for the constprop results
+        tmp = self.cp_results.get(node.coll, None)
+
+        if (tmp is None) or (not isinstance(tmp.const, const.Value)):
+            return RewriteResult()
+
+        coll = tmp.const.data
+
+        # rewrite to directly inline:
+        # get the method:
+        tpl_elem = []
+        curr = node
+        for i in coll:
+            new_c = py.Constant(value=i)
+            newstmt = func.Call(callee=node.fn, inputs=(new_c.result,), kwargs=())
+
+            new_c.insert_after(curr)
+            newstmt.insert_after(new_c)
+            tpl_elem.append(newstmt.result)
+            curr = newstmt
+
+        # assemble tuple:
+        tpl = py.tuple.New(values=tuple(tpl_elem))
+        tpl.insert_after(curr)
+        node.result.replace_by(tpl.result)
+        node.delete()
+
+        return RewriteResult(has_done_something=True)

src/kirin/dialects/fcf/rewrite/fcfmap_inline.py

@@ -46,7 +46,7 @@ def main():
     """The kernel function to apply. The function should have signature `fn(x: int) -> Any`."""
     coll: ir.SSAValue = info.argument(ir.types.Any)
     """The iterable to map over."""
-    result: ir.ResultValue = info.result(ir.types.List)
+    result: ir.ResultValue = info.result(ir.types.Tuple)
     """The list of results."""
 
 

src/kirin/dialects/fcf/stmts.py

@@ -43,6 +43,6 @@ def test_fold():
     assert foldl(xs) == sum(xs)
     assert foldr(xs) == sum(xs)
     assert map_list.return_type.is_subseteq(types.List[types.Float])
-    assert map_list([1, 2, 3]) == [2.0, 3.0, 4.0]
+    assert map_list([1, 2, 3]) == (2.0, 3.0, 4.0)
     assert scan([1, 2, 3, 4, 5]) == (15, [1, 2, 3, 4, 5])
     assert scan.return_type.is_subseteq(types.Tuple[types.Int, types.List[types.Int]])

test/dialects/fcf/test_fold.py

@@ -11,4 +11,4 @@ def map_func(i):
 
 
 def test_enumerate_kirin():
-    assert enumerate_kirin([1, 2, 3, 4, 5]) == [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)]
+    assert enumerate_kirin([1, 2, 3, 4, 5]) == ((0, 1), (1, 2), (2, 3), (3, 4), (4, 5))

test/dialects/fcf/test_map.py

@@ -44,4 +44,4 @@ def map_range(x: range):
 
 def test_map_range():
     assert map_range.return_type.is_subseteq(types.List[types.Float])
-    assert map_range(range(10)) == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
+    assert map_range(range(10)) == (1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0)

test/dialects/pystmts/test_range.py

@@ -0,0 +1,53 @@
+from kirin.prelude import basic
+from kirin.rewrite import Walk
+from kirin.analysis import const
+from kirin.dialects import fcf
+from kirin.analysis.const.prop import Propagate
+from kirin.dialects.fcf.rewrite.fcfmap_inline import InlineFcfMap
+
+
+def test_fcfmap_rewrite():
+
+    @basic(fold=False)
+    def fcf_map_rewrite():
+
+        def _simple(i: int):
+            return i
+
+        tmp = fcf.Map(_simple, range(5))
+        return tmp
+
+    fcf_map_rewrite.code.print()
+    cp = Propagate(dialects=fcf_map_rewrite.dialects)
+    cp.eval(fcf_map_rewrite, ())
+    Walk(InlineFcfMap(cp.results)).rewrite(fcf_map_rewrite.code)
+    fcf_map_rewrite.code.print()
+
+    val = fcf_map_rewrite()
+
+    assert val == (0, 1, 2, 3, 4)
+
+
+def test_fcfmap_rewrite_with_arg():
+
+    @basic(fold=False)
+    def fcf_map_rewrite_with_arg(x: int):
+
+        def _simple(i: int):
+            return i
+
+        tmp = fcf.Map(_simple, (x, x + 1, x + 2))
+        return tmp
+
+    fcf_map_rewrite_with_arg.code.print()
+    cp = Propagate(dialects=fcf_map_rewrite_with_arg.dialects)
+    cp.eval(fcf_map_rewrite_with_arg, (const.JointResult.top(),))
+    Walk(InlineFcfMap(cp.results)).rewrite(fcf_map_rewrite_with_arg.code)
+    fcf_map_rewrite_with_arg.code.print()
+
+    val = fcf_map_rewrite_with_arg(4)
+    print(val)
+    assert val == (4, 5, 6)
+
+
+test_fcfmap_rewrite_with_arg()