Dl/peephole optimize (#486)

adds Peephole optimization rewrites and test cases. For issue
#140

Author: zhenrongliew

src/kirin/rewrite/__init__.py

@@ -8,6 +8,7 @@
 from .getitem import InlineGetItem as InlineGetItem
 from .fixpoint import Fixpoint as Fixpoint
 from .getfield import InlineGetField as InlineGetField
+from .peephole import PeepholeOptimize as PeepholeOptimize
 from .apply_type import ApplyType as ApplyType
 from .compactify import CFGCompactify as CFGCompactify
 from .wrap_const import WrapConst as WrapConst

src/kirin/rewrite/peephole.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass
+
+from kirin import ir, types
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class PeepholeOptimize(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        result_types = node.results.types
+        if not all(
+            t.is_subseteq(types.Union(types.Float, types.Int)) for t in result_types
+        ):
+            return RewriteResult(has_done_something=False)
+
+        if isinstance(node, py.binop.Add):
+            #   add(%a, %a) -> mul(2, %a)
+            if node.lhs is node.rhs:
+                x = py.Constant(2)
+                x.insert_before(node)
+                new_stmt = py.binop.Mult(x.result, node.rhs)
+                node.replace_by(new_stmt)
+                return RewriteResult(has_done_something=True)
+
+            #   add(mul(2, %a), %a) -> mul(3, %a)
+            elif isinstance(mult_node := node.lhs.owner, py.binop.Mult) and isinstance(
+                const_node := mult_node.lhs.owner, py.Constant
+            ):
+                x = const_node.value.unwrap()
+                const_node.replace_by(py.Constant(x + 1))
+                node.replace_by(py.binop.Mult(mult_node.lhs, node.rhs))
+                mult_node.delete()
+                return RewriteResult(has_done_something=True)
+
+            #   add(%a, mul(2, %a)) -> mul(3, %a)
+            elif isinstance(mult_node := node.rhs.owner, py.binop.Mult) and isinstance(
+                const_node := mult_node.lhs.owner, py.Constant
+            ):
+                x = const_node.value.unwrap()
+                const_node.replace_by(py.Constant(x + 1))
+                node.replace_by(py.binop.Mult(mult_node.lhs, node.lhs))
+                mult_node.delete()
+                return RewriteResult(has_done_something=True)
+
+        return RewriteResult(has_done_something=False)

test/rewrite/test_peephole.py

@@ -0,0 +1,55 @@
+from kirin.prelude import basic_no_opt
+from kirin.rewrite import Walk, Fixpoint
+from kirin.rewrite.peephole import PeepholeOptimize
+
+
+# add(%a, %a) -> mul(2, %a)
+@basic_no_opt
+def peephole1(a: int):
+    x = a + a
+    return x
+
+
+# add(mul(2, %a), %a) -> mul(3, %a)
+@basic_no_opt
+def peephole2(a: int):
+    x = 2 * a + a
+    return x
+
+
+# add(%a, mul(2, %a)) -> mul(3, %a)
+@basic_no_opt
+def peephole3(a: int):
+    x = a + 2 * a
+    return x
+
+
+# add(%a, add(%a, mul(2, %a))) -> mul(4, %a)
+@basic_no_opt
+def peephole4(a: int):
+    x = a + a + 2 * a
+    return x
+
+
+def aux(program):
+    for i in range(5):
+        before = program(i)
+        Fixpoint(Walk(PeepholeOptimize())).rewrite(program.code)
+        after = program(i)
+        assert before == after
+
+
+def test_peephole_opt1():
+    aux(peephole1)
+
+
+def test_peephole_opt2():
+    aux(peephole2)
+
+
+def test_peephole_opt3():
+    aux(peephole3)
+
+
+def test_peephole_opt4():
+    aux(peephole4)