Fix Scale rewrite

Author: david-pl

src/bloqade/squin/op/rewrite.py

@@ -6,7 +6,7 @@
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from .stmts import Mult, Scale
+from .stmts import Mult, Scale, Operator
 from .types import OpType
 
 
@@ -17,11 +17,17 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
             return RewriteResult()
 
         lhs_is_op = node.lhs.type.is_subseteq(OpType)
-        rhs_is_op = node.lhs.type.is_subseteq(OpType)
+        rhs_is_op = node.rhs.type.is_subseteq(OpType)
 
         if not lhs_is_op and not rhs_is_op:
             return RewriteResult()
 
+        if isinstance(node.lhs, ir.ResultValue):
+            lhs_is_op = isinstance(node.lhs.stmt, Operator)
+
+        if isinstance(node.rhs, ir.ResultValue):
+            rhs_is_op = isinstance(node.rhs.stmt, Operator)
+
         if lhs_is_op and rhs_is_op:
             mult = Mult(node.lhs, node.rhs)
             node.replace_by(mult)

src/bloqade/squin/op/types.py

@@ -1,3 +1,5 @@
+from typing import overload
+
 from kirin import types
 
 
@@ -6,7 +8,16 @@ 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":
+    @overload
+    def __mul__(self, other: "Op") -> "Op": ...
+
+    @overload
+    def __mul__(self, other: int | float | complex) -> "Op": ...
+
+    def __mul__(self, other) -> "Op":
+        raise NotImplementedError("@ can only be used within a squin kernel program")
+
+    def __rmul__(self, other: int | float | complex) -> "Op":
         raise NotImplementedError("@ can only be used within a squin kernel program")
 
 

test/squin/test_mult_rewrite.py

@@ -21,6 +21,7 @@ def main():
         return q
 
     helper.print()
+    main.print()
 
     assert isinstance(helper.code, func.Function)
 
@@ -37,3 +38,82 @@ def main():
         count_mults_in_main += isinstance(stmt, squin.op.stmts.Mult)
 
     assert count_mults_in_main == 1
+
+
+def test_scale_rewrite():
+
+    @squin.kernel
+    def simple_rmul():
+        x = squin.op.x()
+        y = 2 * x
+        return y
+
+    simple_rmul.print()
+
+    assert isinstance(simple_rmul.code, func.Function)
+
+    simple_rmul_stmts = list(simple_rmul.code.body.stmts())
+    assert any(
+        map(lambda stmt: isinstance(stmt, squin.op.stmts.Scale), simple_rmul_stmts)
+    )
+    assert not any(
+        map(lambda stmt: isinstance(stmt, squin.op.stmts.Mult), simple_rmul_stmts)
+    )
+    assert not any(map(lambda stmt: isinstance(stmt, py.Mult), simple_rmul_stmts))
+
+    @squin.kernel
+    def simple_lmul():
+        x = squin.op.x()
+        y = x * 2
+        return y
+
+    simple_lmul.print()
+
+    assert isinstance(simple_lmul.code, func.Function)
+
+    simple_lmul_stmts = list(simple_lmul.code.body.stmts())
+    assert any(
+        map(lambda stmt: isinstance(stmt, squin.op.stmts.Scale), simple_lmul_stmts)
+    )
+    assert not any(
+        map(lambda stmt: isinstance(stmt, squin.op.stmts.Mult), simple_lmul_stmts)
+    )
+    assert not any(map(lambda stmt: isinstance(stmt, py.Mult), simple_lmul_stmts))
+
+    @squin.kernel
+    def scale_mult():
+        x = squin.op.x()
+        y = squin.op.y()
+        return 2 * (x * y)
+
+    assert isinstance(scale_mult.code, func.Function)
+
+    scale_mult_stmts = list(scale_mult.code.body.stmts())
+    assert (
+        sum(map(lambda stmt: isinstance(stmt, squin.op.stmts.Scale), scale_mult_stmts))
+        == 1
+    )
+    assert (
+        sum(map(lambda stmt: isinstance(stmt, squin.op.stmts.Mult), scale_mult_stmts))
+        == 1
+    )
+
+    @squin.kernel
+    def scale_mult2():
+        x = squin.op.x()
+        y = squin.op.y()
+        return 2 * x * y
+
+    scale_mult2.print()
+
+    assert isinstance(scale_mult2.code, func.Function)
+
+    scale_mult2_stmts = list(scale_mult2.code.body.stmts())
+    assert (
+        sum(map(lambda stmt: isinstance(stmt, squin.op.stmts.Scale), scale_mult2_stmts))
+        == 1
+    )
+    assert (
+        sum(map(lambda stmt: isinstance(stmt, squin.op.stmts.Mult), scale_mult2_stmts))
+        == 1
+    )