bloqade-circuit 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

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/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/analysis/fidelity/__init__.py

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

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/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/noise/__init__.py

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

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/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/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/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/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/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
     )

bloqade/pyqrack/device.py

@@ -0,0 +1,166 @@
+from typing import Any, TypeVar, ParamSpec
+from dataclasses import field, dataclass
+
+import numpy as np
+from kirin import ir
+
+from pyqrack.pauli import Pauli
+from bloqade.device import AbstractSimulatorDevice
+from bloqade.pyqrack.reg import Measurement, PyQrackQubit
+from bloqade.pyqrack.base import (
+    MemoryABC,
+    StackMemory,
+    DynamicMemory,
+    PyQrackOptions,
+    PyQrackInterpreter,
+    _default_pyqrack_args,
+)
+from bloqade.pyqrack.task import PyQrackSimulatorTask
+from bloqade.analysis.address.lattice import AnyAddress
+from bloqade.analysis.address.analysis import AddressAnalysis
+
+RetType = TypeVar("RetType")
+Params = ParamSpec("Params")
+
+
+@dataclass
+class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
+    options: PyQrackOptions = field(default_factory=_default_pyqrack_args)
+    loss_m_result: Measurement = field(default=Measurement.One, kw_only=True)
+    rng_state: np.random.Generator = field(
+        default_factory=np.random.default_rng, kw_only=True
+    )
+
+    MemoryType = TypeVar("MemoryType", bound=MemoryABC)
+
+    def __post_init__(self):
+        self.options = PyQrackOptions({**_default_pyqrack_args(), **self.options})
+
+    def new_task(
+        self,
+        mt: ir.Method[Params, RetType],
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+        memory: MemoryType,
+    ) -> PyQrackSimulatorTask[Params, RetType, MemoryType]:
+        interp = PyQrackInterpreter(
+            mt.dialects,
+            memory=memory,
+            rng_state=self.rng_state,
+            loss_m_result=self.loss_m_result,
+        )
+        return PyQrackSimulatorTask(
+            kernel=mt, args=args, kwargs=kwargs, pyqrack_interp=interp
+        )
+
+    def state_vector(
+        self,
+        kernel: ir.Method[Params, RetType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> list[complex]:
+        """Runs task and returns the state vector."""
+        task = self.task(kernel, args, kwargs)
+        task.run()
+        return task.state.sim_reg.out_ket()
+
+    @staticmethod
+    def pauli_expectation(pauli: list[Pauli], qubits: list[PyQrackQubit]) -> float:
+        """Returns the expectation value of the given Pauli operator given a list of Pauli operators and qubits.
+
+        Args:
+            pauli (list[Pauli]):
+                List of Pauli operators to compute the expectation value for.
+            qubits (list[PyQrackQubit]):
+                List of qubits corresponding to the Pauli operators.
+
+        returns:
+            float:
+                The expectation value of the Pauli operator.
+
+        """
+
+        if len(pauli) == 0:
+            return 0.0
+
+        if len(pauli) != len(qubits):
+            raise ValueError("Length of Pauli and qubits must match.")
+
+        sim_reg = qubits[0].sim_reg
+
+        if any(qubit.sim_reg is not sim_reg for qubit in qubits):
+            raise ValueError("All qubits must belong to the same simulator register.")
+
+        qubit_ids = [qubit.addr for qubit in qubits]
+
+        if len(qubit_ids) != len(set(qubit_ids)):
+            raise ValueError("Qubits must be unique.")
+
+        return sim_reg.pauli_expectation(pauli, qubit_ids)
+
+
+@dataclass
+class StackMemorySimulator(PyQrackSimulatorBase):
+    """PyQrack simulator device with precalculated stack of qubits."""
+
+    min_qubits: int = field(default=0, kw_only=True)
+
+    def task(
+        self,
+        kernel: ir.Method[Params, RetType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ):
+        if kwargs is None:
+            kwargs = {}
+
+        address_analysis = AddressAnalysis(dialects=kernel.dialects)
+        frame, _ = address_analysis.run_analysis(kernel)
+        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.options.copy()
+        options["qubitCount"] = num_qubits
+        memory = StackMemory(
+            options,
+            total=num_qubits,
+        )
+
+        return self.new_task(kernel, args, kwargs, memory)
+
+
+@dataclass
+class DynamicMemorySimulator(PyQrackSimulatorBase):
+    """PyQrack simulator device with dynamic qubit allocation."""
+
+    def task(
+        self,
+        kernel: ir.Method[Params, RetType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ):
+        if kwargs is None:
+            kwargs = {}
+
+        memory = DynamicMemory(self.options.copy())
+        return self.new_task(kernel, args, kwargs, memory)
+
+
+def test():
+    from bloqade.qasm2 import extended
+
+    @extended
+    def main():
+        return 1
+
+    @extended
+    def obs(result: int) -> int:
+        return result
+
+    res = DynamicMemorySimulator().task(main)
+    return res.run()

bloqade/pyqrack/noise/native.py

@@ -93,8 +93,7 @@ class PyQrackMethods(interp.MethodTable):
 
         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.sim_reg.m(qarg.addr)
                 qarg.drop()
 
         return ()