bloqade-circuit 0.6.4__py3-none-any.whl → 0.9.1__py3-none-any.whl

bloqade/native/upstream/squin2native.py

@@ -0,0 +1,79 @@
+from itertools import chain
+
+from kirin import ir, rewrite
+from kirin.dialects import py, func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.analysis.callgraph import CallGraph
+
+from bloqade.native import kernel, broadcast
+from bloqade.squin.gate import stmts, dialect as gate_dialect
+from bloqade.rewrite.passes import CallGraphPass, UpdateDialectsOnCallGraph
+
+
+class GateRule(RewriteRule):
+    SQUIN_MAPPING: dict[type[ir.Statement], tuple[ir.Method, ...]] = {
+        stmts.X: (broadcast.x,),
+        stmts.Y: (broadcast.y,),
+        stmts.Z: (broadcast.z,),
+        stmts.H: (broadcast.h,),
+        stmts.S: (broadcast.s, broadcast.s_adj),
+        stmts.T: (broadcast.t, broadcast.t_adj),
+        stmts.SqrtX: (broadcast.sqrt_x, broadcast.sqrt_x_adj),
+        stmts.SqrtY: (broadcast.sqrt_y, broadcast.sqrt_y_adj),
+        stmts.Rx: (broadcast.rx,),
+        stmts.Ry: (broadcast.ry,),
+        stmts.Rz: (broadcast.rz,),
+        stmts.CX: (broadcast.cx,),
+        stmts.CY: (broadcast.cy,),
+        stmts.CZ: (broadcast.cz,),
+        stmts.U3: (broadcast.u3,),
+    }
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if (native_methods := self.SQUIN_MAPPING.get(type(node))) is None:
+            return RewriteResult()
+
+        if isinstance(node, stmts.SingleQubitNonHermitianGate):
+            native_method = native_methods[1] if node.adjoint else native_methods[0]
+        else:
+            native_method = native_methods[0]
+
+        # do not rewrite in invoke because callgraph pass will be looking for invoke statements
+        (callee := py.Constant(native_method)).insert_before(node)
+        node.replace_by(func.Call(callee.result, tuple(node.args), kwargs=()))
+
+        return RewriteResult(has_done_something=True)
+
+
+class SquinToNative:
+    """A Target that converts Squin gates to native gates."""
+
+    def emit(self, mt: ir.Method, *, no_raise=True) -> ir.Method:
+        """Convert Squin gates to native gates.
+
+        Args:
+            mt (ir.Method): The method to convert.
+            no_raise (bool, optional): Whether to suppress errors. Defaults to True.
+
+        Returns:
+            ir.Method: The converted method.
+        """
+        old_callgraph = CallGraph(mt)
+        all_dialects = chain.from_iterable(
+            ker.dialects.data for kers in old_callgraph.defs.values() for ker in kers
+        )
+        combined_dialects = mt.dialects.union(all_dialects).union(kernel)
+
+        out = mt.similar(combined_dialects)
+        UpdateDialectsOnCallGraph(combined_dialects, no_raise=no_raise)(out)
+        CallGraphPass(combined_dialects, rewrite.Walk(GateRule()), no_raise=no_raise)(
+            out
+        )
+        # verify all kernels in the callgraph and discard gate dialect
+        out.dialects.discard(gate_dialect)
+        new_callgraph = CallGraph(out)
+        for ker in new_callgraph.edges.keys():
+            ker.dialects.discard(gate_dialect)
+            ker.verify()
+
+        return out

bloqade/pyqrack/__init__.py

