bloqade-circuit 0.1.0__tar.gz → 0.2.0__tar.gz

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/.github/workflows/ci.yml

@@ -21,7 +21,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Install uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v6
         with:
           # Install a specific version of uv.
           version: "0.5.1"

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/.github/workflows/release.yml

@@ -10,7 +10,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Install uv
-      uses: astral-sh/setup-uv@v5
+      uses: astral-sh/setup-uv@v6
       with:
         # Install a specific version of uv.
         version: "0.5.5"

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/PKG-INFO

@@ -1,19 +1,25 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.1.0
+Version: 0.2.0
 Summary: The software development toolkit for neutral atom arrays.
 Author-email: Roger-luo <rluo@quera.com>, kaihsin <khwu@quera.com>, weinbe58 <pweinberg@quera.com>, johnzl-777 <jlong@quera.com>
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: kirin-toolchain~=0.16.0
+Requires-Dist: kirin-toolchain~=0.17.0
 Requires-Dist: numpy>=1.22.0
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: pydantic<2.11.0,>=1.3.0
+Requires-Dist: pyqrack-cpu>=1.38.2; sys_platform != 'darwin'
+Requires-Dist: pyqrack>=1.38.2; sys_platform == 'darwin'
 Requires-Dist: rich>=13.9.4
 Requires-Dist: scipy>=1.13.1
 Provides-Extra: cirq
 Requires-Dist: cirq-core>=1.4.1; extra == 'cirq'
 Requires-Dist: cirq-core[contrib]>=1.4.1; extra == 'cirq'
+Provides-Extra: pyqrack-cuda
+Requires-Dist: pyqrack-cuda>=1.38.2; extra == 'pyqrack-cuda'
+Provides-Extra: pyqrack-opencl
+Requires-Dist: pyqrack>=1.38.2; (sys_platform != 'darwin') and extra == 'pyqrack-opencl'
 Provides-Extra: qasm2
 Requires-Dist: lark>=1.2.2; extra == 'qasm2'
 Provides-Extra: qbraid

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bloqade-circuit"
-version = "0.1.0"
+version = "0.2.0"
 description = "The software development toolkit for neutral atom arrays."
 readme = "README.md"
 authors = [
@@ -13,10 +13,12 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy>=1.22.0",
     "scipy>=1.13.1",
-    "kirin-toolchain~=0.16.0",
+    "kirin-toolchain~=0.17.0",
     "rich>=13.9.4",
     "pydantic>=1.3.0,<2.11.0",
     "pandas>=2.2.3",
+    "pyqrack>=1.38.2 ; sys_platform == 'darwin'",
+    "pyqrack-cpu>=1.38.2 ; sys_platform != 'darwin'",
 ]
 
 [project.optional-dependencies]
@@ -36,6 +38,12 @@ cirq = [
     "cirq-core>=1.4.1",
     "cirq-core[contrib]>=1.4.1",
 ]
+pyqrack-opencl = [
+    "pyqrack>=1.38.2 ; sys_platform != 'darwin'",
+]
+pyqrack-cuda = [
+    "pyqrack-cuda>=1.38.2",
+]
 
 [build-system]
 requires = ["hatchling"]

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/src/bloqade/analysis/address/impls.py

@@ -192,7 +192,10 @@ class SquinWireMethodTable(interp.MethodTable):
 
         origin_qubit = frame.get(stmt.qubit)
 
-        return (AddressWire(origin_qubit=origin_qubit),)
+        if isinstance(origin_qubit, AddressQubit):
+            return (AddressWire(origin_qubit=origin_qubit),)
+        else:
+            return (Address.top(),)
 
     @interp.impl(squin.wire.Apply)
     def apply(
@@ -201,14 +204,7 @@ class SquinWireMethodTable(interp.MethodTable):
         frame: ForwardFrame[Address],
         stmt: squin.wire.Apply,
     ):
-
-        origin_qubits = tuple(
-            [frame.get(input_elem).origin_qubit for input_elem in stmt.inputs]
-        )
-        new_address_wires = tuple(
-            [AddressWire(origin_qubit=origin_qubit) for origin_qubit in origin_qubits]
-        )
-        return new_address_wires
+        return frame.get_values(stmt.inputs)
 
 
 @squin.qubit.dialect.register(key="qubit.address")

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/src/bloqade/analysis/address/lattice.py

