Implement pyqrack interpreter methods for squin dialect

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 (
@@ -14,4 +15,5 @@
 # NOTE: The following import is for registering the method tables
 from .noise import native as native
 from .qasm2 import uop as uop, core as core, glob as glob, parallel as parallel
+from .squin import op as op, qubit as qubit
 from .target import PyQrack as PyQrack

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/op.py

@@ -0,0 +1,154 @@
+from kirin import interp
+
+from bloqade.squin import op
+from bloqade.pyqrack.base import PyQrackInterpreter
+
+from .runtime import (
+    SnRuntime,
+    SpRuntime,
+    U3Runtime,
+    RotRuntime,
+    KronRuntime,
+    MultRuntime,
+    ScaleRuntime,
+    AdjointRuntime,
+    ControlRuntime,
+    PhaseOpRuntime,
+    IdentityRuntime,
+    OperatorRuntime,
+    ProjectorRuntime,
+    OperatorRuntimeABC,
+    PauliStringRuntime,
+)
+
+
+@op.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+
+    @interp.impl(op.stmts.Kron)
+    def kron(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Kron
+    ) -> tuple[OperatorRuntimeABC]:
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        return (KronRuntime(lhs, rhs),)
+
+    @interp.impl(op.stmts.Mult)
+    def mult(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Mult
+    ) -> tuple[OperatorRuntimeABC]:
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        return (MultRuntime(lhs, rhs),)
+
+    @interp.impl(op.stmts.Adjoint)
+    def adjoint(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Adjoint
+    ) -> tuple[OperatorRuntimeABC]:
+        op = frame.get(stmt.op)
+        return (AdjointRuntime(op),)
+
+    @interp.impl(op.stmts.Scale)
+    def scale(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Scale
+    ) -> tuple[OperatorRuntimeABC]:
+        op = frame.get(stmt.op)
+        factor = frame.get(stmt.factor)
+        return (ScaleRuntime(op, factor),)
+
+    @interp.impl(op.stmts.Control)
+    def control(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Control
+    ) -> tuple[OperatorRuntimeABC]:
+        op = frame.get(stmt.op)
+        n_controls = stmt.n_controls
+        rt = ControlRuntime(
+            op=op,
+            n_controls=n_controls,
+        )
+        return (rt,)
+
+    @interp.impl(op.stmts.Rot)
+    def rot(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Rot
+    ) -> tuple[OperatorRuntimeABC]:
+        axis = frame.get(stmt.axis)
+        angle = frame.get(stmt.angle)
+        return (RotRuntime(axis, angle),)
+
+    @interp.impl(op.stmts.Identity)
+    def identity(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Identity
+    ) -> tuple[OperatorRuntimeABC]:
+        return (IdentityRuntime(sites=stmt.sites),)
+
+    @interp.impl(op.stmts.PhaseOp)
+    @interp.impl(op.stmts.ShiftOp)
+    def phaseop(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: op.stmts.PhaseOp | op.stmts.ShiftOp,
+    ) -> tuple[OperatorRuntimeABC]:
+        theta = frame.get(stmt.theta)
+        global_ = isinstance(stmt, op.stmts.PhaseOp)
+        return (PhaseOpRuntime(theta, global_=global_),)
+
+    @interp.impl(op.stmts.X)
+    @interp.impl(op.stmts.Y)
+    @interp.impl(op.stmts.Z)
+    @interp.impl(op.stmts.H)
+    @interp.impl(op.stmts.S)
+    @interp.impl(op.stmts.T)
+    def operator(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: (
+            op.stmts.X | op.stmts.Y | op.stmts.Z | op.stmts.H | op.stmts.S | op.stmts.T
+        ),
+    ) -> tuple[OperatorRuntimeABC]:
+        return (OperatorRuntime(method_name=stmt.name.lower()),)
+
+    @interp.impl(op.stmts.P0)
+    @interp.impl(op.stmts.P1)
+    def projector(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: op.stmts.P0 | op.stmts.P1,
+    ) -> tuple[OperatorRuntimeABC]:
+        state = isinstance(stmt, op.stmts.P1)
+        return (ProjectorRuntime(to_state=state),)
+
+    @interp.impl(op.stmts.Sp)
+    def sp(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Sp
+    ) -> tuple[OperatorRuntimeABC]:
+        return (SpRuntime(),)
+
+    @interp.impl(op.stmts.Sn)
+    def sn(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Sn
+    ) -> tuple[OperatorRuntimeABC]:
+        return (SnRuntime(),)
+
+    @interp.impl(op.stmts.U3)
+    def u3(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.U3
+    ) -> tuple[OperatorRuntimeABC]:
+        theta = frame.get(stmt.theta)
+        phi = frame.get(stmt.phi)
+        lam = frame.get(stmt.lam)
+        return (U3Runtime(theta, phi, lam),)
+
+    @interp.impl(op.stmts.PauliString)
+    def clifford_string(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: op.stmts.PauliString,
+    ) -> tuple[OperatorRuntimeABC]:
+        string = stmt.string
+        ops = [OperatorRuntime(method_name=name.lower()) for name in stmt.string]
+        return (PauliStringRuntime(string, ops),)

src/bloqade/pyqrack/squin/qubit.py

@@ -0,0 +1,111 @@
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+from kirin.interp.exceptions import InterpreterError
+
+from bloqade.squin import qubit
+from bloqade.pyqrack.reg import QubitState, PyQrackQubit
+from bloqade.pyqrack.base import PyQrackInterpreter
+
+from .runtime import OperatorRuntimeABC
+
+
+@qubit.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+    @interp.impl(qubit.New)
+    def new(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.New):
+        n_qubits: int = frame.get(stmt.n_qubits)
+        qreg = ilist.IList(
+            [
+                PyQrackQubit(i, interp.memory.sim_reg, QubitState.Active)
+                for i in interp.memory.allocate(n_qubits=n_qubits)
+            ]
+        )
+        return (qreg,)
+
+    @interp.impl(qubit.Apply)
+    def apply(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Apply):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+
+        for qbit in qubits:
+            if not qbit.is_active():
+                return ()
+
+        operator: OperatorRuntimeABC = frame.get(stmt.operator)
+        operator.apply(*qubits)
+
+    @interp.impl(qubit.Broadcast)
+    def broadcast(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Broadcast
+    ):
+        operator: OperatorRuntimeABC = frame.get(stmt.operator)
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        operator.broadcast_apply(qubits)
+
+    def _measure_qubit(self, qbit: PyQrackQubit):
+        if qbit.is_active():
+            return qbit.sim_reg.m(qbit.addr)
+
+    @interp.impl(qubit.MeasureQubitList)
+    def measure_qubit_list(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: qubit.MeasureQubitList,
+    ):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        result = [self._measure_qubit(qbit) for qbit in qubits]
+        return (result,)
+
+    @interp.impl(qubit.MeasureQubit)
+    def measure_qubit(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.MeasureQubit
+    ):
+        qbit: PyQrackQubit = frame.get(stmt.qubit)
+        result = self._measure_qubit(qbit)
+        return (result,)
+
+    @interp.impl(qubit.MeasureAny)
+    def measure_any(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.MeasureAny
+    ):
+        input = frame.get(stmt.input)
+
+        if isinstance(input, PyQrackQubit) and input.is_active():
+            result = self._measure_qubit
+        elif isinstance(input, ilist.IList):
+            result = []
+            for qbit in input:
+                if not isinstance(qbit, PyQrackQubit):
+                    raise InterpreterError(f"Cannot measure {type(qbit).__name__}")
+
+                result.append(self._measure_qubit(qbit))
+        else:
+            raise InterpreterError(f"Cannot measure {type(input).__name__}")
+
+        return (result,)
+
+    @interp.impl(qubit.MeasureAndReset)
+    def measure_and_reset(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: qubit.MeasureAndReset,
+    ):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        result = []
+        for qbit in qubits:
+            if qbit.is_active():
+                result.append(qbit.sim_reg.m(qbit.addr))
+            else:
+                result.append(None)
+            qbit.sim_reg.force_m(qbit.addr, 0)
+
+        return (result,)
+
+    @interp.impl(qubit.Reset)
+    def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Reset):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            qbit.sim_reg.force_m(qbit.addr, 0)