@@ -3,7 +3,6 @@ from .reg import (
     CRegister as CRegister,
     QubitState as QubitState,
     Measurement as Measurement,
-    PyQrackWire as PyQrackWire,
     PyQrackQubit as PyQrackQubit,
 )
 from .base import (
@@ -16,9 +15,10 @@ 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, glob as glob, parallel as parallel
-from .squin import op as op, noise as noise, qubit as qubit
+from .squin import gate as gate, noise as noise, qubit as qubit
 from .device import (
     StackMemorySimulator as StackMemorySimulator,
     DynamicMemorySimulator as DynamicMemorySimulator,
 )
+from .native import NativeMethods as NativeMethods
 from .target import PyQrack as PyQrack

bloqade/pyqrack/base.py

@@ -48,7 +48,7 @@ def _default_pyqrack_args() -> PyQrackOptions:
         isSchmidtDecomposeMulti=True,
         isSchmidtDecompose=True,
         isStabilizerHybrid=False,
-        isBinaryDecisionTree=True,
+        isBinaryDecisionTree=False,
         isPaged=True,
         isCpuGpuHybrid=True,
         isOpenCL=True,
@@ -146,7 +146,13 @@ class PyQrackInterpreter(Interpreter, typing.Generic[MemoryType]):
     loss_m_result: Measurement = field(default=Measurement.One, kw_only=True)
     """The value of a measurement result when a qubit is lost."""
 
+    global_measurement_id: int = field(init=False, default=0)
+
     def initialize(self) -> Self:
         super().initialize()
         self.memory.reset()  # reset allocated qubits
         return self
+
+    def set_global_measurement_id(self, m: Measurement):
+        m.measurement_id = self.global_measurement_id
+        self.global_measurement_id += 1

bloqade/pyqrack/device.py

@@ -1,11 +1,10 @@
-from typing import Any, TypeVar, ParamSpec
+from typing import Any, TypeVar, ParamSpec, NamedTuple
 from dataclasses import field, dataclass
 
 import numpy as np
 from kirin import ir
-from kirin.passes import fold
+from kirin.dialects.ilist import IList
 
-from bloqade.squin import noise as squin_noise
 from pyqrack.pauli import Pauli
 from bloqade.device import AbstractSimulatorDevice
 from bloqade.pyqrack.reg import Measurement, PyQrackQubit
@@ -18,14 +17,153 @@ from bloqade.pyqrack.base import (
     _default_pyqrack_args,
 )
 from bloqade.pyqrack.task import PyQrackSimulatorTask
-from bloqade.squin.noise.rewrite import RewriteNoiseStmts
-from bloqade.analysis.address.lattice import AnyAddress
+from pyqrack.qrack_simulator import QrackSimulator
+from bloqade.analysis.address.lattice import UnknownReg, UnknownQubit
 from bloqade.analysis.address.analysis import AddressAnalysis
 
 RetType = TypeVar("RetType")
 Params = ParamSpec("Params")
 
 
+class QuantumState(NamedTuple):
+    """
+    A representation of a quantum state as a density matrix, where the density matrix is
+    rho = sum_i eigenvalues[i] |eigenvectors[:,i]><eigenvectors[:,i]|.
+
+    This representation is efficient for low-rank density matrices by only storing
+    the non-zero eigenvalues and corresponding eigenvectors of the density matrix.
+    For example, a pure state has only one non-zero eigenvalue equal to 1.0.
+
+    Endianness and qubit ordering of the state vector is consistent with Cirq, where
+    eigenvectors[0,0] corresponds to the amplitude of the |00..000> element of the zeroth eigenvector;
+    eigenvectors[1,0] corresponds to the amplitude of the |00..001> element of the zeroth eigenvector;
+    eigenvectors[3,0] corresponds to the amplitude of the |00..011> element of the zeroth eigenvector;
+    eigenvectors[-1,0] corresponds to the amplitude of the |11..111> element of the zeroth eigenvector.
+    A flip of the LAST bit |00..000><00..001| corresponds to applying a PauliX gate to the FIRST qubit.
+    A flip of the FIRST bit |00..000><10..000| corresponds to applying a PauliX gate to the LAST qubit.
+
+    Attributes:
+        eigenvalues (1d np.ndarray):
+            The non-zero eigenvalues of the density matrix.
+        eigenvectors (2d np.ndarray):
+            The corresponding eigenvectors of the density matrix,
+            where eigenvectors[:,i] is the i-th eigenvector.
+    Methods:
+        Not Implemented, pending https://github.com/QuEraComputing/bloqade-circuit/issues/447
+    """
+
+    eigenvalues: np.ndarray
+    eigenvectors: np.ndarray
+
+    def canonicalize(self, tol: float = 1e-12) -> "QuantumState":
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def __add__(self, other: "QuantumState") -> "QuantumState":
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def __mul__(self, scalar: float) -> "QuantumState":
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    @property
+    def dense(self) -> np.ndarray[tuple[int, int], np.complexfloating]:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def __matmul__(self, right: "cirq.Circuit") -> "QuantumState":  # noqa: F821
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def expect(self, operator: Any) -> float:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def probability(self) -> np.ndarray[tuple[int], np.floating]:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def von_neumann_entropy(self) -> float:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    @property
+    def qubit_basis(self) -> list[PyQrackQubit]:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def reduced_density_matrix(
+        self, qubits: list[PyQrackQubit], tol: float = 1e-12
+    ) -> "QuantumState":
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def overlap(self, other: "QuantumState") -> complex:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+
+def _pyqrack_reduced_density_matrix(
+    inds: tuple[int, ...], sim_reg: QrackSimulator, tol: float = 1e-12
+) -> QuantumState:
+    """
+    Extract the reduced density matrix representing the state of a list
+    of qubits from a PyQRack simulator register.
+
+    Inputs:
+        inds: A list of integers labeling the qubit registers to extract the reduced density matrix for
+        sim_reg: The PyQRack simulator register to extract the reduced density matrix from
+        tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+    Outputs:
+        An eigh result containing the eigenvalues and eigenvectors of the reduced density matrix.
+    """
+    # Identify the rest of the qubits in the register
+    N = sim_reg.num_qubits()
+    other = tuple(set(range(N)).difference(inds))
+
+    if len(set(inds)) != len(inds):
+        raise ValueError("Qubits must be unique.")
+
+    if max(inds) > N - 1:
+        raise ValueError(
+            f"Qubit indices {inds} exceed the number of qubits in the register {N}."
+        )
+
+    reordering = inds + other
+    # Fix pyqrack edannes to be consistent with Cirq.
+    reordering = tuple(N - 1 - x for x in reordering)
+    # Extract the statevector from the PyQRack qubits
+    statevector = np.array(sim_reg.out_ket())
+    # Reshape into a (2,2,2, ..., 2) tensor
+    vec_f = np.reshape(statevector, (2,) * N)
+    # Reorder the indexes to obey the order of the qubits
+    vec_p = np.transpose(vec_f, reordering)
+    # Rehape into a 2^N by 2^M matrix to compute the singular value decomposition
+    vec_svd = np.reshape(vec_p, (2 ** len(inds), 2 ** len(other)))
+    # The singular values and vectors are the eigenspace of the reduced density matrix
+    s, v, d = np.linalg.svd(vec_svd, full_matrices=False)
+
+    # Remove the negligible singular values
+    nonzero_inds = np.where(np.abs(v) > tol)[0]
+    s = s[:, nonzero_inds]
+    v = v[nonzero_inds] ** 2
+    # Forge into the correct result type
+    result = QuantumState(eigenvalues=v, eigenvectors=s)
+    return result
+
+
 @dataclass
 class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
     """PyQrack simulation device base class."""
@@ -50,23 +188,14 @@ class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
         kwargs: dict[str, Any],
         memory: MemoryType,
     ) -> PyQrackSimulatorTask[Params, RetType, MemoryType]:
