Implement mult runtime -- TODO: don't wrap mult

Author: david-pl

src/bloqade/pyqrack/squin/op.py

@@ -5,7 +5,13 @@
 # from bloqade.pyqrack.reg import QubitState, PyQrackQubit
 from bloqade.pyqrack.base import PyQrackInterpreter
 
-from .runtime import ControlRuntime, IdentityRuntime, OperatorRuntime, ProjectorRuntime
+from .runtime import (
+    MultRuntime,
+    ControlRuntime,
+    IdentityRuntime,
+    OperatorRuntime,
+    ProjectorRuntime,
+)
 
 # from kirin.dialects import ilist
 
@@ -19,11 +25,13 @@ class PyQrackMethods(interp.MethodTable):
     # ):
     #     is_unitary: bool = info.attribute(default=False)
 
-    # @interp.impl(op.stmts.Mult)
-    # def mult(
-    #     self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Mult
-    # ):
-    #     is_unitary: bool = info.attribute(default=False)
+    @interp.impl(op.stmts.Mult)
+    def mult(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Mult
+    ):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        return (MultRuntime(lhs, rhs),)
 
     # @interp.impl(op.stmts.Adjoint)
     # def adjoint(

src/bloqade/pyqrack/squin/runtime.py

@@ -15,10 +15,7 @@ def apply(self, *qubits: PyQrackQubit) -> None:
 class OperatorRuntime(OperatorRuntimeABC):
     method_name: str
 
-    def apply(
-        self,
-        *qubits: PyQrackQubit,
-    ) -> None:
+    def apply(self, *qubits: PyQrackQubit) -> None:
         getattr(qubits[-1].sim_reg, self.method_name)(qubits[-1].addr)
 
 
@@ -27,10 +24,7 @@ class ControlRuntime(OperatorRuntimeABC):
     method_name: str
     n_controls: int
 
-    def apply(
-        self,
-        *qubits: PyQrackQubit,
-    ) -> None:
+    def apply(self, *qubits: PyQrackQubit) -> None:
         # NOTE: this is a bit odd, since you can "skip" qubits by making n_controls < len(qubits)
         ctrls = [qbit.addr for qbit in qubits[: self.n_controls]]
         target = qubits[-1]
@@ -41,10 +35,7 @@ def apply(
 class ProjectorRuntime(OperatorRuntimeABC):
     to_state: bool
 
-    def apply(
-        self,
-        *qubits: PyQrackQubit,
-    ) -> None:
+    def apply(self, *qubits: PyQrackQubit) -> None:
         qubits[-1].sim_reg.force_m(qubits[-1].addr, self.to_state)
 
 
@@ -55,3 +46,13 @@ class IdentityRuntime(OperatorRuntimeABC):
 
     def apply(self, *qubits: PyQrackQubit) -> None:
         pass
+
+
+@dataclass
+class MultRuntime(OperatorRuntimeABC):
+    lhs: OperatorRuntimeABC
+    rhs: OperatorRuntimeABC
+
+    def apply(self, *qubits: PyQrackQubit) -> None:
+        self.rhs.apply(*qubits)
+        self.lhs.apply(*qubits)

src/bloqade/squin/op/__init__.py

@@ -11,6 +11,11 @@
 def kron(lhs: types.Op, rhs: types.Op) -> types.Op: ...
 
 
+# FIXME: should we just rewrite the py.binop.mult instead?
+@_wraps(stmts.Mult)
+def mult(lhs: types.Op, rhs: types.Op) -> types.Op: ...
+
+
 @_wraps(stmts.Adjoint)
 def adjoint(op: types.Op) -> types.Op: ...
 

test/pyqrack/test_squin.py

@@ -101,8 +101,26 @@ def main():
     target.run(main)
 
 
+def test_mult():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(1)
+        x = squin.op.x()
+        id = squin.op.mult(x, x)
+        squin.qubit.apply(id, q)
+        return squin.qubit.measure(q)
+
+    main.print()
+
+    target = PyQrack(1)
+    result = target.run(main)
+
+    assert result == [0]
+
+
 # TODO: remove
-test_qubit()
-test_x()
-test_basic_ops("x")
-test_cx()
+# test_qubit()
+# test_x()
+# test_basic_ops("x")
+# test_cx()
+# test_mult()