Move generalized tableau interpreter and simulator over from bloqade-internal

ppvm-python/pyproject.toml

@@ -8,6 +8,8 @@ authors = [
 ]
 requires-python = ">=3.10"
 dependencies = [
+    "bloqade-circuit>=0.14.1",
+    "kirin>=0.0.15",
     "ppvm-python-native",
 ]
 

ppvm-python/src/ppvm/__init__.py

@@ -2,3 +2,7 @@
 from .generalized_tableau import MeasurementResult as MeasurementResult
 from .paulisum import LossyPauliSum as LossyPauliSum
 from .paulisum import PauliSum as PauliSum
+from .squin_interpreter.device import (
+    GeneralizedTableauSimulator as GeneralizedTableauSimulator,
+    GeneralizedTableauSimulatorTask as GeneralizedTableauSimulatorTask,
+)

ppvm-python/src/ppvm/squin_interpreter/__init__.py

@@ -0,0 +1,8 @@
+"""TODO: once we open-source, all of this will be moved into bloqade-circuit"""
+
+from .impls import gate as gate, noise as noise
+from .device import (
+    GeneralizedTableauSimulator as GeneralizedTableauSimulator,
+    GeneralizedTableauSimulatorTask as GeneralizedTableauSimulatorTask,
+)
+from ._interp import GeneralizedTableauInterpreter as GeneralizedTableauInterpreter

ppvm-python/src/ppvm/squin_interpreter/_interp.py

@@ -0,0 +1,30 @@
+from dataclasses import field, dataclass
+
+from ..generalized_tableau import GeneralizedTableau
+from kirin import interp
+from typing_extensions import Self
+
+
+@dataclass
+class GeneralizedTableauInterpreter(interp.Interpreter):
+    keys = ("generalized_tableau", "main")
+    backend: GeneralizedTableau
+
+    rng_seed: int | None = None
+
+    current_qubit_addr: int = field(init=False, default=0)
+
+    def initialize(self) -> Self:
+        super().initialize()
+        self.backend = GeneralizedTableau(
+            n_qubits=self.backend.n_qubits,
+            min_abs_coeff=self.backend.min_abs_coeff,
+            seed=self.rng_seed,
+        )
+        self.current_qubit_addr = 0
+        return self
+
+    def allocate_qubit(self) -> int:
+        addr = self.current_qubit_addr
+        self.current_qubit_addr += 1
+        return addr

ppvm-python/src/ppvm/squin_interpreter/device.py

@@ -0,0 +1,64 @@
+from typing import Any, TypeVar, ParamSpec, cast
+from dataclasses import field, dataclass
+
+from ..generalized_tableau import GeneralizedTableau
+from kirin import ir
+from bloqade.device import AbstractSimulatorTask, AbstractSimulatorDevice
+from bloqade.analysis.address import UnknownReg, UnknownQubit, AddressAnalysis
+
+from ._interp import GeneralizedTableauInterpreter
+
+RetType = TypeVar("RetType")
+Param = ParamSpec("Param")
+InternalState = TypeVar("InternalState")
+
+
+@dataclass
+class GeneralizedTableauSimulatorTask(
+    AbstractSimulatorTask[Param, RetType, GeneralizedTableau]
+):
+    generalized_tableau_interp: GeneralizedTableauInterpreter
+
+    def run(self) -> RetType:
+        _, ret = self.generalized_tableau_interp.run(
+            self.kernel,
+            *self.args,
+            **self.kwargs,
+        )
+        return cast(RetType, ret)
+
+    @property
+    def state(self) -> GeneralizedTableau:
+        return self.generalized_tableau_interp.backend
+
+
+@dataclass
+class GeneralizedTableauSimulator(
+    AbstractSimulatorDevice[GeneralizedTableauSimulatorTask]
+):
+    n_qubits: int | None = None
+    options: dict[str, Any] = field(default_factory=dict)
+
+    def task(
+        self,
+        kernel: ir.Method[Param, RetType],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ):
+        address_analysis = AddressAnalysis(dialects=kernel.dialects)
+        frame, _ = address_analysis.run(kernel)
+
+        if self.n_qubits is None and any(
+            isinstance(a, (UnknownQubit, UnknownReg)) for a in frame.entries.values()
+        ):
+            raise ValueError(
+                "All addresses must be resolved. Or set n_qubits to a positive integer."
+            )
+
+        n_qubits = max(self.n_qubits or 0, address_analysis.qubit_count)
+
+        tab = GeneralizedTableau(n_qubits=n_qubits, **self.options)
+        interp = GeneralizedTableauInterpreter(
+            kernel.dialects, backend=tab, rng_seed=self.options.get("seed")
+        )
+        return GeneralizedTableauSimulatorTask(kernel, args, kwargs or {}, interp)

