Refactor & simplify noise dialect (#512)

Author: david-pl

src/bloqade/cirq_utils/emit/noise.py

@@ -16,16 +16,6 @@
 @noise.dialect.register(key="emit.cirq")
 class EmitCirqNoiseMethods(MethodTable):
 
-    @impl(noise.stmts.PauliError)
-    def pauli_error(
-        self, interp: EmitCirq, frame: EmitCirqFrame, stmt: noise.stmts.PauliError
-    ):
-        op = frame.get(stmt.basis)
-        p = frame.get(stmt.p)
-        error_probabilities = {self._op_to_key(op): p}
-        gate = cirq.asymmetric_depolarize(error_probabilities=error_probabilities)
-        return (BasicOpRuntime(gate=gate),)
-
     @impl(noise.stmts.Depolarize)
     def depolarize(
         self, interp: EmitCirq, frame: EmitCirqFrame, stmt: noise.stmts.Depolarize
@@ -72,20 +62,6 @@ def two_qubit_pauli_channel(
         gate = cirq.asymmetric_depolarize(error_probabilities=error_probabilities)
         return (BasicOpRuntime(gate),)
 
-    @impl(noise.stmts.StochasticUnitaryChannel)
-    def stochastic_unitary_channel(
-        self,
-        emit: EmitCirq,
-        frame: EmitCirqFrame,
-        stmt: noise.stmts.StochasticUnitaryChannel,
-    ):
-        ops = frame.get(stmt.operators)
-        ps = frame.get(stmt.probabilities)
-
-        error_probabilities = {self._op_to_key(op_): p for op_, p in zip(ops, ps)}
-        cirq_op = cirq.asymmetric_depolarize(error_probabilities=error_probabilities)
-        return (BasicOpRuntime(cirq_op),)
-
     @staticmethod
     def _op_to_key(operator: OperatorRuntimeABC) -> str:
         match operator:

src/bloqade/pyqrack/squin/noise/native.py

@@ -1,93 +1,57 @@
-import random
-import typing
-from functools import cached_property
-from dataclasses import dataclass
-
 from kirin import interp
-from kirin.dialects import ilist
 
-from bloqade.pyqrack import QubitState, PyQrackQubit, PyQrackInterpreter
+from bloqade.pyqrack import PyQrackQubit, PyQrackInterpreter
 from bloqade.squin.noise.stmts import (
     QubitLoss,
     Depolarize,
-    PauliError,
     Depolarize2,
     TwoQubitPauliChannel,
     SingleQubitPauliChannel,
-    StochasticUnitaryChannel,
 )
 from bloqade.squin.noise._dialect import dialect as squin_noise_dialect
 
-from ..runtime import KronRuntime, IdentityRuntime, OperatorRuntime, OperatorRuntimeABC
-
-
-@dataclass(frozen=True)
-class StochasticUnitaryChannelRuntime(OperatorRuntimeABC):
-    operators: (
-        ilist.IList[OperatorRuntimeABC, typing.Any] | tuple[OperatorRuntimeABC, ...]
-    )
-    probabilities: ilist.IList[float, typing.Any] | tuple[float, ...]
-
-    @property
-    def n_sites(self) -> int:
-        n = self.operators[0].n_sites
-        for op in self.operators[1:]:
-            assert (
-                op.n_sites == n
-            ), "Encountered a stochastic unitary channel with operators of different size!"
-        return n
-
-    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
-        # NOTE: probabilities don't necessarily sum to 1; could be no noise event should occur
-        p_no_op = 1 - sum(self.probabilities)
-        if random.uniform(0.0, 1.0) < p_no_op:
-            return
-
-        selected_ops = random.choices(self.operators, weights=self.probabilities)
-        for op in selected_ops:
-            op.apply(*qubits, adjoint=adjoint)
-
-
-@dataclass(frozen=True)
-class QubitLossRuntime(OperatorRuntimeABC):
-    p: float
-
-    @property
-    def n_sites(self) -> int:
-        return 1
-
-    def apply(self, qubit: PyQrackQubit, adjoint: bool = False) -> None:
-        if random.uniform(0.0, 1.0) <= self.p:
-            qubit.state = QubitState.Lost
-
 
 @squin_noise_dialect.register(key="pyqrack")
 class PyQrackMethods(interp.MethodTable):
-    @interp.impl(PauliError)
-    def pauli_error(
-        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: PauliError
-    ):
-        op = frame.get(stmt.basis)
-        p = frame.get(stmt.p)
-        return (StochasticUnitaryChannelRuntime((op,), (p,)),)
+
+    single_pauli_choices = ("i", "x", "y", "z")
+    two_pauli_choices = (
+        "ii",
+        "ix",
+        "iy",
+        "iz",
+        "xi",
+        "xx",
+        "xy",
+        "xz",
+        "yi",
+        "yx",
+        "yy",
+        "yz",
+        "zi",
+        "zx",
+        "zy",
+        "zz",
+    )
 
     @interp.impl(Depolarize)
     def depolarize(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: Depolarize
     ):
         p = frame.get(stmt.p)
-        ps = (p / 3.0,) * 3
-        ops = self.single_qubit_paulis
-        return (StochasticUnitaryChannelRuntime(ops, ps),)
+        ps = [p / 3.0] * 3
+        qubits = frame.get(stmt.qubits)
+        self.apply_single_qubit_pauli_error(interp, ps, qubits)
 
     @interp.impl(Depolarize2)
     def depolarize2(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: Depolarize2
     ):
         p = frame.get(stmt.p)
-        ps = (p / 15.0,) * 15
-        ops = self.two_qubit_paulis
-        return (StochasticUnitaryChannelRuntime(ops, ps),)
+        ps = [p / 15.0] * 15
+        controls = frame.get(stmt.controls)
+        targets = frame.get(stmt.targets)
+        self.apply_two_qubit_pauli_error(interp, ps, controls, targets)
 
     @interp.impl(SingleQubitPauliChannel)
     def single_qubit_pauli_channel(
@@ -96,9 +60,11 @@ def single_qubit_pauli_channel(
         frame: interp.Frame,
         stmt: SingleQubitPauliChannel,
     ):
-        ps = frame.get(stmt.params)
-        ops = self.single_qubit_paulis
-        return (StochasticUnitaryChannelRuntime(ops, ps),)
+        px = frame.get(stmt.px)
+        py = frame.get(stmt.py)
+        pz = frame.get(stmt.pz)
+        qubits = frame.get(stmt.qubits)
+        self.apply_single_qubit_pauli_error(interp, [px, py, pz], qubits)
 
     @interp.impl(TwoQubitPauliChannel)
     def two_qubit_pauli_channel(
@@ -107,43 +73,54 @@ def two_qubit_pauli_channel(
         frame: interp.Frame,
         stmt: TwoQubitPauliChannel,
     ):
-        ps = frame.get(stmt.params)
-        ops = self.two_qubit_paulis
-        return (StochasticUnitaryChannelRuntime(ops, ps),)
-
-    @interp.impl(StochasticUnitaryChannel)
-    def stochastic_unitary_channel(
-        self,
-        interp: PyQrackInterpreter,
-        frame: interp.Frame,
-        stmt: StochasticUnitaryChannel,
-    ):
-        operators = frame.get(stmt.operators)
-        probabilities = frame.get(stmt.probabilities)
-
-        return (StochasticUnitaryChannelRuntime(operators, probabilities),)
+        ps = frame.get(stmt.probabilities)
+        controls = frame.get(stmt.controls)
+        targets = frame.get(stmt.targets)
+        self.apply_two_qubit_pauli_error(interp, ps, controls, targets)
 
     @interp.impl(QubitLoss)
     def qubit_loss(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: QubitLoss
     ):
         p = frame.get(stmt.p)
-        return (QubitLossRuntime(p),)
+        qubits: list[PyQrackQubit] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            if interp.rng_state.uniform(0.0, 1.0) <= p:
+                qbit.drop()
 
-    @cached_property
-    def single_qubit_paulis(self):
-        return (OperatorRuntime("x"), OperatorRuntime("y"), OperatorRuntime("z"))
+    def apply_single_qubit_pauli_error(
+        self,
+        interp: PyQrackInterpreter,
+        ps: list[float],
+        qubits: list[PyQrackQubit],
+    ):
+        pi = 1 - sum(ps)
+        probs = [pi] + ps
 
-    @cached_property
-    def two_qubit_paulis(self):
-        paulis = (IdentityRuntime(sites=1), *self.single_qubit_paulis)
-        ops: list[KronRuntime] = []
-        for idx1, pauli1 in enumerate(paulis):
-            for idx2, pauli2 in enumerate(paulis):
-                if idx1 == idx2 == 0:
-                    # NOTE: 'II'
-                    continue
+        assert all(0 <= x <= 1 for x in probs), "Invalid Pauli error probabilities"
 
-                ops.append(KronRuntime(pauli1, pauli2))
+        for qbit in qubits:
+            which = interp.rng_state.choice(self.single_pauli_choices, p=probs)
+            self.apply_pauli_error(which, qbit)
+
+    def apply_two_qubit_pauli_error(
+        self,
+        interp: PyQrackInterpreter,
+        ps: list[float],
+        controls: list[PyQrackQubit],
+        targets: list[PyQrackQubit],
+    ):
+        pii = 1 - sum(ps)
+        probs = [pii] + ps
+        assert all(0 <= x <= 1 for x in probs), "Invalid Pauli error probabilities"
+
+        for control, target in zip(controls, targets):
+            which = interp.rng_state.choice(self.two_pauli_choices, p=probs)
+            self.apply_pauli_error(which[0], control)
+            self.apply_pauli_error(which[1], target)
+
+    def apply_pauli_error(self, which: str, qbit: PyQrackQubit):
+        if not qbit.is_active() or which == "i":
+            return
 
-        return tuple(ops)
+        getattr(qbit.sim_reg, which)(qbit.addr)

src/bloqade/squin/__init__.py

@@ -29,16 +29,21 @@
     sqrt_x as sqrt_x,
     sqrt_y as sqrt_y,
     sqrt_z as sqrt_z,
+    bit_flip as bit_flip,
+    depolarize as depolarize,
+    qubit_loss as qubit_loss,
     sqrt_x_adj as sqrt_x_adj,
     sqrt_y_adj as sqrt_y_adj,
     sqrt_z_adj as sqrt_z_adj,
+    depolarize2 as depolarize2,
+    two_qubit_pauli_channel as two_qubit_pauli_channel,
+    single_qubit_pauli_channel as single_qubit_pauli_channel,
 )
 
 # NOTE: it's important to keep these imports here since they import squin.kernel
 # we skip isort here
 from . import parallel as parallel  # isort: skip
 from .stdlib import (  # isort: skip
     gate as gate,
-    channel as channel,
     broadcast as broadcast,
 )

src/bloqade/squin/analysis/nsites/impls.py

@@ -124,34 +124,6 @@ def two_qubit_noise(
     ):
         return (NumberSites(sites=2),)
 
-    @interp.impl(noise.stmts.PauliError)
-    def pauli_error(
-        self, interp: NSitesAnalysis, frame: interp.Frame, stmt: noise.stmts.PauliError
-    ):
-        pauli_ops_sites = frame.get(stmt.basis)
-        return (pauli_ops_sites,)
-
-    @interp.impl(noise.stmts.StochasticUnitaryChannel)
-    def stochastic_unitary_noise(
-        self,
-        interp: NSitesAnalysis,
-        frame: interp.Frame,
-        stmt: noise.stmts.StochasticUnitaryChannel,
-    ):
-        ops = frame.get(stmt.operators)
-
-        # StochasticUnitaryChannel always accepts an IList of Operators
-        # but it's the number of sites of the individual operator themselves that should
-        # represent the sites the channel acts on.
-        if (
-            isinstance(ops, tuple)
-            and all(isinstance(op, NumberSites) for op in ops)
-            and len(set([op.sites for op in ops])) == 1
-        ):
-            return (ops[0],)
-
-        return (NoSites(),)
-
 
 @scf.dialect.register(key="op.nsites")
 class ScfSquinOp(ScfTypeInfer):

src/bloqade/squin/noise/__init__.py

@@ -1,11 +1,2 @@
-from . import stmts as stmts
+from . import stmts as stmts, _interface as _interface
 from ._dialect import dialect as dialect
-from ._wrapper import (
-    depolarize as depolarize,
-    qubit_loss as qubit_loss,
-    depolarize2 as depolarize2,
-    pauli_error as pauli_error,
-    two_qubit_pauli_channel as two_qubit_pauli_channel,
-    single_qubit_pauli_channel as single_qubit_pauli_channel,
-    stochastic_unitary_channel as stochastic_unitary_channel,
-)

src/bloqade/squin/noise/_dialect.py

@@ -1,3 +1,3 @@
 from kirin import ir
 
-dialect = ir.Dialect(name="squin.noise")
+dialect = ir.Dialect("squin.noise")

src/bloqade/squin/noise/_interface.py

@@ -0,0 +1,39 @@
+from typing import Any, Literal, TypeVar
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.types import Qubit
+
+from . import stmts
+
+
+@wraps(stmts.Depolarize)
+def depolarize(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+N = TypeVar("N", bound=int)
+
+
+@wraps(stmts.Depolarize2)
+def depolarize2(
+    p: float, controls: ilist.IList[Qubit, N], targets: ilist.IList[Qubit, N]
+) -> None: ...
+
+
+@wraps(stmts.SingleQubitPauliChannel)
+def single_qubit_pauli_channel(
+    px: float, py: float, pz: float, qubits: ilist.IList[Qubit, Any]
+) -> None: ...
+
+
+@wraps(stmts.TwoQubitPauliChannel)
+def two_qubit_pauli_channel(
+    probabilities: ilist.IList[float, Literal[15]],
+    controls: ilist.IList[Qubit, N],
+    targets: ilist.IList[Qubit, N],
+) -> None: ...
+
+
+@wraps(stmts.QubitLoss)
+def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ...