bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/pyqrack/qasm2/core.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)),)

bloqade/pyqrack/qasm2/parallel.py

@@ -0,0 +1,46 @@
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from bloqade.pyqrack.reg import PyQrackQubit
+from bloqade.pyqrack.base import PyQrackInterpreter
+from bloqade.qasm2.dialects import parallel
+
+
+@parallel.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+
+    @interp.impl(parallel.CZ)
+    def cz(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: parallel.CZ):
+
+        qargs: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qargs)
+        ctrls: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.ctrls)
+        for qarg, ctrl in zip(qargs, ctrls):
+            if qarg.is_active() and ctrl.is_active():
+                interp.memory.sim_reg.mcz([ctrl.addr], qarg.addr)
+        return ()
+
+    @interp.impl(parallel.UGate)
+    def ugate(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: parallel.UGate
+    ):
+        qargs: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qargs)
+        theta, phi, lam = (
+            frame.get(stmt.theta),
+            frame.get(stmt.phi),
+            frame.get(stmt.lam),
+        )
+        for qarg in qargs:
+            if qarg.is_active():
+                interp.memory.sim_reg.u(qarg.addr, theta, phi, lam)
+        return ()
+
+    @interp.impl(parallel.RZ)
+    def rz(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: parallel.RZ):
+        qargs: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qargs)
+        phi = frame.get(stmt.theta)
+        for qarg in qargs:
+            if qarg.is_active():
+                interp.memory.sim_reg.r(3, phi, qarg.addr)
+        return ()

bloqade/pyqrack/qasm2/uop.py

