bloqade-circuit 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

bloqade/pyqrack/__init__.py

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

bloqade/pyqrack/squin/noise/__init__.py

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

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

bloqade/squin/__init__.py

@@ -6,3 +6,12 @@ from . import (
     lowering as lowering,
 )
 from .groups import wired as wired, kernel as kernel
+
+try:
+    # NOTE: make sure optional cirq dependency is installed
+    import cirq as cirq_package  # noqa: F401
+except ImportError:
+    pass
+else:
+    from . import cirq as cirq
+    from .cirq import load_circuit as load_circuit

bloqade/squin/cirq/__init__.py

@@ -0,0 +1,89 @@
+from typing import Any
+
+import cirq
+from kirin import ir, types
+from kirin.dialects import func
+
+from . import lowering as lowering
+from .. import kernel
+from .lowering import Squin
+
+
+def load_circuit(
+    circuit: cirq.Circuit,
+    kernel_name: str = "main",
+    dialects: ir.DialectGroup = kernel,
+    globals: dict[str, Any] | None = None,
+    file: str | None = None,
+    lineno_offset: int = 0,
+    col_offset: int = 0,
+    compactify: bool = True,
+):
+    """Converts a cirq.Circuit object into a squin kernel.
+
+    Args:
+        circuit (cirq.Circuit): The circuit to load.
+
+    Keyword Args:
+        kernel_name (str): The name of the kernel to load. Defaults to "main".
+        dialects (ir.DialectGroup | None): The dialects to use. Defaults to `squin.kernel`.
+        globals (dict[str, Any] | None): The global variables to use. Defaults to None.
+        file (str | None): The file name for error reporting. Defaults to None.
+        lineno_offset (int): The line number offset for error reporting. Defaults to 0.
+        col_offset (int): The column number offset for error reporting. Defaults to 0.
+        compactify (bool): Whether to compactify the output. Defaults to True.
+
+    Example:
+
+    ```python
+    # from cirq's "hello qubit" example
+    import cirq
+    from bloqade import squin
+
+    # Pick a qubit.
+    qubit = cirq.GridQubit(0, 0)
+
+    # Create a circuit.
+    circuit = cirq.Circuit(
+        cirq.X(qubit)**0.5,  # Square root of NOT.
+        cirq.measure(qubit, key='m')  # Measurement.
+    )
+
+    # load the circuit as squin
+    main = squin.load_circuit(circuit)
+
+    # print the resulting IR
+    main.print()
+    ```
+    """
+
+    target = Squin(dialects=dialects, circuit=circuit)
+    body = target.run(
+        circuit,
+        source=str(circuit),  # TODO: proper source string
+        file=file,
+        globals=globals,
+        lineno_offset=lineno_offset,
+        col_offset=col_offset,
+        compactify=compactify,
+    )
+
+    # NOTE: no return value
+    return_value = func.ConstantNone()
+    body.blocks[0].stmts.append(return_value)
+    body.blocks[0].stmts.append(func.Return(value_or_stmt=return_value))
+
+    code = func.Function(
+        sym_name=kernel_name,
+        signature=func.Signature((), types.NoneType),
+        body=body,
+    )
+
+    return ir.Method(
+        mod=None,
+        py_func=None,
+        sym_name=kernel_name,
+        arg_names=[],
+        dialects=dialects,
+        code=code,
+    )

bloqade/squin/cirq/lowering.py

