Implement runtime for squin noise statements (#273)

Closes #227

Author: david-pl

src/bloqade/pyqrack/__init__.py

@@ -16,7 +16,7 @@
 # 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, qubit as qubit
+from .squin import op as op, noise as noise, qubit as qubit
 from .device import (
     StackMemorySimulator as StackMemorySimulator,
     DynamicMemorySimulator as DynamicMemorySimulator,

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

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

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

@@ -0,0 +1,72 @@
+import random
+import typing
+from dataclasses import dataclass
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from bloqade.pyqrack import QubitState, PyQrackQubit, PyQrackInterpreter
+from bloqade.squin.noise.stmts import QubitLoss, StochasticUnitaryChannel
+from bloqade.squin.noise._dialect import dialect as squin_noise_dialect
+
+from ..runtime import OperatorRuntimeABC
+
+
+@dataclass(frozen=True)
+class StochasticUnitaryChannelRuntime(OperatorRuntimeABC):
+    operators: ilist.IList[OperatorRuntimeABC, typing.Any]
+    probabilities: ilist.IList[float, typing.Any]
+
+    @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(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),)
+
+    @interp.impl(QubitLoss)
+    def qubit_loss(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: QubitLoss
+    ):
+        p = frame.get(stmt.p)
+        return (QubitLossRuntime(p),)

src/bloqade/squin/groups.py

@@ -3,12 +3,12 @@
 from kirin.rewrite import Walk, Chain
 from kirin.dialects import ilist
 
-from . import op, wire, qubit
+from . import op, wire, noise, qubit
 from .op.rewrite import PyMultToSquinMult
 from .rewrite.desugar import ApplyDesugarRule, MeasureDesugarRule
 
 
-@ir.dialect_group(structural_no_opt.union([op, qubit]))
+@ir.dialect_group(structural_no_opt.union([op, qubit, noise]))
 def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
@@ -36,7 +36,7 @@ def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
     return run_pass
 
 
-@ir.dialect_group(structural_no_opt.union([op, wire]))
+@ir.dialect_group(structural_no_opt.union([op, wire, noise]))
 def wired(self):
     py_mult_to_mult_pass = PyMultToSquinMult(self)
 

src/bloqade/squin/noise/__init__.py

@@ -4,5 +4,7 @@
     pp_error as pp_error,
     depolarize as depolarize,
     qubit_loss as qubit_loss,
-    pauli_channel as pauli_channel,
+    pauli_error as pauli_error,
+    two_qubit_pauli_channel as two_qubit_pauli_channel,
+    single_qubit_pauli_channel as single_qubit_pauli_channel,
 )

src/bloqade/squin/noise/_wrapper.py

@@ -14,11 +14,15 @@ def pp_error(op: Op, p: float) -> Op: ...
 
 
 @wraps(stmts.Depolarize)
-def depolarize(n_qubits: int, p: float) -> Op: ...
+def depolarize(p: float) -> Op: ...
 
 
-@wraps(stmts.PauliChannel)
-def pauli_channel(n_qubits: int, params: tuple[float, ...]) -> Op: ...
+@wraps(stmts.SingleQubitPauliChannel)
+def single_qubit_pauli_channel(params: tuple[float, float, float]) -> Op: ...
+
+
+@wraps(stmts.TwoQubitPauliChannel)
+def two_qubit_pauli_channel(params: tuple[float, ...]) -> Op: ...
 
 
 @wraps(stmts.QubitLoss)

src/bloqade/squin/noise/rewrite.py

