Draft implementation for wire dialect

Author: david-pl

src/bloqade/pyqrack/__init__.py

@@ -3,6 +3,7 @@
     CRegister as CRegister,
     QubitState as QubitState,
     Measurement as Measurement,
+    PyQrackWire as PyQrackWire,
     PyQrackQubit as PyQrackQubit,
 )
 from .base import (

src/bloqade/pyqrack/reg.py

@@ -70,3 +70,8 @@ def is_active(self) -> bool:
     def drop(self):
         """Drop the qubit in-place."""
         self.state = QubitState.Lost
+
+
+@dataclass
+class PyQrackWire:
+    qubit: PyQrackQubit

src/bloqade/pyqrack/squin/wire.py

@@ -0,0 +1,69 @@
+from kirin import interp
+
+from bloqade.squin import wire
+from bloqade.pyqrack.reg import PyQrackWire, PyQrackQubit
+from bloqade.pyqrack.base import PyQrackInterpreter
+
+from .runtime import OperatorRuntimeABC
+
+
+@wire.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+    # @interp.impl(wire.Wrap)
+    # def wrap(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: wire.Wrap):
+    #     traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
+    #     wire: ir.SSAValue = info.argument(WireType)
+    #     qubit: ir.SSAValue = info.argument(QubitType)
+
+    @interp.impl(wire.Unwrap)
+    def unwrap(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: wire.Unwrap
+    ):
+        q: PyQrackQubit = frame.get(stmt.qubit)
+        return (PyQrackWire(q),)
+
+    @interp.impl(wire.Apply)
+    def apply(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: wire.Apply):
+        ws = stmt.inputs
+        assert isinstance(ws, tuple)
+        qubits: list[PyQrackQubit] = []
+        for w in ws:
+            assert isinstance(w, PyQrackWire)
+            qubits.append(w.qubit)
+        op: OperatorRuntimeABC = frame.get(stmt.operator)
+
+        op.apply(*qubits)
+
+        out_ws = [PyQrackWire(qbit) for qbit in qubits]
+        return (out_ws,)
+
+    @interp.impl(wire.Measure)
+    def measure(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: wire.Measure
+    ):
+        w: PyQrackWire = frame.get(stmt.wire)
+        qbit = w.qubit
+        res: int = qbit.sim_reg.m(qbit.addr)
+        return (res,)
+
+    @interp.impl(wire.MeasureAndReset)
+    def measure_and_reset(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: wire.MeasureAndReset,
+    ):
+        w: PyQrackWire = frame.get(stmt.wire)
+        qbit = w.qubit
+        res: int = qbit.sim_reg.m(qbit.addr)
+        qbit.sim_reg.force_m(qbit.addr, False)
+        new_w = PyQrackWire(qbit)
+        return (new_w, res)
+
+    @interp.impl(wire.Reset)
+    def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: wire.Reset):
+        w: PyQrackWire = frame.get(stmt.wire)
+        qbit = w.qubit
+        qbit.sim_reg.force_m(qbit.addr, False)
+        new_w = PyQrackWire(qbit)
+        return (new_w,)

src/bloqade/squin/wire.py

@@ -8,10 +8,11 @@
 
 from kirin import ir, types, interp, lowering
 from kirin.decl import info, statement
+from kirin.lowering import wraps
 
-from bloqade.types import QubitType
+from bloqade.types import Qubit, QubitType
 
-from .op.types import OpType
+from .op.types import Op, OpType
 
 # from kirin.lowering import wraps
 
@@ -101,3 +102,11 @@ class ConstPropWire(interp.MethodTable):
     def apply(self, interp, frame, stmt: Apply):
 
         return frame.get_values(stmt.inputs)
+
+
+@wraps(Unwrap)
+def unwrap(qubit: Qubit) -> Wire: ...
+
+
+@wraps(Apply)
+def apply(op: Op, w: Wire) -> Wire: ...

test/pyqrack/test_squin.py

@@ -4,7 +4,7 @@
 from kirin.dialects import ilist
 
 from bloqade import squin
-from bloqade.pyqrack import PyQrack, PyQrackQubit
+from bloqade.pyqrack import PyQrack, PyQrackWire, PyQrackQubit
 
 
 def test_qubit():
@@ -311,6 +311,20 @@ def main():
     assert result == [1, 1, 1]
 
 
+def test_wire():
+    @squin.wired
+    def main():
+        w = squin.wire.unwrap(1)
+        x = squin.op.x()
+        squin.wire.apply(x, w)
+        return w
+
+    target = PyQrack(1)
+    result = target.run(main)
+    assert isinstance(result, PyQrackWire)
+    assert result.qubit.sim_reg.out_ket() == [0, 1]
+
+
 # TODO: remove
 # test_qubit()
 # test_x()
@@ -329,3 +343,4 @@ def main():
 # test_broadcast()
 # test_u3()
 # test_clifford_str()
+# test_wire()