@@ -81,5 +81,5 @@ class AddressWire(Address):
 
     def is_subseteq(self, other: Address) -> bool:
         if isinstance(other, AddressWire):
-            return self.origin_qubit == self.origin_qubit
+            return self.origin_qubit == other.origin_qubit
         return False

bloqade_circuit-0.2.0/src/bloqade/analysis/fidelity/__init__.py

@@ -0,0 +1 @@
+from .analysis import FidelityAnalysis as FidelityAnalysis

bloqade_circuit-0.2.0/src/bloqade/analysis/fidelity/analysis.py

@@ -0,0 +1,69 @@
+from typing import Any
+from dataclasses import field
+
+from kirin import ir
+from kirin.lattice import EmptyLattice
+from kirin.analysis import Forward
+from kirin.interp.value import Successor
+from kirin.analysis.forward import ForwardFrame
+
+from ..address import AddressAnalysis
+
+
+class FidelityAnalysis(Forward):
+    """
+    This analysis pass can be used to track the global addresses of qubits and wires.
+    """
+
+    keys = ["circuit.fidelity"]
+    lattice = EmptyLattice
+
+    """
+    The fidelity of the gate set described by the analysed program. It reduces whenever a noise channel is encountered.
+    """
+    gate_fidelity: float = 1.0
+
+    _current_gate_fidelity: float = field(init=False)
+
+    """
+    The probabilities that each of the atoms in the register survive the duration of the analysed program. The order of the list follows the order they are in the register.
+    """
+    atom_survival_probability: list[float] = field(init=False)
+
+    _current_atom_survival_probability: list[float] = field(init=False)
+
+    addr_frame: ForwardFrame = field(init=False)
+
+    def initialize(self):
+        super().initialize()
+        self._current_gate_fidelity = 1.0
+        self._current_atom_survival_probability = [
+            1.0 for _ in range(len(self.atom_survival_probability))
+        ]
+        return self
+
+    def posthook_succ(self, frame: ForwardFrame, succ: Successor):
+        self.gate_fidelity *= self._current_gate_fidelity
+        for i, _current_survival in enumerate(self._current_atom_survival_probability):
+            self.atom_survival_probability[i] *= _current_survival
+
+    def eval_stmt_fallback(self, frame: ForwardFrame, stmt: ir.Statement):
+        # NOTE: default is to conserve fidelity, so do nothing here
+        return
+
+    def run_method(self, method: ir.Method, args: tuple[EmptyLattice, ...]):
+        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+
+    def run_analysis(
+        self, method: ir.Method, args: tuple | None = None, *, no_raise: bool = True
+    ) -> tuple[ForwardFrame, Any]:
+        self._run_address_analysis(method, no_raise=no_raise)
+        return super().run_analysis(method, args, no_raise=no_raise)
+
+    def _run_address_analysis(self, method: ir.Method, no_raise: bool):
+        addr_analysis = AddressAnalysis(self.dialects)
+        addr_frame, _ = addr_analysis.run_analysis(method=method, no_raise=no_raise)
+        self.addr_frame = addr_frame
+
+        # NOTE: make sure we have as many probabilities as we have addresses
+        self.atom_survival_probability = [1.0] * addr_analysis.qubit_count

bloqade_circuit-0.2.0/src/bloqade/device.py