@@ -0,0 +1,111 @@
+import itertools
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import Walk
+from kirin.dialects import ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from .stmts import (
+    PPError,
+    QubitLoss,
+    Depolarize,
+    PauliError,
+    NoiseChannel,
+    TwoQubitPauliChannel,
+    SingleQubitPauliChannel,
+    StochasticUnitaryChannel,
+)
+from ..op.stmts import X, Y, Z, Kron, Identity
+
+
+class _RewriteNoiseStmts(RewriteRule):
+    """Rewrites squin noise statements to StochasticUnitaryChannel"""
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, NoiseChannel) or isinstance(node, QubitLoss):
+            return RewriteResult()
+
+        return getattr(self, "rewrite_" + node.name)(node)
+
+    def rewrite_pauli_error(self, node: PauliError) -> RewriteResult:
+        (operators := ilist.New(values=(node.basis,))).insert_before(node)
+        (ps := ilist.New(values=(node.p,))).insert_before(node)
+        stochastic_channel = StochasticUnitaryChannel(
+            operators=operators.result, probabilities=ps.result
+        )
+
+        node.replace_by(stochastic_channel)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_single_qubit_pauli_channel(
+        self, node: SingleQubitPauliChannel
+    ) -> RewriteResult:
+        paulis = (X(), Y(), Z())
+        paulis_ssa: list[ir.SSAValue] = []
+        for op in paulis:
+            op.insert_before(node)
+            paulis_ssa.append(op.result)
+
+        (pauli_ops := ilist.New(values=paulis_ssa)).insert_before(node)
+
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=pauli_ops.result, probabilities=node.params
+        )
+        node.replace_by(stochastic_unitary)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_two_qubit_pauli_channel(
+        self, node: TwoQubitPauliChannel
+    ) -> RewriteResult:
+        paulis = (X(), Y(), Z(), Identity(sites=1))
+        for op in paulis:
+            op.insert_before(node)
+
+        # NOTE: collect list so we can skip the last entry, which will be two identities
+        combinations = list(itertools.product(paulis, repeat=2))[:-1]
+        operators: list[ir.SSAValue] = []
+        for pauli_1, pauli_2 in combinations:
+            op = Kron(pauli_1.result, pauli_2.result)
+            op.insert_before(node)
+            operators.append(op.result)
+
+        (operator_list := ilist.New(values=operators)).insert_before(node)
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=operator_list.result, probabilities=node.params
+        )
+
+        node.replace_by(stochastic_unitary)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_p_p_error(self, node: PPError) -> RewriteResult:
+        (operators := ilist.New(values=(node.op,))).insert_before(node)
+        (ps := ilist.New(values=(node.p,))).insert_before(node)
+        stochastic_channel = StochasticUnitaryChannel(
+            operators=operators.result, probabilities=ps.result
+        )
+
+        node.replace_by(stochastic_channel)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_depolarize(self, node: Depolarize) -> RewriteResult:
+        paulis = (X(), Y(), Z())
+        operators: list[ir.SSAValue] = []
+        for op in paulis:
+            op.insert_before(node)
+            operators.append(op.result)
+
+        (operator_list := ilist.New(values=operators)).insert_before(node)
+        (ps := ilist.New(values=[node.p for _ in range(3)])).insert_before(node)
+
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=operator_list.result, probabilities=ps.result
+        )
+        node.replace_by(stochastic_unitary)
+
+        return RewriteResult(has_done_something=True)
+
+
+class RewriteNoiseStmts(Pass):
+    def unsafe_run(self, mt: ir.Method):
+        return Walk(_RewriteNoiseStmts()).rewrite(mt.code)

src/bloqade/squin/noise/stmts.py

@@ -1,21 +1,23 @@
-from kirin import ir, types
+from kirin import ir, types, lowering
 from kirin.decl import info, statement
+from kirin.dialects import ilist
 
 from bloqade.squin.op.types import OpType
 
 from ._dialect import dialect
+from ..op.types import NumOperators
 
 
 @statement
 class NoiseChannel(ir.Statement):
-    pass
+    traits = frozenset({lowering.FromPythonCall()})
+    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
 class PauliError(NoiseChannel):
     basis: ir.SSAValue = info.argument(OpType)
     p: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -26,34 +28,37 @@ class PPError(NoiseChannel):
 
     op: ir.SSAValue = info.argument(OpType)
     p: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
 class Depolarize(NoiseChannel):
     """
-    Apply n-qubit depolaize error to qubits
-    NOTE For Stim, this can only accept 1 or 2 qubits
+    Apply depolarize error to qubit
     """
 
-    n_qubits: int = info.attribute(types.Int)
     p: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
-class PauliChannel(NoiseChannel):
-    # NOTE:
-    # 1-qubit 3 params px, py, pz
-    # 2-qubit 15 params pix, piy, piz, pxi, pxx, pxy, pxz, pyi, pyx ..., pzz
-    # TODO add validation for params (maybe during lowering via custom lower?)
-    n_qubits: int = info.attribute()
-    params: ir.SSAValue = info.argument(types.Tuple[types.Vararg(types.Float)])
-    result: ir.ResultValue = info.result(OpType)
+class SingleQubitPauliChannel(NoiseChannel):
+    params: ir.SSAValue = info.argument(ilist.IListType[types.Float, types.Literal(3)])
+
+
+@statement(dialect=dialect)
+class TwoQubitPauliChannel(NoiseChannel):
+    params: ir.SSAValue = info.argument(ilist.IListType[types.Float, types.Literal(15)])
 
 
 @statement(dialect=dialect)
 class QubitLoss(NoiseChannel):
     # NOTE: qubit loss error (not supported by Stim)
     p: ir.SSAValue = info.argument(types.Float)
+
+
+@statement(dialect=dialect)
+class StochasticUnitaryChannel(ir.Statement):
+    operators: ir.SSAValue = info.argument(ilist.IListType[OpType, NumOperators])
+    probabilities: ir.SSAValue = info.argument(
+        ilist.IListType[types.Float, NumOperators]
+    )
     result: ir.ResultValue = info.result(OpType)

src/bloqade/squin/op/types.py

@@ -22,3 +22,5 @@ def __rmul__(self, other: complex) -> "Op":
 
 
 OpType = types.PyClass(Op)
+
+NumOperators = types.TypeVar("NumOperators")