-
-        if squin_noise in mt.dialects:
-            # NOTE: rewrite noise statements
-            mt_ = mt.similar(mt.dialects)
-            RewriteNoiseStmts(mt_.dialects)(mt_)
-            fold.Fold(mt_.dialects)(mt_)
-        else:
-            mt_ = mt
-
         interp = PyQrackInterpreter(
-            mt_.dialects,
+            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
+            kernel=mt, args=args, kwargs=kwargs, pyqrack_interp=interp
         )
 
     def state_vector(
@@ -112,6 +241,51 @@ class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
 
         return sim_reg.pauli_expectation(qubit_ids, pauli)
 
+    @staticmethod
+    def quantum_state(
+        qubits: list[PyQrackQubit] | IList[PyQrackQubit, Any], tol: float = 1e-12
+    ) -> "QuantumState":
+        """
+        Extract the reduced density matrix representing the state of a list
+        of qubits from a PyQRack simulator register.
+
+        Inputs:
+            qubits: A list of PyQRack qubits to extract the reduced density matrix for
+            tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+        Outputs:
+            An eigh result containing the eigenvalues and eigenvectors of the reduced density matrix.
+        """
+        if len(qubits) == 0:
+            return QuantumState(
+                eigenvalues=np.array([]), eigenvectors=np.array([]).reshape(0, 0)
+            )
+        sim_reg = qubits[0].sim_reg
+
+        if not all([x.sim_reg is sim_reg for x in qubits]):
+            raise ValueError("All qubits must be from the same simulator register.")
+        inds: tuple[int, ...] = tuple(qubit.addr for qubit in qubits)
+
+        return _pyqrack_reduced_density_matrix(inds, sim_reg, tol)
+
+    @classmethod
+    def reduced_density_matrix(
+        cls, qubits: list[PyQrackQubit] | IList[PyQrackQubit, Any], tol: float = 1e-12
+    ) -> np.ndarray:
+        """
+        Extract the reduced density matrix representing the state of a list
+        of qubits from a PyQRack simulator register.
+
+        Inputs:
+            qubits: A list of PyQRack qubits to extract the reduced density matrix for
+            tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+        Outputs:
+            A dense 2^n x 2^n numpy array representing the reduced density matrix.
+        """
+        rdm = cls.quantum_state(qubits, tol)
+        return np.einsum(
+            "ax,x,bx", rdm.eigenvectors, rdm.eigenvalues, rdm.eigenvectors.conj()
+        )
+
 
 @dataclass
 class StackMemorySimulator(PyQrackSimulatorBase):
@@ -179,9 +353,9 @@ class StackMemorySimulator(PyQrackSimulatorBase):
             kwargs = {}
 
         address_analysis = AddressAnalysis(dialects=kernel.dialects)
-        frame, _ = address_analysis.run_analysis(kernel)
+        frame, _ = address_analysis.run(kernel)
         if self.min_qubits == 0 and any(
-            isinstance(a, AnyAddress) for a in frame.entries.values()
+            isinstance(a, (UnknownQubit, UnknownReg)) for a in frame.entries.values()
         ):
             raise ValueError(
                 "All addresses must be resolved. Or set min_qubits to a positive integer."

bloqade/pyqrack/native.py

@@ -0,0 +1,49 @@
+import math
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from pyqrack import Pauli
+from bloqade.pyqrack import PyQrackQubit
+from bloqade.pyqrack.base import PyQrackInterpreter
+from bloqade.native.dialects.gate import stmts
+
+
+@stmts.dialect.register(key="pyqrack")
+class NativeMethods(interp.MethodTable):
+
+    @interp.impl(stmts.CZ)
+    def cz(self, _interp: PyQrackInterpreter, frame: interp.Frame, stmt: stmts.CZ):
+        controls = frame.get_casted(stmt.controls, ilist.IList[PyQrackQubit, Any])
+        targets = frame.get_casted(stmt.targets, ilist.IList[PyQrackQubit, Any])
+
+        for ctrl, trgt in zip(controls, targets):
+            if ctrl.is_active() and trgt.is_active():
+                ctrl.sim_reg.mcz([ctrl.addr], trgt.addr)
+
+        return ()
+
+    @interp.impl(stmts.R)
+    def r(self, _interp: PyQrackInterpreter, frame: interp.Frame, stmt: stmts.R):
+        qubits = frame.get_casted(stmt.qubits, ilist.IList[PyQrackQubit, Any])
+        rotation_angle = 2 * math.pi * frame.get_casted(stmt.rotation_angle, float)
+        axis_angle = 2 * math.pi * frame.get_casted(stmt.axis_angle, float)
+        for qubit in qubits:
+            if qubit.is_active():
+                qubit.sim_reg.r(Pauli.PauliZ, axis_angle, qubit.addr)
+                qubit.sim_reg.r(Pauli.PauliX, rotation_angle, qubit.addr)
+                qubit.sim_reg.r(Pauli.PauliZ, -axis_angle, qubit.addr)
+
+        return ()
+
+    @interp.impl(stmts.Rz)
+    def rz(self, _interp: PyQrackInterpreter, frame: interp.Frame, stmt: stmts.Rz):
+        qubits = frame.get_casted(stmt.qubits, ilist.IList[PyQrackQubit, Any])
+        rotation_angle = 2 * math.pi * frame.get_casted(stmt.rotation_angle, float)
+
+        for qubit in qubits:
+            if qubit.is_active():
+                qubit.sim_reg.r(Pauli.PauliZ, rotation_angle, qubit.addr)
+
+        return ()

bloqade/pyqrack/reg.py

@@ -2,15 +2,20 @@ import enum
 from typing import TYPE_CHECKING
 from dataclasses import dataclass
 
+from bloqade.types import MeasurementResult
 from bloqade.qasm2.types import Qubit
 
 if TYPE_CHECKING:
     from pyqrack import QrackSimulator
 
 
-class Measurement(enum.IntEnum):
+class Measurement(MeasurementResult, enum.IntEnum):
     """Enumeration of measurement results."""
 
+    def __init__(self, measurement_id: int = 0) -> None:
+        super().__init__()
+        self.measurement_id = measurement_id
+
     Zero = 0
     One = 1
     Lost = enum.auto()
@@ -70,8 +75,3 @@ class PyQrackQubit(Qubit):
     def drop(self):
         """Drop the qubit in-place."""
         self.state = QubitState.Lost
-
-
-@dataclass
-class PyQrackWire:
-    qubit: PyQrackQubit

bloqade/pyqrack/squin/gate/__init__.py

@@ -0,0 +1 @@
+from . import gate as gate

bloqade/pyqrack/squin/gate/gate.py

@@ -0,0 +1,136 @@
+import math
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from bloqade.squin import gate
+from pyqrack.pauli import Pauli
+from bloqade.pyqrack.reg import PyQrackQubit
+from bloqade.pyqrack.target import PyQrackInterpreter
+from bloqade.squin.gate.stmts import (
+    CX,
+    CY,
+    CZ,
+    U3,
+    H,
+    S,
+    T,
+    X,
+    Y,
+    Z,
+    Rx,
+    Ry,
+    Rz,
+    SqrtX,
+    SqrtY,
+)
+
+
+@gate.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+
+    @interp.impl(X)
+    @interp.impl(Y)
+    @interp.impl(Z)
+    @interp.impl(H)
+    def single_qubit_gate(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: X | Y | Z | H
+    ):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        method_name = stmt.name.lower()
+        for qbit in qubits:
+            if qbit.is_active():
+                getattr(qbit.sim_reg, method_name)(qbit.addr)
+
+    @interp.impl(T)
+    @interp.impl(S)
+    def single_qubit_nh_gate(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: S | T
+    ):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+
+        method_name = stmt.name.lower()
+        if stmt.adjoint:
+            method_name = "adj" + method_name
+
+        for qbit in qubits:
+            if qbit.is_active():
+                getattr(qbit.sim_reg, method_name)(qbit.addr)
+                qbit.sim_reg.r
+
+    @interp.impl(SqrtX)
+    @interp.impl(SqrtY)
+    def sqrt_x(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: SqrtX | SqrtY
+    ):
+        angle = math.pi / 2
+
+        if isinstance(stmt, SqrtX):
+            axis = Pauli.PauliX
+        else:
+            angle *= -1
+            axis = Pauli.PauliY
+
+        if stmt.adjoint:
+            angle *= -1
+
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            if qbit.is_active():
+                qbit.sim_reg.r(axis, angle, qbit.addr)
+
+    @interp.impl(Rx)
+    @interp.impl(Ry)
+    @interp.impl(Rz)
+    def rot(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: Rx | Ry | Rz):
+        match stmt:
+            case Rx():
+                axis = Pauli.PauliX
+            case Ry():
+                axis = Pauli.PauliY
+            case Rz():
+                axis = Pauli.PauliZ
+
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+
+        # NOTE: convert turns to radians
+        angle = frame.get(stmt.angle) * 2 * math.pi
+
+        for qbit in qubits:
+            if qbit.is_active():
+                qbit.sim_reg.r(axis, angle, qbit.addr)
+
+    @interp.impl(CX)
+    @interp.impl(CY)
+    @interp.impl(CZ)
+    def control(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: CX | CY | CZ
+    ):
+        controls: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.controls)
+        targets: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.targets)
+
+        if len(controls) != len(targets):
+            raise RuntimeError(
+                f"Found {len(controls)} controls but {len(targets)} targets when trying to evaluate {stmt}."
+            )
+
+        # NOTE: pyqrack convention "multi-control-x"
+        method_name = "m" + stmt.name.lower()
+
+        for control, target in zip(controls, targets):
+            if control.is_active() and target.is_active():
+                getattr(control.sim_reg, method_name)([control.addr], target.addr)
+
+    @interp.impl(U3)
+    def u3(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: U3):
+        theta = frame.get(stmt.theta) * 2 * math.pi
+        phi = frame.get(stmt.phi) * 2 * math.pi
+        lam = frame.get(stmt.lam) * 2 * math.pi
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+
+        for qbit in qubits:
+            if not qbit.is_active():
+                continue
+
+            qbit.sim_reg.u(qbit.addr, theta, phi, lam)