@@ -0,0 +1,303 @@
+import math
+from typing import Any
+from dataclasses import field, dataclass
+
+import cirq
+from kirin import ir, lowering
+from kirin.rewrite import Walk, CFGCompactify
+from kirin.dialects import py, ilist
+
+from .. import op, noise, qubit
+
+CirqNode = cirq.Circuit | cirq.Moment | cirq.Gate | cirq.Qid | cirq.Operation
+
+DecomposeNode = (
+    cirq.SwapPowGate
+    | cirq.ISwapPowGate
+    | cirq.PhasedXPowGate
+    | cirq.PhasedXZGate
+    | cirq.CSwapGate
+)
+
+
+@dataclass
+class Squin(lowering.LoweringABC[CirqNode]):
+    """Lower a cirq.Circuit object to a squin kernel"""
+
+    circuit: cirq.Circuit
+    qreg: qubit.New = field(init=False)
+    qreg_index: dict[cirq.Qid, int] = field(init=False, default_factory=dict)
+    next_qreg_index: int = field(init=False, default=0)
+
+    def lower_qubit_getindex(self, state: lowering.State[CirqNode], qid: cirq.Qid):
+        index = self.qreg_index.get(qid)
+
+        if index is None:
+            index = self.next_qreg_index
+            self.qreg_index[qid] = index
+            self.next_qreg_index += 1
+
+        index_ssa = state.current_frame.push(py.Constant(index)).result
+        qbit_getitem = state.current_frame.push(py.GetItem(self.qreg.result, index_ssa))
+        return qbit_getitem.result
+
+    def lower_qubit_getindices(
+        self, state: lowering.State[CirqNode], qids: list[cirq.Qid]
+    ):
+        qbits_getitem = [self.lower_qubit_getindex(state, qid) for qid in qids]
+        qbits_stmt = ilist.New(values=qbits_getitem)
+        qbits_result = state.current_frame.get(qbits_stmt.name)
+
+        if qbits_result is not None:
+            return qbits_result
+
+        state.current_frame.push(qbits_stmt)
+        return qbits_stmt.result
+
+    def run(
+        self,
+        stmt: CirqNode,
+        *,
+        source: str | None = None,
+        globals: dict[str, Any] | None = None,
+        file: str | None = None,
+        lineno_offset: int = 0,
+        col_offset: int = 0,
+        compactify: bool = True,
+    ) -> ir.Region:
+
+        state = lowering.State(
+            self,
+            file=file,
+            lineno_offset=lineno_offset,
+            col_offset=col_offset,
+        )
+
+        with state.frame(
+            [stmt],
+            globals=globals,
+            finalize_next=False,
+        ) as frame:
+            # NOTE: create a global register of qubits first
+            # TODO: can there be a circuit without qubits?
+            n_qubits = cirq.num_qubits(self.circuit)
+            n = frame.push(py.Constant(n_qubits))
+            self.qreg = frame.push(qubit.New(n_qubits=n.result))
+
+            self.visit(state, stmt)
+
+            if compactify:
+                Walk(CFGCompactify()).rewrite(frame.curr_region)
+
+            region = frame.curr_region
+
+        return region
+
+    def visit(self, state: lowering.State[CirqNode], node: CirqNode) -> lowering.Result:
+        name = node.__class__.__name__
+        return getattr(self, f"visit_{name}", self.generic_visit)(state, node)
+
+    def generic_visit(self, state: lowering.State[CirqNode], node: CirqNode):
+        if isinstance(node, CirqNode):
+            raise lowering.BuildError(
+                f"Cannot lower {node.__class__.__name__} node: {node}"
+            )
+        raise lowering.BuildError(
+            f"Unexpected `{node.__class__.__name__}` node: {repr(node)} is not an AST node"
+        )
+
+    def lower_literal(self, state: lowering.State[CirqNode], value) -> ir.SSAValue:
+        raise lowering.BuildError("Literals not supported in cirq circuit")
+
+    def lower_global(
+        self, state: lowering.State[CirqNode], node: CirqNode
+    ) -> lowering.LoweringABC.Result:
+        raise lowering.BuildError("Literals not supported in cirq circuit")
+
+    def visit_Circuit(
+        self, state: lowering.State[CirqNode], node: cirq.Circuit
+    ) -> lowering.Result:
+        for moment in node:
+            state.lower(moment)
+
+    def visit_Moment(
+        self, state: lowering.State[CirqNode], node: cirq.Moment
+    ) -> lowering.Result:
+        for op_ in node.operations:
+            state.lower(op_)
+
+    def visit_GateOperation(
+        self, state: lowering.State[CirqNode], node: cirq.GateOperation
+    ):
+        if isinstance(node.gate, cirq.MeasurementGate):
+            # NOTE: special dispatch here, since measurement is a gate + a qubit in cirq,
+            # but a single statement in squin
+            return self.lower_measurement(state, node)
+
+        if isinstance(node.gate, DecomposeNode):
+            # NOTE: easier to decompose these, but for that we need the qubits too,
+            # so we need to do this within this method
+            for subnode in cirq.decompose_once(node):
+                state.lower(subnode)
+            return
+
+        op_ = state.lower(node.gate).expect_one()
+        qbits = self.lower_qubit_getindices(state, node.qubits)
+        return state.current_frame.push(qubit.Apply(operator=op_, qubits=qbits))
+
+    def lower_measurement(
+        self, state: lowering.State[CirqNode], node: cirq.GateOperation
+    ):
+        if len(node.qubits) == 1:
+            qbit = self.lower_qubit_getindex(state, node.qubits[0])
+            return state.current_frame.push(qubit.MeasureQubit(qbit))
+
+        qbits = self.lower_qubit_getindices(state, node.qubits)
+        return state.current_frame.push(qubit.MeasureQubitList(qbits))
+
+    def visit_SingleQubitPauliStringGateOperation(
+        self,
+        state: lowering.State[CirqNode],
+        node: cirq.SingleQubitPauliStringGateOperation,
+    ):
+
+        match node.pauli:
+            case cirq.X:
+                op_ = op.stmts.X()
+            case cirq.Y:
+                op_ = op.stmts.Y()
+            case cirq.Z:
+                op_ = op.stmts.Z()
+            case cirq.I:
+                op_ = op.stmts.Identity(sites=1)
+            case _:
+                raise lowering.BuildError(f"Unexpected Pauli operation {node.pauli}")
+
+        state.current_frame.push(op_)
+        qargs = self.lower_qubit_getindices(state, [node.qubit])
+        return state.current_frame.push(qubit.Apply(op_.result, qargs))
+
+    def visit_HPowGate(self, state: lowering.State[CirqNode], node: cirq.HPowGate):
+        if node.exponent == 1:
+            return state.current_frame.push(op.stmts.H())
+
+        return state.lower(node.in_su2())
+
+    def visit_XPowGate(self, state: lowering.State[CirqNode], node: cirq.XPowGate):
+        if node.exponent == 1:
+            return state.current_frame.push(op.stmts.X())
+
+        return self.visit(state, node.in_su2())
+
+    def visit_YPowGate(self, state: lowering.State[CirqNode], node: cirq.YPowGate):
+        if node.exponent == 1:
+            return state.current_frame.push(op.stmts.Y())
+
+        return self.visit(state, node.in_su2())
+
+    def visit_ZPowGate(self, state: lowering.State[CirqNode], node: cirq.ZPowGate):
+        if node.exponent == 0.5:
+            return state.current_frame.push(op.stmts.S())
+
+        if node.exponent == 0.25:
+            return state.current_frame.push(op.stmts.T())
+
+        if node.exponent == 1:
+            return state.current_frame.push(op.stmts.Z())
+
+        # NOTE: just for the Z gate, an arbitrary exponent is equivalent to the ShiftOp
+        t = node.exponent
+        theta = state.current_frame.push(py.Constant(math.pi * t))
+        return state.current_frame.push(op.stmts.ShiftOp(theta=theta.result))
+
+    def visit_Rx(self, state: lowering.State[CirqNode], node: cirq.Rx):
+        x = state.current_frame.push(op.stmts.X())
+        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
+        return state.current_frame.push(op.stmts.Rot(axis=x.result, angle=angle.result))
+
+    def visit_Ry(self, state: lowering.State[CirqNode], node: cirq.Ry):
+        y = state.current_frame.push(op.stmts.Y())
+        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
+        return state.current_frame.push(op.stmts.Rot(axis=y.result, angle=angle.result))
+
+    def visit_Rz(self, state: lowering.State[CirqNode], node: cirq.Rz):
+        z = state.current_frame.push(op.stmts.Z())
+        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
+        return state.current_frame.push(op.stmts.Rot(axis=z.result, angle=angle.result))
+
+    def visit_CXPowGate(self, state: lowering.State[CirqNode], node: cirq.CXPowGate):
+        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Control(x, n_controls=1))
+
+    def visit_CZPowGate(self, state: lowering.State[CirqNode], node: cirq.CZPowGate):
+        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Control(z, n_controls=1))
+
+    def visit_ControlledOperation(
+        self, state: lowering.State[CirqNode], node: cirq.ControlledOperation
+    ):
+        return self.visit_GateOperation(state, node)
+
+    def visit_ControlledGate(
+        self, state: lowering.State[CirqNode], node: cirq.ControlledGate
+    ):
+        op_ = state.lower(node.sub_gate).expect_one()
+        n_controls = node.num_controls()
+        return state.current_frame.push(op.stmts.Control(op_, n_controls=n_controls))
+
+    def visit_XXPowGate(self, state: lowering.State[CirqNode], node: cirq.XXPowGate):
+        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Kron(x, x))
+
+    def visit_YYPowGate(self, state: lowering.State[CirqNode], node: cirq.YYPowGate):
+        y = state.lower(cirq.YPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Kron(y, y))
+
+    def visit_ZZPowGate(self, state: lowering.State[CirqNode], node: cirq.ZZPowGate):
+        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Kron(z, z))
+
+    def visit_CCXPowGate(self, state: lowering.State[CirqNode], node: cirq.CCXPowGate):
+        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Control(x, n_controls=2))
+
+    def visit_CCZPowGate(self, state: lowering.State[CirqNode], node: cirq.CCZPowGate):
+        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Control(z, n_controls=2))
+
+    def visit_BitFlipChannel(
+        self, state: lowering.State[CirqNode], node: cirq.BitFlipChannel
+    ):
+        x = state.current_frame.push(op.stmts.X())
+        p = state.current_frame.push(py.Constant(node.p))
+        return state.current_frame.push(
+            noise.stmts.PauliError(basis=x.result, p=p.result)
+        )
+
+    def visit_AmplitudeDampingChannel(
+        self, state: lowering.State[CirqNode], node: cirq.AmplitudeDampingChannel
+    ):
+        r = state.current_frame.push(op.stmts.Reset())
+        p = state.current_frame.push(py.Constant(node.gamma))
+
+        # TODO: do we need a dedicated noise stmt for this? Using PauliError
+        # with this basis feels like a hack
+        noise_channel = state.current_frame.push(
+            noise.stmts.PauliError(basis=r.result, p=p.result)
+        )
+
+        return noise_channel
+
+    def visit_GeneralizedAmplitudeDampingChannel(
+        self,
+        state: lowering.State[CirqNode],
+        node: cirq.GeneralizedAmplitudeDampingChannel,
+    ):
+        raise NotImplementedError("TODO: needs a new operator statement")
+        # p = state.current_frame.push(py.Constant(node.p))
+        # gamma = state.current_frame.push(py.Constant(node.gamma))
+
+        # p1 =
+
+        # x = state.current_frame.push(op.stmts.X())
+        # noise_channel1 = noise.stmts.PauliError(basis=x.result, p=)