@@ -0,0 +1,247 @@
+import math
+
+from kirin import interp
+
+from bloqade.pyqrack.reg import PyQrackQubit
+from bloqade.qasm2.dialects import uop
+
+
+@uop.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+    GATE_TO_METHOD = {
+        "x": "x",
+        "y": "y",
+        "z": "z",
+        "h": "h",
+        "s": "s",
+        "t": "t",
+        "cx": "mcx",
+        "CX": "mcx",
+        "cz": "mcz",
+        "cy": "mcy",
+        "ch": "mch",
+        "sdag": "adjs",
+        "sdg": "adjs",
+        "tdag": "adjt",
+        "tdg": "adjt",
+    }
+
+    AXIS_MAP = {"rx": 1, "ry": 2, "rz": 3, "crx": 1, "cry": 2, "crz": 3}
+
+    @interp.impl(uop.Barrier)
+    def barrier(
+        self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.Barrier
+    ):
+        return ()
+
+    @interp.impl(uop.X)
+    @interp.impl(uop.Y)
+    @interp.impl(uop.Z)
+    @interp.impl(uop.H)
+    @interp.impl(uop.S)
+    @interp.impl(uop.Sdag)
+    @interp.impl(uop.T)
+    @interp.impl(uop.Tdag)
+    def single_qubit_gate(
+        self,
+        interp: interp.Interpreter,
+        frame: interp.Frame,
+        stmt: uop.SingleQubitGate,
+    ):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active():
+            getattr(qarg.sim_reg, self.GATE_TO_METHOD[stmt.name])(qarg.addr)
+        return ()
+
+    @interp.impl(uop.UGate)
+    def ugate(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.UGate):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active():
+            qarg.sim_reg.u(
+                qarg.addr,
+                frame.get(stmt.theta),
+                frame.get(stmt.phi),
+                frame.get(stmt.lam),
+            )
+        return ()
+
+    @interp.impl(uop.Id)
+    def id(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.Id):
+        return ()
+
+    @interp.impl(uop.SX)
+    def sx(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.SX):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active():
+            qarg.sim_reg.u(qarg.addr, math.pi / 2, math.pi / 2, -math.pi / 2)
+        return ()
+
+    @interp.impl(uop.SXdag)
+    def sx_dag(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.SX):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active():
+            qarg.sim_reg.u(qarg.addr, math.pi * (1.5), math.pi / 2, math.pi / 2)
+        return ()
+
+    @interp.impl(uop.CSX)
+    def csx(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.CSX):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        ctrl: PyQrackQubit = frame.get(stmt.ctrl)
+        if qarg.is_active() and ctrl.is_active():
+            qarg.sim_reg.mcu(
+                [ctrl.addr], qarg.addr, math.pi / 2, math.pi / 2, -math.pi / 2
+            )
+        return ()
+
+    @interp.impl(uop.Swap)
+    def swap(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.Swap):
+        qarg1: PyQrackQubit = frame.get(stmt.ctrl)
+        qarg2: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg1.is_active() and qarg2.is_active():
+            qarg1.sim_reg.swap(qarg1.addr, qarg2.addr)
+        return ()
+
+    @interp.impl(uop.CSwap)
+    def cswap(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.CSwap):
+        qarg1: PyQrackQubit = frame.get(stmt.qarg1)
+        qarg2: PyQrackQubit = frame.get(stmt.qarg2)
+        ctrl: PyQrackQubit = frame.get(stmt.ctrl)
+        if qarg1.is_active() and qarg2.is_active():
+            qarg1.sim_reg.cswap([ctrl.addr], qarg1.addr, qarg2.addr)
+        return ()
+
+    @interp.impl(uop.CX)
+    @interp.impl(uop.CZ)
+    @interp.impl(uop.CY)
+    @interp.impl(uop.CH)
+    def control_gate(
+        self,
+        interp: interp.Interpreter,
+        frame: interp.Frame,
+        stmt: uop.CX | uop.CZ | uop.CY,
+    ):
+        ctrl: PyQrackQubit = frame.get(stmt.ctrl)
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if ctrl.is_active() and qarg.is_active():
+            getattr(qarg.sim_reg, self.GATE_TO_METHOD[stmt.name])(
+                [ctrl.addr], qarg.addr
+            )
+        return ()
+
+    @interp.impl(uop.CCX)
+    def ccx(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.CCX):
+        ctrl1: PyQrackQubit = frame.get(stmt.ctrl1)
+        ctrl2: PyQrackQubit = frame.get(stmt.ctrl2)
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if ctrl1.is_active() and ctrl2.is_active() and qarg.is_active():
+            qarg.sim_reg.mcx([ctrl1.addr, ctrl2.addr], qarg.addr)
+        return ()
+
+    @interp.impl(uop.RX)
+    @interp.impl(uop.RY)
+    @interp.impl(uop.RZ)
+    def rotation(
+        self,
+        interp: interp.Interpreter,
+        frame: interp.Frame,
+        stmt: uop.RX | uop.RY | uop.RZ,
+    ):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active():
+            qarg.sim_reg.r(self.AXIS_MAP[stmt.name], frame.get(stmt.theta), qarg.addr)
+        return ()
+
+    @interp.impl(uop.U1)
+    def u1(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.U1):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active():
+            qarg.sim_reg.u(qarg.addr, 0, 0, frame.get(stmt.lam))
+        return ()
+
+    @interp.impl(uop.U2)
+    def u2(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.U2):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active():
+            qarg.sim_reg.u(
+                qarg.addr, math.pi / 2, frame.get(stmt.phi), frame.get(stmt.lam)
+            )
+        return ()
+
+    @interp.impl(uop.CRX)
+    @interp.impl(uop.CRY)
+    @interp.impl(uop.CRZ)
+    def crx(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.CRX):
+        ctrl: PyQrackQubit = frame.get(stmt.ctrl)
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active() and ctrl.is_active():
+            qarg.sim_reg.mcr(
+                self.AXIS_MAP[stmt.name], frame.get(stmt.lam), [ctrl.addr], qarg.addr
+            )
+        return ()
+
+    @interp.impl(uop.CU1)
+    def cu1(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.CU1):
+        ctrl: PyQrackQubit = frame.get(stmt.ctrl)
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active() and ctrl.is_active():
+            qarg.sim_reg.mcu([ctrl.addr], qarg.addr, 0, 0, frame.get(stmt.lam))
+        return ()
+
+    @interp.impl(uop.CU3)
+    def cu3(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.CU3):
+        ctrl: PyQrackQubit = frame.get(stmt.ctrl)
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active() and ctrl.is_active():
+            qarg.sim_reg.mcu(
+                [ctrl.addr],
+                qarg.addr,
+                frame.get(stmt.theta),
+                frame.get(stmt.phi),
+                frame.get(stmt.lam),
+            )
+        return ()
+
+    @interp.impl(uop.CU)
+    def cu(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.CU):
+        ctrl: PyQrackQubit = frame.get(stmt.ctrl)
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        if qarg.is_active() and ctrl.is_active():
+            ctrl.sim_reg.u(ctrl.addr, 0, 0, frame.get(stmt.gamma))
+            qarg.sim_reg.mcu(
+                [ctrl.addr],
+                qarg.addr,
+                frame.get(stmt.theta),
+                frame.get(stmt.phi),
+                frame.get(stmt.lam),
+            )
+        return ()
+
+    @interp.impl(uop.RXX)
+    def rxx(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.RXX):
+        a: PyQrackQubit = frame.get(stmt.qarg)
+        b: PyQrackQubit = frame.get(stmt.ctrl)
+        theta = frame.get(stmt.theta)
+        sim_reg = a.sim_reg
+        if a.is_active() and b.is_active():
+            sim_reg.u(a.addr, math.pi / 2, theta, 0)
+            sim_reg.h(b.addr)
+            sim_reg.mcx([a.addr], b.addr)
+            sim_reg.u(b.addr, 0, 0, -theta)
+            sim_reg.mcx([a.addr], b.addr)
+            sim_reg.h(b.addr)
+            sim_reg.u(a.addr, math.pi / 2, -math.pi, math.pi - theta)
+
+        return ()
+
+    @interp.impl(uop.RZZ)
+    def rzz(self, interp: interp.Interpreter, frame: interp.Frame, stmt: uop.RZZ):
+        a: PyQrackQubit = frame.get(stmt.qarg)
+        b: PyQrackQubit = frame.get(stmt.ctrl)
+        theta = frame.get(stmt.theta)
+        sim_reg = a.sim_reg
+        if a.is_active() and b.is_active():
+            sim_reg.mcx([a.addr], b.addr)
+            sim_reg.u(b.addr, 0, 0, theta)
+            sim_reg.mcx([a.addr], b.addr)
+
+        return ()

