Rewrite py.Mult to squin.Mult in squin kernel

Author: david-pl

src/bloqade/squin/groups.py

@@ -3,13 +3,15 @@
 from kirin.dialects import ilist
 
 from . import op, wire, qubit
+from .op.rewrite import PyMultToSquinMult
 
 
 @ir.dialect_group(structural_no_opt.union([op, qubit]))
 def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
     ilist_desugar_pass = ilist.IListDesugar(self)
+    py_mult_to_mult_pass = PyMultToSquinMult(self)
 
     def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
         method.verify()
@@ -18,7 +20,10 @@ def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
 
         if typeinfer:
             typeinfer_pass(method)
+
         ilist_desugar_pass(method)
+        py_mult_to_mult_pass(method)
+
         if typeinfer:
             typeinfer_pass(method)  # fix types after desugaring
             method.verify_type()

src/bloqade/squin/op/__init__.py

@@ -2,7 +2,7 @@
 from kirin.prelude import structural_no_opt as _structural_no_opt
 from kirin.lowering import wraps as _wraps
 
-from . import stmts as stmts, types as types
+from . import stmts as stmts, types as types, rewrite as rewrite
 from .traits import Unitary as Unitary, MaybeUnitary as MaybeUnitary
 from ._dialect import dialect as dialect
 

src/bloqade/squin/op/rewrite.py

@@ -0,0 +1,48 @@
+"""Rewrite py.binop.mult to Mult stmt"""
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import Walk
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from .stmts import Mult, Scale
+from .types import OpType
+
+
+class _PyMultToSquinMult(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, py.Mult):
+            return RewriteResult()
+
+        lhs_is_op = node.lhs.type.is_subseteq(OpType)
+        rhs_is_op = node.lhs.type.is_subseteq(OpType)
+
+        if not lhs_is_op and not rhs_is_op:
+            return RewriteResult()
+
+        if lhs_is_op and rhs_is_op:
+            mult = Mult(node.lhs, node.rhs)
+            node.replace_by(mult)
+            return RewriteResult(has_done_something=True)
+
+        if lhs_is_op:
+            scale = Scale(node.lhs, node.rhs)
+            node.replace_by(scale)
+            return RewriteResult(has_done_something=True)
+
+        if rhs_is_op:
+            scale = Scale(node.rhs, node.lhs)
+            node.replace_by(scale)
+            return RewriteResult(has_done_something=True)
+
+        raise ValueError(
+            "Rewrite of py.binop.mult failed. This exception should not be reachable, please report this issue."
+        )
+
+
+class PyMultToSquinMult(Pass):
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        return Walk(_PyMultToSquinMult()).rewrite(mt.code)

src/bloqade/squin/op/types.py

@@ -6,5 +6,8 @@ class Op:
     def __matmul__(self, other: "Op") -> "Op":
         raise NotImplementedError("@ can only be used within a squin kernel program")
 
+    def __mul__(self, other: "Op") -> "Op":
+        raise NotImplementedError("@ can only be used within a squin kernel program")
+
 
 OpType = types.PyClass(Op)

test/squin/test_mult_rewrite.py

@@ -0,0 +1,39 @@
+from kirin.dialects import py, func
+
+from bloqade import squin
+
+
+def test_mult_rewrite():
+
+    @squin.kernel
+    def helper(x: squin.op.types.Op, y: squin.op.types.Op):
+        return x * y
+
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(1)
+        x = squin.op.x()
+        y = squin.op.y()
+        z = x * y
+        t = helper(x, z)
+
+        squin.qubit.apply(t, q)
+        return q
+
+    helper.print()
+
+    assert isinstance(helper.code, func.Function)
+
+    helper_stmts = list(helper.code.body.stmts())
+    assert len(helper_stmts) == 2  # [Mult(), Return()]
+    assert isinstance(helper_stmts[0], squin.op.stmts.Mult)
+
+    assert isinstance(main.code, func.Function)
+
+    count_mults_in_main = 0
+    for stmt in main.code.body.stmts():
+        assert not isinstance(stmt, py.Mult)
+
+        count_mults_in_main += isinstance(stmt, squin.op.stmts.Mult)
+
+    assert count_mults_in_main == 1