Fix type inference

Author: david-pl

src/bloqade/squin/groups.py

@@ -18,11 +18,12 @@ def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
         if fold:
             fold_pass.fixpoint(method)
 
+        py_mult_to_mult_pass(method)
+
         if typeinfer:
             typeinfer_pass(method)
 
         ilist_desugar_pass(method)
-        py_mult_to_mult_pass(method)
 
         if typeinfer:
             typeinfer_pass(method)  # fix types after desugaring

src/bloqade/squin/op/complex.py

@@ -0,0 +1,5 @@
+import numbers
+
+from kirin.ir.attrs.types import PyClass
+
+Number = PyClass(numbers.Number)

src/bloqade/squin/op/number.py

@@ -6,7 +6,7 @@
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from .stmts import Mult, Scale, Operator
+from .stmts import Mult, Scale
 from .types import OpType
 
 
@@ -22,12 +22,6 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         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/rewrite.py

@@ -2,8 +2,8 @@
 from kirin.decl import info, statement
 
 from .types import OpType
+from .number import Number
 from .traits import Unitary, HasSites, FixedSites, MaybeUnitary
-from .complex import Complex
 from ._dialect import dialect
 
 
@@ -54,7 +54,7 @@ class Scale(CompositeOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
     is_unitary: bool = info.attribute(default=False)
     op: ir.SSAValue = info.argument(OpType)
-    factor: ir.SSAValue = info.argument(Complex)
+    factor: ir.SSAValue = info.argument(Number)
     result: ir.ResultValue = info.result(OpType)
 
 

src/bloqade/squin/op/stmts.py

@@ -12,12 +12,12 @@ def __matmul__(self, other: "Op") -> "Op":
     def __mul__(self, other: "Op") -> "Op": ...
 
     @overload
-    def __mul__(self, other: int | float | complex) -> "Op": ...
+    def __mul__(self, other: 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":
+    def __rmul__(self, other: complex) -> "Op":
         raise NotImplementedError("@ can only be used within a squin kernel program")
 
 

src/bloqade/squin/op/types.py

@@ -1,3 +1,4 @@
+from kirin.types import PyClass
 from kirin.dialects import py, func
 
 from bloqade import squin
@@ -117,3 +118,42 @@ def scale_mult2():
         sum(map(lambda stmt: isinstance(stmt, squin.op.stmts.Mult), scale_mult2_stmts))
         == 1
     )
+
+
+def test_scale_types():
+    @squin.kernel
+    def simple_lmul():
+        x = squin.op.x()
+        y = x * (2 + 0j)
+        return y
+
+    @squin.kernel
+    def simple_rmul():
+        x = squin.op.x()
+        y = 2.1 * x
+        return y
+
+    @squin.kernel
+    def nested_rmul():
+        x = squin.op.x()
+        y = squin.op.y()
+        return 2 * x * y
+
+    @squin.kernel
+    def nested_lmul():
+        x = squin.op.x()
+        y = squin.op.y()
+        return x * y * 2.0j
+
+    def check_stmt_type(code, typ):
+        for stmt in code.body.stmts():
+            if isinstance(stmt, func.Return):
+                continue
+            is_op = stmt.result.type.is_subseteq(squin.op.types.OpType)
+            is_num = stmt.result.type.is_equal(PyClass(typ))
+            assert is_op or is_num
+
+    check_stmt_type(simple_lmul.code, complex)
+    check_stmt_type(simple_rmul.code, float)
+    check_stmt_type(nested_rmul.code, int)
+    check_stmt_type(nested_lmul.code, complex)