Fix types in qbraid lowering

Author: david-pl

src/bloqade/qasm2/dialects/expr/stmts.py

@@ -87,7 +87,7 @@ def print_impl(self, printer: Printer) -> None:
 
 
 # QASM 2.0 arithmetic operations
-PyNum = types.Union(types.Int, types.Float)
+PyNum = types.TypeVar("PyNum", bound=types.Union(types.Int, types.Float))
 
 
 @statement(dialect=dialect)
@@ -110,7 +110,7 @@ class Sin(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the sine of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The sine of the number."""
 
 
@@ -122,7 +122,7 @@ class Cos(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the cosine of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The cosine of the number."""
 
 
@@ -134,7 +134,7 @@ class Tan(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the tangent of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The tangent of the number."""
 
 
@@ -146,7 +146,7 @@ class Exp(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the exponential of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The exponential of the number."""
 
 
@@ -158,7 +158,7 @@ class Log(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the natural log of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The natural log of the number."""
 
 
@@ -170,7 +170,7 @@ class Sqrt(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the square root of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The square root of the number."""
 
 

src/bloqade/qbraid/lowering.py

@@ -320,5 +320,6 @@ def lower_full_turns(self, value: float) -> ir.SSAValue:
         self.block_list.append(const_pi)
         turns = self.lower_number(2 * value)
         mul = qasm2.expr.Mul(const_pi.result, turns)
+        mul.result.type = types.Float
         self.block_list.append(mul)
         return mul.result

src/bloqade/qbraid/schema.py

@@ -238,13 +238,13 @@ def decompiled_circuit(self) -> str:
             str: The decompiled circuit from hardware execution.
 
         """
-        from bloqade.noise import native
         from bloqade.qasm2.emit import QASM2
         from bloqade.qasm2.passes import glob, parallel
+        from bloqade.qasm2.rewrite.noise import remove_noise
 
         mt = self.lower_noise_model("method")
 
-        native.RemoveNoisePass(mt.dialects)(mt)
+        remove_noise.RemoveNoisePass(mt.dialects)(mt)
         parallel.ParallelToUOp(mt.dialects)(mt)
         glob.GlobalToUOP(mt.dialects)(mt)
         return QASM2(qelib1=True).emit_str(mt)

test/qbraid/test_lowering.py

@@ -56,11 +56,8 @@ def run_assert(noise_model: schema.NoiseModel, expected_stmts: List[ir.Statement
     )
 
     expected_mt = ir.Method(
-        mod=None,
-        py_func=None,
         dialects=lowering.qbraid_noise,
         sym_name="test",
-        arg_names=[],
         code=expected_func_stmt,
     )
 
@@ -242,7 +239,10 @@ def test_lowering_global_w():
         (lam_num := as_float(2 * -(0.5 + phi_val))),
         (lam := qasm2.expr.Mul(pi_lam.result, lam_num.result)),
         parallel.UGate(
-            theta=theta.result, phi=phi.result, lam=lam.result, qargs=qargs.result
+            theta=ir.ResultValue(theta, 0, type=types.Float),
+            phi=ir.ResultValue(phi, 0, type=types.Float),
+            lam=ir.ResultValue(lam, 0, type=types.Float),
+            qargs=qargs.result,
         ),
         func.Return(creg.result),
     ]
@@ -304,7 +304,10 @@ def test_lowering_local_w():
         (lam_num := as_float(2 * -(0.5 + phi_val))),
         (lam := qasm2.expr.Mul(pi_lam.result, lam_num.result)),
         parallel.UGate(
-            qargs=qargs.result, theta=theta.result, phi=phi.result, lam=lam.result
+            qargs=qargs.result,
+            theta=ir.ResultValue(theta, 0, type=types.Float),
+            phi=ir.ResultValue(phi, 0, type=types.Float),
+            lam=ir.ResultValue(lam, 0, type=types.Float),
         ),
         func.Return(creg.result),
     ]
@@ -348,7 +351,9 @@ def test_lowering_global_rz():
         (theta_pi := qasm2.expr.ConstPI()),
         (theta_num := as_float(2 * phi_val)),
         (theta := qasm2.expr.Mul(theta_pi.result, theta_num.result)),
-        parallel.RZ(theta=theta.result, qargs=qargs.result),
+        parallel.RZ(
+            theta=ir.ResultValue(theta, 0, type=types.Float), qargs=qargs.result
+        ),
         func.Return(creg.result),
     ]
 
@@ -401,7 +406,9 @@ def test_lowering_local_rz():
         (theta_pi := qasm2.expr.ConstPI()),
         (theta_num := as_float(2 * phi_val)),
         (theta := qasm2.expr.Mul(theta_pi.result, theta_num.result)),
-        parallel.RZ(theta=theta.result, qargs=qargs.result),
+        parallel.RZ(
+            theta=ir.ResultValue(theta, 0, type=types.Float), qargs=qargs.result
+        ),
         func.Return(creg.result),
     ]