ppvm-python/src/ppvm/squin_interpreter/impls/gate.py

@@ -0,0 +1,126 @@
+import math
+from typing import Any
+
+from kirin import interp
+from bloqade.squin import gate
+from kirin.dialects import ilist
+
+from ..qubit import GeneralizedTableauQubit
+from .._interp import GeneralizedTableauInterpreter
+
+
+def _turns_to_radian(turns: float):
+    return 2 * turns * math.pi
+
+
+@gate.dialect.register(key="generalized_tableau")
+class GateMethods(interp.MethodTable):
+
+    @interp.impl(gate.stmts.X)
+    @interp.impl(gate.stmts.Y)
+    @interp.impl(gate.stmts.Z)
+    @interp.impl(gate.stmts.H)
+    def single_qubit_gate(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: gate.stmts.X | gate.stmts.Y | gate.stmts.Z | gate.stmts.H,
+    ):
+        qubits: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.qubits)
+        method = getattr(interp_.backend, stmt.name.lower())
+        for qbit in qubits:
+            method(qbit.addr)
+
+    @interp.impl(gate.stmts.S)
+    @interp.impl(gate.stmts.T)
+    @interp.impl(gate.stmts.SqrtX)
+    @interp.impl(gate.stmts.SqrtY)
+    def single_qubit_nh_gate(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: gate.stmts.S | gate.stmts.T | gate.stmts.SqrtX | gate.stmts.SqrtY,
+    ):
+        method_name = stmt.name.lower()
+
+        if stmt.adjoint:
+            method_name += "_adj"
+
+        method = getattr(interp_.backend, method_name)
+        qubits: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            method(qbit.addr)
+
+    @interp.impl(gate.stmts.Rx)
+    @interp.impl(gate.stmts.Ry)
+    @interp.impl(gate.stmts.Rz)
+    def rot(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: gate.stmts.Rx | gate.stmts.Ry | gate.stmts.Rz,
+    ):
+        method = getattr(interp_.backend, stmt.name.lower())
+
+        qubits: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.qubits)
+        angle = _turns_to_radian(frame.get(stmt.angle))
+
+        for qbit in qubits:
+            method(qbit.addr, angle)
+
+    @interp.impl(gate.stmts.CX)
+    @interp.impl(gate.stmts.CY)
+    @interp.impl(gate.stmts.CZ)
+    def control(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: gate.stmts.CX | gate.stmts.CY | gate.stmts.CZ,
+    ):
+        method_name = stmt.name.lower()
+        if method_name == "cx":
+            method_name = "cnot"
+
+        method = getattr(interp_.backend, method_name)
+
+        controls: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.controls)
+        targets: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.targets)
+
+        for ctrl, target in zip(controls, targets):
+            method(ctrl.addr, target.addr)
+
+    @interp.impl(gate.stmts.U3)
+    def u3(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: gate.stmts.U3,
+    ):
+        qubits: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.qubits)
+        theta = _turns_to_radian(frame.get(stmt.theta))
+        phi = _turns_to_radian(frame.get(stmt.phi))
+        lam = _turns_to_radian(frame.get(stmt.lam))
+
+        for qbit in qubits:
+            interp_.backend.u3(qbit.addr, theta, phi, lam)
+
+    @interp.impl(gate.stmts.PhasedXZ)
+    def phased_xz(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: gate.stmts.PhasedXZ,
+    ):
+        x_exponent = frame.get(stmt.x_exponent)
+        z_exponent = frame.get(stmt.z_exponent)
+        axis_phase_exponent = frame.get(stmt.axis_phase_exponent)
+        qubits: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.qubits)
+
+        angle_rz_pre = -axis_phase_exponent * math.pi * 2
+        angle_rx = x_exponent * math.pi * 2
+        angle_rz_post = (axis_phase_exponent + z_exponent) * math.pi * 2
+
+        for qbit in qubits:
+            interp_.backend.rz(qbit.addr, angle_rz_pre)
+            interp_.backend.rx(qbit.addr, angle_rx)
+            interp_.backend.rz(qbit.addr, angle_rz_post)

