adding pyqrack impl into circuits repo

Author: weinbe58

_typos.toml

@@ -12,3 +12,4 @@ AttributeIDSupressMenu = "AttributeIDSupressMenu"
 Braket = "Braket"
 mch = "mch"
 IY = "IY"
+ket = "ket"

pyproject.toml

@@ -17,6 +17,8 @@ dependencies = [
     "rich>=13.9.4",
     "pydantic>=1.3.0,<2.11.0",
     "pandas>=2.2.3",
+    "pyqrack[pyqrack]>=1.38.2",
+    "pyqrack-cpu[pyqrack-cpu]>=1.38.2",
 ]
 
 [project.optional-dependencies]
@@ -76,6 +78,26 @@ doc = [
 examples = [
     "networkx>=3.4.2",
 ]
+dev-linux = [
+    "cirq-core[contrib]>=1.4.1",
+    "lark>=1.2.2",
+    "pyqrack-cpu>=1.38.2",
+    "qbraid>=0.9.5",
+    "ffmpeg>=1.4",
+    "matplotlib>=3.10.0",
+    "pyqt5>=5.15.11",
+    "tqdm>=4.67.1",
+]
+dev-mac-arm = [
+    "cirq-core[contrib]>=1.4.1",
+    "ffmpeg>=1.4",
+    "lark>=1.2.2",
+    "matplotlib>=3.10.0",
+    "pyqrack>=1.38.2",
+    "pyqt5>=5.15.11",
+    "qbraid>=0.9.5",
+    "tqdm>=4.67.1",
+]
 
 [tool.isort]
 profile = "black"

src/bloqade/pyqrack/__init__.py

@@ -0,0 +1,18 @@
+from .reg import (
+    CBitRef as CBitRef,
+    CRegister as CRegister,
+    PyQrackReg as PyQrackReg,
+    QubitState as QubitState,
+    Measurement as Measurement,
+    PyQrackQubit as PyQrackQubit,
+)
+from .base import (
+    StackMemory as StackMemory,
+    DynamicMemory as DynamicMemory,
+    PyQrackInterpreter as PyQrackInterpreter,
+)
+
+# 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, parallel as parallel
+from .target import PyQrack as PyQrack

src/bloqade/pyqrack/base.py

@@ -0,0 +1,131 @@
+import abc
+import typing
+from dataclasses import field, dataclass
+from unittest.mock import Mock
+
+import numpy as np
+from kirin.interp import Interpreter
+from typing_extensions import Self
+from kirin.interp.exceptions import InterpreterError
+
+from bloqade.pyqrack.reg import Measurement
+
+if typing.TYPE_CHECKING:
+    from pyqrack import QrackSimulator
+
+
+class PyQrackOptions(typing.TypedDict):
+    qubitCount: int
+    isTensorNetwork: bool
+    isSchmidtDecomposeMulti: bool
+    isSchmidtDecompose: bool
+    isStabilizerHybrid: bool
+    isBinaryDecisionTree: bool
+    isPaged: bool
+    isCpuGpuHybrid: bool
+    isOpenCL: bool
+
+
+def _default_pyqrack_args() -> PyQrackOptions:
+    return PyQrackOptions(
+        qubitCount=-1,
+        isTensorNetwork=False,
+        isSchmidtDecomposeMulti=True,
+        isSchmidtDecompose=True,
+        isStabilizerHybrid=True,
+        isBinaryDecisionTree=True,
+        isPaged=True,
+        isCpuGpuHybrid=True,
+        isOpenCL=True,
+    )
+
+
+@dataclass
+class MemoryABC(abc.ABC):
+    pyqrack_options: PyQrackOptions = field(default_factory=_default_pyqrack_args)
+    sim_reg: "QrackSimulator" = field(init=False)
+
+    @abc.abstractmethod
+    def allocate(self, n_qubits: int) -> tuple[int, ...]:
+        """Allocate `n_qubits` qubits and return their ids."""
+        ...
+
+    def reset(self):
+        """Reset the memory, releasing all qubits."""
+        from pyqrack import QrackSimulator
+
+        # do not reset the simulator it might be used by
+        # results of the simulation
+        self.sim_reg = QrackSimulator(**self.pyqrack_options)
+
+
+@dataclass
+class MockMemory(MemoryABC):
+    """Mock memory for testing purposes."""
+
+    allocated: int = field(init=False, default=0)
+
+    def allocate(self, n_qubits: int):
+        allocated = self.allocated + n_qubits
+        result = tuple(range(self.allocated, allocated))
+        self.allocated = allocated
+        return result
+
+    def reset(self):
+        self.allocated = 0
+        self.sim_reg = Mock()
+
+
+@dataclass
+class StackMemory(MemoryABC):
+    total: int = field(kw_only=True)
+    allocated: int = field(init=False, default=0)
+
+    def allocate(self, n_qubits: int):
+        curr_allocated = self.allocated
+        self.allocated += n_qubits
+
+        if self.allocated > self.total:
+            raise InterpreterError(
+                f"qubit allocation exceeds memory, "
+                f"{self.total} qubits, "
+                f"{self.allocated} allocated"
+            )
+
+        return tuple(range(curr_allocated, self.allocated))
+
+    def reset(self):
+        super().reset()
+        self.allocated = 0
+
+
+@dataclass
+class DynamicMemory(MemoryABC):
+    def __post_init__(self):
+        self.reset()
+
+        if self.sim_reg.is_tensor_network:
+            raise ValueError("DynamicMemory does not support tensor networks")
+
+    def allocate(self, n_qubits: int):
+        start = self.sim_reg.num_qubits()
+        for i in range(start, start + n_qubits):
+            self.sim_reg.allocate_qubit(i)
+
+        return tuple(range(start, start + n_qubits))
+
+
+@dataclass
+class PyQrackInterpreter(Interpreter):
+    keys = ["pyqrack", "main"]
+    memory: MemoryABC = field(kw_only=True)
+    rng_state: np.random.Generator = field(
+        default_factory=np.random.default_rng, kw_only=True
+    )
+    loss_m_result: Measurement = field(default=Measurement.One, kw_only=True)
+    """The value of a measurement result when a qubit is lost."""
+
+    def initialize(self) -> Self:
+        super().initialize()
+        self.memory.reset()  # reset allocated qubits
+        return self

src/bloqade/pyqrack/noise/__init__.py

@@ -0,0 +1,103 @@
+from typing import TYPE_CHECKING, List
+
+from kirin import interp
+
+from bloqade.noise import native
+from bloqade.pyqrack import PyQrackInterpreter, reg
+
+if TYPE_CHECKING:
+    from pyqrack import QrackSimulator
+
+
+@native.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+    def apply_pauli_error(
+        self,
+        interp: PyQrackInterpreter,
+        qarg: reg.PyQrackQubit,
+        px: float,
+        py: float,
+        pz: float,
+    ):
+        p = [1 - (px + py + pz), px, py, pz]
+
+        assert all(0 <= x <= 1 for x in p), "Invalid Pauli error probabilities"
+
+        which = interp.rng_state.choice(["i", "x", "y", "z"], p=p)
+
+        if which == "i":
+            return
+
+        getattr(qarg.sim_reg, which)(qarg.addr)
+
+    @interp.impl(native.PauliChannel)
+    def single_qubit_error_channel(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: native.PauliChannel,
+    ):
+        qargs: List[reg.PyQrackQubit] = frame.get(stmt.qargs)
+
+        active_qubits = (qarg for qarg in qargs if qarg.is_active())
+
+        for qarg in active_qubits:
+            self.apply_pauli_error(interp, qarg, stmt.px, stmt.py, stmt.pz)
+
+        return ()
+
+    @interp.impl(native.CZPauliChannel)
+    def cz_pauli_channel(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: native.CZPauliChannel,
+    ):
+
+        qargs: List[reg.PyQrackQubit] = frame.get(stmt.qargs)
+        ctrls: List[reg.PyQrackQubit] = frame.get(stmt.ctrls)
+
+        if stmt.paired:
+            valid_pairs = (
+                (ctrl, qarg)
+                for ctrl, qarg in zip(ctrls, qargs)
+                if ctrl.is_active() and qarg.is_active()
+            )
+        else:
+            valid_pairs = (
+                (ctrl, qarg)
+                for ctrl, qarg in zip(ctrls, qargs)
+                if ctrl.is_active() ^ qarg.is_active()
+            )
+
+        for ctrl, qarg in valid_pairs:
+            if ctrl.is_active():
+                self.apply_pauli_error(
+                    interp, ctrl, stmt.px_ctrl, stmt.py_ctrl, stmt.pz_ctrl
+                )
+
+            if qarg.is_active():
+                self.apply_pauli_error(
+                    interp, qarg, stmt.px_qarg, stmt.py_qarg, stmt.pz_qarg
+                )
+
+        return ()
+
+    @interp.impl(native.AtomLossChannel)
+    def atom_loss_channel(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: native.AtomLossChannel,
+    ):
+        qargs: List[reg.PyQrackQubit["QrackSimulator"]] = frame.get(stmt.qargs)
+
+        active_qubits = (qarg for qarg in qargs if qarg.is_active())
+
+        for qarg in active_qubits:
+            if interp.rng_state.uniform() <= stmt.prob:
+                sim_reg = qarg.ref.sim_reg
+                sim_reg.force_m(qarg.addr, 0)
+                qarg.drop()
+
+        return ()

src/bloqade/pyqrack/noise/native.py

@@ -0,0 +1,79 @@
+from kirin import interp
+
+from bloqade.pyqrack.reg import (
+    CBitRef,
+    CRegister,
+    PyQrackReg,
+    QubitState,
+    Measurement,
+    PyQrackQubit,
+)
+from bloqade.pyqrack.base import PyQrackInterpreter
+from bloqade.qasm2.dialects import core
+
+
+@core.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+
+    @interp.impl(core.QRegNew)
+    def qreg_new(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.QRegNew
+    ):
+        n_qubits: int = frame.get(stmt.n_qubits)
+        return (
+            PyQrackReg(
+                size=n_qubits,
+                sim_reg=interp.memory.sim_reg,
+                addrs=interp.memory.allocate(n_qubits),
+                qubit_state=[QubitState.Active] * n_qubits,
+            ),
+        )
+
+    @interp.impl(core.CRegNew)
+    def creg_new(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.CRegNew
+    ):
+        n_bits: int = frame.get(stmt.n_bits)
+        return (CRegister(size=n_bits),)
+
+    @interp.impl(core.QRegGet)
+    def qreg_get(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.QRegGet
+    ):
+        return (PyQrackQubit(ref=frame.get(stmt.reg), pos=frame.get(stmt.idx)),)
+
+    @interp.impl(core.CRegGet)
+    def creg_get(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.CRegGet
+    ):
+        return (CBitRef(ref=frame.get(stmt.reg), pos=frame.get(stmt.idx)),)
+
+    @interp.impl(core.Measure)
+    def measure(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.Measure
+    ):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        carg: CBitRef = frame.get(stmt.carg)
+        if qarg.is_active():
+            carg.set_value(Measurement(qarg.sim_reg.m(qarg.addr)))
+        else:
+            carg.set_value(interp.loss_m_result)
+
+        return ()
+
+    @interp.impl(core.Reset)
+    def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.Reset):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        qarg.sim_reg.force_m(qarg.addr, 0)
+        return ()
+
+    @interp.impl(core.CRegEq)
+    def creg_eq(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.CRegEq
+    ):
+        lhs: CRegister = frame.get(stmt.lhs)
+        rhs: CRegister = frame.get(stmt.rhs)
+        if len(lhs) != len(rhs):
+            return (False,)
+
+        return (all(left is right for left, right in zip(lhs, rhs)),)