bloqade/pyqrack/reg.py

@@ -0,0 +1,109 @@
+import enum
+from typing import TYPE_CHECKING, List
+from dataclasses import dataclass
+
+from bloqade.qasm2.types import QReg, Qubit
+
+if TYPE_CHECKING:
+    from pyqrack import QrackSimulator
+
+
+class Measurement(enum.IntEnum):
+    """Enumeration of measurement results."""
+
+    Zero = 0
+    One = 1
+    Lost = enum.auto()
+
+
+class CRegister(list[Measurement]):
+    """Runtime representation of a classical register."""
+
+    def __init__(self, size: int):
+        super().__init__(Measurement.Zero for _ in range(size))
+
+
+@dataclass(frozen=True)
+class CBitRef:
+    """Object representing a reference to a classical bit."""
+
+    ref: CRegister
+    """The classical register that is holding this bit."""
+
+    pos: int
+    """The position of this bit in the classical register."""
+
+    def set_value(self, value: bool):
+        self.ref[self.pos] = value
+
+    def get_value(self):
+        return self.ref[self.pos]
+
+
+class QubitState(enum.Enum):
+    Active = enum.auto()
+    Lost = enum.auto()
+
+
+@dataclass(frozen=True)
+class PyQrackReg(QReg):
+    """Simulation runtime value of a quantum register."""
+
+    size: int
+    """The number of qubits in this register."""
+
+    sim_reg: "QrackSimulator"
+    """The register of the simulator."""
+
+    addrs: tuple[int, ...]
+    """The global addresses of the qubits in this register."""
+
+    qubit_state: List[QubitState]
+    """The state of each qubit in this register."""
+
+    def drop(self, pos: int):
+        """Drop the qubit at the given position in-place.
+
+        Args
+            pos (int): The position of the qubit to drop.
+
+        """
+        assert self.qubit_state[pos] is QubitState.Active, "Qubit already lost"
+        self.qubit_state[pos] = QubitState.Lost
+
+    def __getitem__(self, pos: int):
+        return PyQrackQubit(self, pos)
+
+
+@dataclass(frozen=True)
+class PyQrackQubit(Qubit):
+    """The runtime representation of a qubit reference."""
+
+    ref: PyQrackReg
+    """The quantum register that is holding this qubit."""
+
+    pos: int
+    """The position of this qubit in the quantum register."""
+
+    @property
+    def sim_reg(self):
+        """The register of the simulator."""
+        return self.ref.sim_reg
+
+    @property
+    def addr(self) -> int:
+        """The global address of the qubit."""
+        return self.ref.addrs[self.pos]
+
+    def is_active(self) -> bool:
+        """Check if the qubit is active.
+
+        Returns
+            True if the qubit is active, False otherwise.
+
+        """
+        return self.ref.qubit_state[self.pos] is QubitState.Active
+
+    def drop(self):
+        """Drop the qubit in-place."""
+        self.ref.drop(self.pos)