bloqade/squin/groups.py

@@ -3,12 +3,12 @@ from kirin.prelude import structural_no_opt
 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 kernel(self):
     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)
 

bloqade/squin/noise/__init__.py

@@ -4,5 +4,7 @@ from ._wrapper import (
     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,
 )

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)

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)

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)

bloqade/squin/op/stmts.py

@@ -9,7 +9,7 @@ from ._dialect import dialect
 
 @statement
 class Operator(ir.Statement):
-    pass
+    result: ir.ResultValue = info.result(OpType)
 
 
 @statement
@@ -26,7 +26,6 @@ class CompositeOp(Operator):
 class BinaryOp(CompositeOp):
     lhs: ir.SSAValue = info.argument(OpType)
     rhs: ir.SSAValue = info.argument(OpType)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -46,7 +45,6 @@ class Adjoint(CompositeOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
     is_unitary: bool = info.attribute(default=False)
     op: ir.SSAValue = info.argument(OpType)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -55,7 +53,6 @@ class Scale(CompositeOp):
     is_unitary: bool = info.attribute(default=False)
     op: ir.SSAValue = info.argument(OpType)
     factor: ir.SSAValue = info.argument(NumberType)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -64,7 +61,6 @@ class Control(CompositeOp):
     is_unitary: bool = info.attribute(default=False)
     op: ir.SSAValue = info.argument(OpType)
     n_controls: int = info.attribute()
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -72,14 +68,12 @@ class Rot(CompositeOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary()})
     axis: ir.SSAValue = info.argument(OpType)
     angle: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
 class Identity(CompositeOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), HasSites()})
     sites: int = info.attribute()
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement
@@ -87,7 +81,6 @@ class ConstantOp(PrimitiveOp):
     traits = frozenset(
         {ir.Pure(), lowering.FromPythonCall(), ir.ConstantLike(), FixedSites(1)}
     )
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement
@@ -109,7 +102,6 @@ class U3(PrimitiveOp):
     theta: ir.SSAValue = info.argument(types.Float)
     phi: ir.SSAValue = info.argument(types.Float)
     lam: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -124,7 +116,6 @@ class PhaseOp(PrimitiveOp):
 
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), FixedSites(1)})
     theta: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -139,7 +130,6 @@ class ShiftOp(PrimitiveOp):
 
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), FixedSites(1)})
     theta: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -149,7 +139,6 @@ class Reset(PrimitiveOp):
     """
 
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), FixedSites(1)})
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement

bloqade/squin/op/types.py

@@ -22,3 +22,5 @@ class Op:
 
 
 OpType = types.PyClass(Op)
+
+NumOperators = types.TypeVar("NumOperators")

{bloqade_circuit-0.4.0.dist-info → bloqade_circuit-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.4.0
+Version: 0.4.1
 Summary: The software development toolkit for neutral atom arrays.
 Author-email: Roger-luo <rluo@quera.com>, kaihsin <khwu@quera.com>, weinbe58 <pweinberg@quera.com>, johnzl-777 <jlong@quera.com>
 License-File: LICENSE

{bloqade_circuit-0.4.0.dist-info → bloqade_circuit-0.4.1.dist-info}/RECORD

@@ -9,7 +9,7 @@ bloqade/analysis/address/impls.py,sha256=c3FMF2LLkV_fwWaolIRTV7ueXC9qwAR6PUW9kDQ
 bloqade/analysis/address/lattice.py,sha256=dUq999feqPoBYkqEXe1hjHOn4TP_bkvKip8fyWQ-2-8,1755
 bloqade/analysis/fidelity/__init__.py,sha256=iJkhoHvCMU9bKxQqgxIWKQWvpqNFRgNBI5DK8-4RAB8,59
 bloqade/analysis/fidelity/analysis.py,sha256=G6JEYc8eeWJ9mwsbUAIzXuU2nrnTU4te41c04xE71gM,3218
-bloqade/pyqrack/__init__.py,sha256=_EQJ4blQxqS1_Z6OTjmgSquFo1R_9rEjECRg4fH97hQ,794
+bloqade/pyqrack/__init__.py,sha256=lonTS-luJkTVujCCtgdZRC12V7FQdoFcozAI-byXwN0,810
 bloqade/pyqrack/base.py,sha256=9z61PaaAFqCBBwkgsDZSr-qr9IQ5OJ_JUvltmJ7Bgls,4407
 bloqade/pyqrack/device.py,sha256=cOdyT1k0b73QbOsEIu5KqxHW2OAWP86fi_3XGPhaWGA,7134
 bloqade/pyqrack/reg.py,sha256=uTL07CT1R0xUsInLmwU9YuuNdV6lV0lCs1zhdUz1qIs,1660
@@ -27,6 +27,8 @@ bloqade/pyqrack/squin/op.py,sha256=CY2qTYqFPwEceWt-iSiax6i_Bs2i9hQhzGR8ni_PmmA,4
 bloqade/pyqrack/squin/qubit.py,sha256=svQMbsLxv3yjiFMSRc4C7QGllzjtmlSWOsMY1mjTI8Q,2223
 bloqade/pyqrack/squin/runtime.py,sha256=806S23bJKbh2xBuZ0yPAgjs7ZxDHB-zqQtNeA--x9fw,15349
 bloqade/pyqrack/squin/wire.py,sha256=rqlAeU-r_EHOwJMqHrEAxpZ_rKsvUpwGG7MP4BW75Nw,1658
+bloqade/pyqrack/squin/noise/__init__.py,sha256=uXgRQPOrHNRp3k2ff2HD8mheUEaqxZPKEnwV-s4BiV4,31
+bloqade/pyqrack/squin/noise/native.py,sha256=KF4VGzU5Ps92DeLcIDIMsxQQtQ97z_3KUHqBPPkZFaM,2286
 bloqade/qasm2/__init__.py,sha256=W9dR4Qnvigc7e7Ay7puSJHAIuiQk8vWqY-W64SMu5oU,515
 bloqade/qasm2/_qasm_loading.py,sha256=1EFTt1YDkL8fsoSgSuqD1QcKO4EMFIGuBTX9HCnb6S0,4724
 bloqade/qasm2/_wrappers.py,sha256=4x3fldC4sV2K_XZ0FPZOorQKAbs_7pualListXtak4A,11148
@@ -104,8 +106,8 @@ bloqade/qbraid/lowering.py,sha256=84RsPONWeQ_2beyOWBNoJbGFKuaMhsWtRiPivTCZ-Q0,11
 bloqade/qbraid/schema.py,sha256=dTPexUFOiBNBnFv0GEbGh6jpIbMIFHk4hFXmXbeihxA,7854
 bloqade/qbraid/simulation_result.py,sha256=zdCJcAdbQkEDzFFuC2q3gqOFTOLAXHk4wh8RRDB6cgc,3956
 bloqade/qbraid/target.py,sha256=LcFHHyLe74yBmrHI9251xHgLN_nUz35lN8RPNwrT6mI,3149
-bloqade/squin/__init__.py,sha256=8W5Z26BFrR3XRP6yb6IWDA5tZiXMhNstevx5kP3pP5k,169
-bloqade/squin/groups.py,sha256=kFS9WtDC0NDRb4XUfNP8HZdrS961uk39XAAKlufJykA,1289
+bloqade/squin/__init__.py,sha256=0NKHHhSFtlITiAMShgKPkh6tvQbfPzqgBi4I2Wm_LCU,398
+bloqade/squin/groups.py,sha256=RXGJnNZUSXF_f5ljjhZ9At8UhaijayoxFoWvxEsUOWc,1310
 bloqade/squin/lowering.py,sha256=w-GyOKYZHHKCGA2slcgWNS97Q_znQU65PeYxEIkvChM,816
 bloqade/squin/qubit.py,sha256=psIZPtbQHsiToCXcT4wuuaZzEifaPVvk_SWrsgPbNwg,5067
 bloqade/squin/wire.py,sha256=kRmpC7P6qIOZOsuFJHJQyeNsEWcU2Of4ZgcdsMESnKA,3746
@@ -115,19 +117,22 @@ bloqade/squin/analysis/nsites/__init__.py,sha256=RlQg7ivczXCXG5lMeL3ipYKj2oJKC4T
 bloqade/squin/analysis/nsites/analysis.py,sha256=rIe1RU1MZRItcE2aB8DYahLrv73HfD3IHCX3E_EGQ1c,1773
 bloqade/squin/analysis/nsites/impls.py,sha256=OaKuAoZ0EAorStYDZxzgc6Dk42kuj19MLkqHWG1MEQM,2592
 bloqade/squin/analysis/nsites/lattice.py,sha256=ruh0808SHtj3ecuT-C3AZTsLY2j3DRhtezGiTZvcuVs,942
-bloqade/squin/noise/__init__.py,sha256=K7wHkzUxWbLF-XQPCAlXY7izz-bS7LJXt2G3GjlkEQc,218
+bloqade/squin/cirq/__init__.py,sha256=aGRwIWqbbKaW-1Pm697uo-wQsM9ZBf0h5Qlo3Y6HJ-Y,2453
+bloqade/squin/cirq/lowering.py,sha256=Crq4oq0O_ucZCjSXV6DNEMvZCo5HCl9Fy0mpYfYOdok,11689
+bloqade/squin/noise/__init__.py,sha256=HQl3FE0SZAGEX3qdveapCaMX391lgLvWeWnoE6Z2pYw,332
 bloqade/squin/noise/_dialect.py,sha256=2IR98J-lXm5Y3srP9g-FD4JC-qTq2seureM6mKKq1xg,63
-bloqade/squin/noise/_wrapper.py,sha256=fHytpR6kINTPnaTfmMa_GRn4geBQkbEGVFRE60V0Zko,474
-bloqade/squin/noise/stmts.py,sha256=0XvBdZFkjMABUkaj5Fry8oQ8QmWzoZggX1aecvWBnVk,1607
+bloqade/squin/noise/_wrapper.py,sha256=0jD5va_go9jEW5rC6bZSWU30kjCha2-axFogPON3-V0,580
+bloqade/squin/noise/rewrite.py,sha256=SxIHgMDqYJXepiZDyukHWpe5yaFDSTG-yJ4JONNVr0o,3917
+bloqade/squin/noise/stmts.py,sha256=rktxkIdjdPUYek0MYh9uh83otkl-7UoADCoWHWf57J8,1678
 bloqade/squin/op/__init__.py,sha256=5OBgT4E44Cy0DNF3yRbXGXkiB8VAJtr48x8hDQEquH4,741
 bloqade/squin/op/_dialect.py,sha256=66G1IYqmsqUEaCTyUqn2shSHmGYduiTU8GfDXcoMvw4,55
 bloqade/squin/op/_wrapper.py,sha256=5pqbKGeNWoYQIPa1xkKBr8z5waxwmAm3AV4efGSRT_s,1714
 bloqade/squin/op/number.py,sha256=yujWUqLrOAr8i8OBDsiS5M882wV7t08u345NgNA6TUc,95
 bloqade/squin/op/rewrite.py,sha256=Itxz_hTAPNLyLYeLS0PCVk143J1Z558UR7N9-urbnoU,1327
 bloqade/squin/op/stdlib.py,sha256=4UFK3wKImpums2v5a9OFKuVvz2TLYbYwidg3JYYEi2o,1073
-bloqade/squin/op/stmts.py,sha256=0QMB3YC_YLr9CfNQiNapSRRT7V5uU8_DVLy-u8Oq9qw,5761
+bloqade/squin/op/stmts.py,sha256=uBQeCWc1JjMTIVcssWkLWi6MyTLtQMoqTWoKgLsYgWQ,5262
 bloqade/squin/op/traits.py,sha256=jjsnzWtPtmQK7K3H_D2fvc8XiW1Y3EMBcgeyPax2sjc,1065
-bloqade/squin/op/types.py,sha256=2am1KC9FHBl__5q_87hTXrSdx1KMkcEUBeZLSrBZQEw,617
+bloqade/squin/op/types.py,sha256=ozUT0Bv9NuUxPjB2vAeqJ9cpdvUaBfP9trB5mybYxgc,663
 bloqade/squin/passes/__init__.py,sha256=Bhog-wZBtToNJXfhlYa6S7tE6OoppyRibjMl5JBfY58,45
 bloqade/squin/passes/stim.py,sha256=VWv3hhfizCWz5sIwwdFt3flWHLzG428evLGIcX8E36Y,1992
 bloqade/squin/rewrite/__init__.py,sha256=0-9m1cbvFRgjZpQ700NEjW1uKvwZPPbrmUwylhgOjUw,457
@@ -186,7 +191,7 @@ bloqade/visual/animation/runtime/atoms.py,sha256=EmjxhujLiHHPS_HtH_B-7TiqeHgvW5u
 bloqade/visual/animation/runtime/ppoly.py,sha256=JB9IP53N1w6adBJEue6J5Nmj818Id9JvrlgrmiQTU1I,1385
 bloqade/visual/animation/runtime/qpustate.py,sha256=rlmxQeJSvaohXrTpXQL5y-NJcpvfW33xPaYM1slv7cc,4270
 bloqade/visual/animation/runtime/utils.py,sha256=ju9IzOWX-vKwfpqUjlUKu3Ssr_UFPFFq-tzH_Nqyo_c,1212
-bloqade_circuit-0.4.0.dist-info/METADATA,sha256=3CjbxNGvQkBKgscX6HpdRqxkGyiAmDLgkgsAYovRvZc,3683
-bloqade_circuit-0.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bloqade_circuit-0.4.0.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
-bloqade_circuit-0.4.0.dist-info/RECORD,,
+bloqade_circuit-0.4.1.dist-info/METADATA,sha256=53teyQGCbiQ0-uxPqwBrNAKdeLcpe7h3Aj7fgVq26HU,3683
+bloqade_circuit-0.4.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bloqade_circuit-0.4.1.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
+bloqade_circuit-0.4.1.dist-info/RECORD,,