@@ -0,0 +1,130 @@
+import abc
+from typing import Any, Generic, TypeVar, ParamSpec
+
+from kirin import ir
+
+from bloqade.task import (
+    BatchFuture,
+    AbstractTask,
+    AbstractRemoteTask,
+    AbstractSimulatorTask,
+    DeviceTaskExpectMixin,
+)
+
+Params = ParamSpec("Params")
+RetType = TypeVar("RetType")
+ObsType = TypeVar("ObsType")
+
+
+TaskType = TypeVar("TaskType", bound=AbstractTask)
+
+
+class AbstractDevice(abc.ABC, Generic[TaskType]):
+    """Abstract base class for devices. Defines the minimum interface for devices."""
+
+    @abc.abstractmethod
+    def task(
+        self,
+        kernel: ir.Method[Params, RetType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> TaskType:
+        """Creates a remote task for the device."""
+
+
+ExpectTaskType = TypeVar("ExpectTaskType", bound=DeviceTaskExpectMixin)
+
+
+class ExpectationDeviceMixin(AbstractDevice[ExpectTaskType]):
+    def expect(
+        self,
+        kernel: ir.Method[Params, RetType],
+        observable: ir.Method[[RetType], ObsType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+        *,
+        shots: int = 1,
+    ) -> ObsType:
+        """Returns the expectation value of the given observable after running the task."""
+        return self.task(kernel, args, kwargs).expect(observable, shots)
+
+
+RemoteTaskType = TypeVar("RemoteTaskType", bound=AbstractRemoteTask)
+
+
+class AbstractRemoteDevice(AbstractDevice[RemoteTaskType]):
+    """Abstract base class for remote devices."""
+
+    def run(
+        self,
+        kernel: ir.Method[Params, RetType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+        *,
+        shots: int = 1,
+        timeout: float | None = None,
+    ) -> list[RetType]:
+        """Runs the kernel and returns the result.
+
+        Args:
+            kernel (ir.Method):
+                The kernel method to run.
+            args (tuple[Any, ...]):
+                Positional arguments to pass to the kernel method.
+            kwargs (dict[str, Any] | None):
+                Keyword arguments to pass to the kernel method.
+            shots (int):
+                The number of times to run the kernel method.
+            timeout (float | None):
+                Timeout in seconds for the asynchronous execution. If None, wait indefinitely.
+
+        Returns:
+            list[RetType]:
+                The result of the kernel method, if any.
+
+        """
+        return self.task(kernel, args, kwargs).run(shots=shots, timeout=timeout)
+
+    def run_async(
+        self,
+        kernel: ir.Method[Params, RetType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+        *,
+        shots: int = 1,
+    ) -> BatchFuture[RetType]:
+        """Runs the kernel asynchronously and returns a Future object.
+
+        Args:
+            kernel (ir.Method):
+                The kernel method to run.
+            args (tuple[Any, ...]):
+                Positional arguments to pass to the kernel method.
+            kwargs (dict[str, Any] | None):
+                Keyword arguments to pass to the kernel method.
+            shots (int):
+                The number of times to run the kernel method.
+
+        Returns:
+            Future[list[RetType]]:
+                The Future for all executions of the kernel method.
+
+
+        """
+        return self.task(kernel, args, kwargs).run_async(shots=shots)
+
+
+SimulatorTaskType = TypeVar("SimulatorTaskType", bound=AbstractSimulatorTask)
+
+
+class AbstractSimulatorDevice(AbstractDevice[SimulatorTaskType]):
+    """Abstract base class for simulator devices."""
+
+    def run(
+        self,
+        kernel: ir.Method[Params, RetType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> RetType:
+        """Runs the kernel and returns the result."""
+        return self.task(kernel, args, kwargs).run()

bloqade_circuit-0.2.0/src/bloqade/noise/__init__.py

@@ -0,0 +1,2 @@
+# NOTE: just to register methods
+from . import native as native, fidelity as fidelity

bloqade_circuit-0.2.0/src/bloqade/noise/fidelity.py

@@ -0,0 +1,51 @@
+from kirin import interp
+from kirin.lattice import EmptyLattice
+
+from bloqade.analysis.fidelity import FidelityAnalysis
+
+from .native import dialect as native
+from .native.stmts import PauliChannel, CZPauliChannel, AtomLossChannel
+from ..analysis.address import AddressQubit, AddressTuple
+
+
+@native.register(key="circuit.fidelity")
+class FidelityMethodTable(interp.MethodTable):
+
+    @interp.impl(PauliChannel)
+    @interp.impl(CZPauliChannel)
+    def pauli_channel(
+        self,
+        interp: FidelityAnalysis,
+        frame: interp.Frame[EmptyLattice],
+        stmt: PauliChannel | CZPauliChannel,
+    ):
+        probs = stmt.probabilities
+        try:
+            ps, ps_ctrl = probs
+        except ValueError:
+            (ps,) = probs
+            ps_ctrl = ()
+
+        p = sum(ps)
+        p_ctrl = sum(ps_ctrl)
+
+        # NOTE: fidelity is just the inverse probability of any noise to occur
+        fid = (1 - p) * (1 - p_ctrl)
+
+        interp._current_gate_fidelity *= fid
+
+    @interp.impl(AtomLossChannel)
+    def atom_loss(
+        self,
+        interp: FidelityAnalysis,
+        frame: interp.Frame[EmptyLattice],
+        stmt: AtomLossChannel,
+    ):
+        # NOTE: since AtomLossChannel acts on IList[Qubit], we know the assigned address is a tuple
+        addresses: AddressTuple = interp.addr_frame.get(stmt.qargs)
+
+        # NOTE: get the corresponding index and reduce survival probability accordingly
+        for qbit_address in addresses.data:
+            assert isinstance(qbit_address, AddressQubit)
+            index = qbit_address.data
+            interp._current_atom_survival_probability[index] *= 1 - stmt.prob

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/src/bloqade/noise/native/model.py

@@ -102,10 +102,9 @@ class MoveNoiseModelABC(abc.ABC):
     params: MoveNoiseParams = field(default_factory=MoveNoiseParams)
     """Parameters for calculating move noise."""
 
-    @classmethod
     @abc.abstractmethod
     def parallel_cz_errors(
-        cls, ctrls: List[int], qargs: List[int], rest: List[int]
+        self, ctrls: List[int], qargs: List[int], rest: List[int]
     ) -> Dict[Tuple[float, float, float, float], List[int]]:
         """Takes a set of ctrls and qargs and returns a noise model for all qubits."""
         pass

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/src/bloqade/noise/native/rewrite.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 
 from kirin import ir
-from kirin.rewrite import abc, dce, walk, result, fixpoint
+from kirin.rewrite import abc, dce, walk, fixpoint
 from kirin.passes.abc import Pass
 
 from .stmts import PauliChannel, CZPauliChannel, AtomLossChannel
@@ -9,19 +9,19 @@ from ._dialect import dialect
 
 
 class RemoveNoiseRewrite(abc.RewriteRule):
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if isinstance(node, (AtomLossChannel, PauliChannel, CZPauliChannel)):
             node.delete()
-            return result.RewriteResult(has_done_something=True)
+            return abc.RewriteResult(has_done_something=True)
 
-        return result.RewriteResult()
+        return abc.RewriteResult()
 
 
 @dataclass
 class RemoveNoisePass(Pass):
     name = "remove-noise"
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         delete_walk = walk.Walk(RemoveNoiseRewrite())
         dce_walk = fixpoint.Fixpoint(walk.Walk(dce.DeadCodeElimination()))
 

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/src/bloqade/noise/native/stmts.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 from kirin import ir, types, lowering
 from kirin.decl import info, statement
 from kirin.dialects import ilist
@@ -7,25 +9,44 @@ from bloqade.qasm2.types import QubitType
 from ._dialect import dialect
 
 
-@statement(dialect=dialect)
-class PauliChannel(ir.Statement):
-
+@statement
+class NativeNoiseStmt(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
 
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        raise NotImplementedError(f"Override the method in {type(self).__name__}")
+
+    def check(self):
+        for probs in self.probabilities:
+            self.check_probability(sum(probs))
+            for p in probs:
+                self.check_probability(p)
+
+    def check_probability(self, p: float):
+        if not 0 <= p <= 1:
+            raise ValueError(
+                f"Invalid noise probability encountered in {type(self).__name__}: {p}"
+            )
+
+
+@statement(dialect=dialect)
+class PauliChannel(NativeNoiseStmt):
     px: float = info.attribute(types.Float)
     py: float = info.attribute(types.Float)
     pz: float = info.attribute(types.Float)
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType])
 
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return ((self.px, self.py, self.pz),)
+
 
 NumQubits = types.TypeVar("NumQubits")
 
 
 @statement(dialect=dialect)
-class CZPauliChannel(ir.Statement):
-
-    traits = frozenset({lowering.FromPythonCall()})
-
+class CZPauliChannel(NativeNoiseStmt):
     paired: bool = info.attribute(types.Bool)
     px_ctrl: float = info.attribute(types.Float)
     py_ctrl: float = info.attribute(types.Float)
@@ -36,11 +57,19 @@ class CZPauliChannel(ir.Statement):
     ctrls: ir.SSAValue = info.argument(ilist.IListType[QubitType, NumQubits])
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType, NumQubits])
 
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return (
+            (self.px_ctrl, self.py_ctrl, self.pz_ctrl),
+            (self.px_qarg, self.py_qarg, self.pz_qarg),
+        )
 
-@statement(dialect=dialect)
-class AtomLossChannel(ir.Statement):
-
-    traits = frozenset({lowering.FromPythonCall()})
 
+@statement(dialect=dialect)
+class AtomLossChannel(NativeNoiseStmt):
     prob: float = info.attribute(types.Float)
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return ((self.prob,),)

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/src/bloqade/pyqrack/__init__.py

@@ -1,9 +1,9 @@
 from .reg import (
     CBitRef as CBitRef,
     CRegister as CRegister,
-    PyQrackReg as PyQrackReg,
     QubitState as QubitState,
     Measurement as Measurement,
+    PyQrackWire as PyQrackWire,
     PyQrackQubit as PyQrackQubit,
 )
 from .base import (
@@ -11,8 +11,14 @@ from .base import (
     DynamicMemory as DynamicMemory,
     PyQrackInterpreter as PyQrackInterpreter,
 )
+from .task import PyQrackSimulatorTask as PyQrackSimulatorTask
 
 # 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 .qasm2 import uop as uop, core as core, glob as glob, parallel as parallel
+from .squin import op as op, qubit as qubit
+from .device import (
+    StackMemorySimulator as StackMemorySimulator,
+    DynamicMemorySimulator as DynamicMemorySimulator,
+)
 from .target import PyQrack as PyQrack

{bloqade_circuit-0.1.0 → bloqade_circuit-0.2.0}/src/bloqade/pyqrack/base.py

@@ -26,13 +26,28 @@ class PyQrackOptions(typing.TypedDict):
     isOpenCL: bool
 
 
+def _validate_pyqrack_options(options: PyQrackOptions) -> None:
+    if options["isBinaryDecisionTree"] and options["isStabilizerHybrid"]:
+        raise ValueError(
+            "Cannot use both isBinaryDecisionTree and isStabilizerHybrid at the same time."
+        )
+    elif options["isTensorNetwork"] and options["isBinaryDecisionTree"]:
+        raise ValueError(
+            "Cannot use both isTensorNetwork and isBinaryDecisionTree at the same time."
+        )
+    elif options["isTensorNetwork"] and options["isStabilizerHybrid"]:
+        raise ValueError(
+            "Cannot use both isTensorNetwork and isStabilizerHybrid at the same time."
+        )
+
+
 def _default_pyqrack_args() -> PyQrackOptions:
     return PyQrackOptions(
         qubitCount=-1,
         isTensorNetwork=False,
         isSchmidtDecomposeMulti=True,
         isSchmidtDecompose=True,
-        isStabilizerHybrid=True,
+        isStabilizerHybrid=False,
         isBinaryDecisionTree=True,
         isPaged=True,
         isCpuGpuHybrid=True,
@@ -45,6 +60,9 @@ class MemoryABC(abc.ABC):
     pyqrack_options: PyQrackOptions = field(default_factory=_default_pyqrack_args)
     sim_reg: "QrackSimulator" = field(init=False)
 
+    def __post_init__(self):
+        _validate_pyqrack_options(self.pyqrack_options)
+
     @abc.abstractmethod
     def allocate(self, n_qubits: int) -> tuple[int, ...]:
         """Allocate `n_qubits` qubits and return their ids."""
@@ -115,10 +133,13 @@ class DynamicMemory(MemoryABC):
         return tuple(range(start, start + n_qubits))
 
 
+MemoryType = typing.TypeVar("MemoryType", bound=MemoryABC)
+
+
 @dataclass
-class PyQrackInterpreter(Interpreter):
+class PyQrackInterpreter(Interpreter, typing.Generic[MemoryType]):
     keys = ["pyqrack", "main"]
-    memory: MemoryABC = field(kw_only=True)
+    memory: MemoryType = field(kw_only=True)
     rng_state: np.random.Generator = field(
         default_factory=np.random.default_rng, kw_only=True
     )