bloqade/pyqrack/target.py

@@ -0,0 +1,112 @@
+from typing import List, TypeVar, ParamSpec
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Fold
+
+from bloqade.pyqrack.base import (
+    StackMemory,
+    DynamicMemory,
+    PyQrackOptions,
+    PyQrackInterpreter,
+    _default_pyqrack_args,
+)
+from bloqade.analysis.address import AnyAddress, AddressAnalysis
+
+Params = ParamSpec("Params")
+RetType = TypeVar("RetType")
+
+
+@dataclass
+class PyQrack:
+    """PyQrack target runtime for Bloqade."""
+
+    min_qubits: int = 0
+    """Minimum number of qubits required for the PyQrack simulator.
+    Useful when address analysis fails to determine the number of qubits.
+    """
+    dynamic_qubits: bool = False
+    """Whether to use dynamic qubit allocation. Cannot use with tensor network simulations."""
+
+    pyqrack_options: PyQrackOptions = field(default_factory=_default_pyqrack_args)
+    """Options to pass to the QrackSimulator object, node `qubitCount` will be overwritten."""
+
+    def __post_init__(self):
+        self.pyqrack_options = PyQrackOptions(
+            {**_default_pyqrack_args(), **self.pyqrack_options}
+        )
+
+    def _get_interp(self, mt: ir.Method[Params, RetType]):
+        if self.dynamic_qubits:
+
+            options = self.pyqrack_options.copy()
+            options["qubitCount"] = -1
+            return PyQrackInterpreter(mt.dialects, memory=DynamicMemory(options))
+        else:
+            address_analysis = AddressAnalysis(mt.dialects)
+            frame, _ = address_analysis.run_analysis(mt)
+            if self.min_qubits == 0 and any(
+                isinstance(a, AnyAddress) for a in frame.entries.values()
+            ):
+                raise ValueError(
+                    "All addresses must be resolved. Or set min_qubits to a positive integer."
+                )
+
+            num_qubits = max(address_analysis.qubit_count, self.min_qubits)
+            options = self.pyqrack_options.copy()
+            options["qubitCount"] = num_qubits
+            memory = StackMemory(
+                options,
+                total=num_qubits,
+            )
+
+            return PyQrackInterpreter(mt.dialects, memory=memory)
+
+    def run(
+        self,
+        mt: ir.Method[Params, RetType],
+        *args: Params.args,
+        **kwargs: Params.kwargs,
+    ) -> RetType:
+        """Run the given kernel method on the PyQrack simulator.
+
+        Args
+            mt (Method):
+                The kernel method to run.
+
+        Returns
+            The result of the kernel method, if any.
+
+        """
+        fold = Fold(mt.dialects)
+        fold(mt)
+        return self._get_interp(mt).run(mt, args, kwargs).expect()
+
+    def multi_run(
+        self,
+        mt: ir.Method[Params, RetType],
+        _shots: int,
+        *args: Params.args,
+        **kwargs: Params.kwargs,
+    ) -> List[RetType]:
+        """Run the given kernel method on the PyQrack `_shots` times, caching analysis results.
+
+        Args
+            mt (Method):
+                The kernel method to run.
+            _shots (int):
+                The number of times to run the kernel method.
+
+        Returns
+            List of results of the kernel method, one for each shot.
+
+        """
+        fold = Fold(mt.dialects)
+        fold(mt)
+
+        interpreter = self._get_interp(mt)
+        batched_results = []
+        for _ in range(_shots):
+            batched_results.append(interpreter.run(mt, args, kwargs).expect())
+
+        return batched_results

bloqade/qasm2/__init__.py

@@ -0,0 +1,19 @@
+from bloqade.types import Qubit as Qubit, QubitType as QubitType
+
+from . import (
+    emit as emit,
+    glob as glob,
+    parse as parse,
+    dialects as dialects,
+    parallel as parallel,
+)
+from .types import (
+    Bit as Bit,
+    CReg as CReg,
+    QReg as QReg,
+    BitType as BitType,
+    CRegType as CRegType,
+    QRegType as QRegType,
+)
+from .groups import gate as gate, main as main, extended as extended
+from ._wrappers import *  # noqa: F403