get pyqrack runtime to work

Author: johnzl-777

src/bloqade/pyqrack/squin/op.py

@@ -10,6 +10,7 @@
     RotRuntime,
     KronRuntime,
     MultRuntime,
+    ResetRuntime,
     ScaleRuntime,
     AdjointRuntime,
     ControlRuntime,
@@ -94,6 +95,12 @@ def phaseop(
         global_ = isinstance(stmt, op.stmts.PhaseOp)
         return (PhaseOpRuntime(theta, global_=global_),)
 
+    @interp.impl(op.stmts.Reset)
+    def reset(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Reset
+    ) -> tuple[OperatorRuntimeABC]:
+        return (ResetRuntime(),)
+
     @interp.impl(op.stmts.X)
     @interp.impl(op.stmts.Y)
     @interp.impl(op.stmts.Z)

src/bloqade/pyqrack/squin/qubit.py

@@ -60,6 +60,8 @@ def measure_qubit(
         result = self._measure_qubit(qbit)
         return (result,)
 
+
+"""
     @interp.impl(qubit.MeasureAndReset)
     def measure_and_reset(
         self,
@@ -83,3 +85,4 @@ def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Res
         qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
         for qbit in qubits:
             qbit.sim_reg.force_m(qbit.addr, 0)
+"""

src/bloqade/pyqrack/squin/runtime.py

@@ -41,6 +41,21 @@ def broadcast_apply(self, qubits: ilist.IList[PyQrackQubit, Any], **kwargs) -> N
             self.apply(*targets, **kwargs)
 
 
+@dataclass(frozen=True)
+class ResetRuntime(OperatorRuntimeABC):
+    """Reset the qubit to |0> state"""
+
+    @property
+    def n_sites(self) -> int:
+        return 1
+
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        for qubit in qubits:
+            if not qubit.is_active():
+                continue
+            qubit.sim_reg.force_m(qubit.addr, 0)
+
+
 @dataclass(frozen=True)
 class OperatorRuntime(OperatorRuntimeABC):
     method_name: str

src/bloqade/squin/op/__init__.py

@@ -105,6 +105,10 @@ def u(theta: float, phi: float, lam: float) -> types.Op: ...
 def pauli_string(*, string: str) -> types.Op: ...
 
 
+@_wraps(stmts.Reset)
+def reset() -> types.Op: ...
+
+
 # stdlibs
 @_ir.dialect_group(_structural_no_opt.add(dialect))
 def op(self):

test/pyqrack/test_squin.py

@@ -28,25 +28,14 @@ def new():
     def m():
         q = squin.qubit.new(3)
         m = squin.qubit.measure(q)
-        squin.qubit.reset(q)
+        squin.qubit.apply(squin.op.reset(), q)
         return m
 
     target = PyQrack(3)
     result = target.run(m)
     assert isinstance(result, ilist.IList)
     assert result.data == [0, 0, 0]
 
-    @squin.kernel
-    def measure_and_reset():
-        q = squin.qubit.new(3)
-        m = squin.qubit.measure_and_reset(q)
-        return m
-
-    target = PyQrack(3)
-    result = target.run(measure_and_reset)
-    assert isinstance(result, ilist.IList)
-    assert result.data == [0, 0, 0]
-
 
 def test_x():
     @squin.kernel