ppvm-python/src/ppvm/squin_interpreter/impls/noise.py

@@ -0,0 +1,94 @@
+from typing import Any
+
+from kirin import interp
+from bloqade.squin import noise
+from kirin.dialects import ilist
+
+from ..qubit import GeneralizedTableauQubit
+from .._interp import GeneralizedTableauInterpreter
+
+
+@noise.dialect.register(key="generalized_tableau")
+class NoiseMethods(interp.MethodTable):
+    @interp.impl(noise.stmts.Depolarize)
+    def depolarize(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: noise.stmts.Depolarize,
+    ):
+        p = frame.get(stmt.p)
+        qubits: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            interp_.backend.depolarize(qbit.addr, p)
+
+    @interp.impl(noise.stmts.Depolarize2)
+    def depolarize2(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: noise.stmts.Depolarize2,
+    ):
+        p = frame.get(stmt.p)
+        controls: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.controls)
+        targets: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.targets)
+        for control, target in zip(controls, targets):
+            interp_.backend.depolarize2(control.addr, target.addr, p)
+
+    @interp.impl(noise.stmts.SingleQubitPauliChannel)
+    def single_qubit_pauli_channel(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: noise.stmts.SingleQubitPauliChannel,
+    ):
+        px = frame.get(stmt.px)
+        py = frame.get(stmt.py)
+        pz = frame.get(stmt.pz)
+        qubits: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.qubits)
+
+        ps = [px, py, pz]
+        for qbit in qubits:
+            interp_.backend.pauli_error(qbit.addr, ps)
+
+    @interp.impl(noise.stmts.TwoQubitPauliChannel)
+    def two_qubit_pauli_channel(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: noise.stmts.TwoQubitPauliChannel,
+    ):
+        ps = frame.get(stmt.probabilities)
+        controls: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.controls)
+        targets: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.targets)
+
+        for control, target in zip(controls, targets):
+            interp_.backend.two_qubit_pauli_error(control.addr, target.addr, ps)
+
+    @interp.impl(noise.stmts.QubitLoss)
+    def qubit_loss(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: noise.stmts.QubitLoss,
+    ):
+        p = frame.get(stmt.p)
+        qubits: ilist.IList[GeneralizedTableauQubit, Any] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            interp_.backend.loss_channel(qbit.addr, p)
+
+    @interp.impl(noise.stmts.CorrelatedQubitLoss)
+    def correlated_qubit_loss(
+        self,
+        interp_: GeneralizedTableauInterpreter,
+        frame: interp.Frame,
+        stmt: noise.stmts.CorrelatedQubitLoss,
+    ):
+        p = frame.get(stmt.p)
+        ps = [p, 0, 0]
+        qubits: list[list[GeneralizedTableauQubit]] = frame.get(stmt.qubits)
+        for qubit_group in qubits:
+            if len(qubit_group) != 2:
+                raise ValueError("Correlated loss only supported for qubit pairs!")
+            addr0, addr1 = [qbit.addr for qbit in qubit_group]
+            interp_.backend.correlated_loss_channel(addr0, addr1, ps)