bloqade-circuit 0.7.13__py3-none-any.whl → 0.8.0__py3-none-any.whl

bloqade/squin/stdlib/simple/noise.py

@@ -0,0 +1,126 @@
+from typing import Any, Literal, TypeVar
+
+from kirin.dialects import ilist
+
+from bloqade.types import Qubit
+
+from .. import broadcast
+from ...groups import kernel
+
+
+@kernel
+def depolarize(p: float, qubit: Qubit) -> None:
+    """
+    Apply a depolarizing noise channel to a qubit with probability `p`.
+
+    This will randomly select one of the Pauli operators X, Y, Z
+    with a probability `p / 3` and apply it to the qubit. No operator is applied
+    with a probability of `1 - p`.
+
+    Args:
+        p (float): The probability with which a Pauli operator is applied.
+        qubit (Qubit): The qubit to which the noise channel is applied.
+    """
+    broadcast.depolarize(p, ilist.IList([qubit]))
+
+
+N = TypeVar("N", bound=int)
+
+
+@kernel
+def depolarize2(p: float, control: Qubit, target: Qubit) -> None:
+    """
+    Symmetric two-qubit depolarization channel applied to a pair of qubits.
+
+    This will randomly select one of the pauli products
+
+    `{IX, IY, IZ, XI, XX, XY, XZ, YI, YX, YY, YZ, ZI, ZX, ZY, ZZ}`
+
+    each with a probability `p / 15`. No noise is applied with a probability of `1 - p`.
+
+    Args:
+        p (float): The probability with which a Pauli product is applied.
+        control (Qubit): The control qubit.
+        target (Qubit): The target qubit.
+    """
+    broadcast.depolarize2(p, ilist.IList([control]), ilist.IList([target]))
+
+
+@kernel
+def single_qubit_pauli_channel(px: float, py: float, pz: float, qubit: Qubit) -> None:
+    """
+    Apply a Pauli error channel with weighted `px, py, pz`. No error is applied with a probability
+    `1 - (px + py + pz)`.
+
+    This randomly selects one of the three Pauli operators X, Y, Z, weighted with the given probabilities in that order.
+
+    Args:
+        probabilities (IList[float, Literal[3]]): A list of 3 probabilities corresponding to the probabilities `(p_x, p_y, p_z)` in that order.
+        qubit (Qubit): The qubit to which the noise channel is applied.
+    """
+    broadcast.single_qubit_pauli_channel(px, py, pz, ilist.IList([qubit]))
+
+
+@kernel
+def two_qubit_pauli_channel(
+    probabilities: ilist.IList[float, Literal[15]], control: Qubit, target: Qubit
+) -> None:
+    """
+    Apply a Pauli product error with weighted `probabilities` to the pair of qubits.
+
+    No error is applied with the probability `1 - sum(probabilities)`.
+
+    This will randomly select one of the pauli products
+
+    `{IX, IY, IZ, XI, XX, XY, XZ, YI, YX, YY, YZ, ZI, ZX, ZY, ZZ}`
+
+    weighted with the corresponding list of probabilities.
+
+    **NOTE**: The order of the given probabilities must match the order of the list of Pauli products above!
+    """
+    broadcast.two_qubit_pauli_channel(
+        probabilities, ilist.IList([control]), ilist.IList([target])
+    )
+
+
+@kernel
+def qubit_loss(p: float, qubit: Qubit) -> None:
+    """
+    Apply a qubit loss channel to the given qubit.
+
+    The qubit is lost with a probability `p`.
+
+    Args:
+        p (float): Probability of the atom being lost.
+        qubit (Qubit): The qubit to which the noise channel is applied.
+    """
+    broadcast.qubit_loss(p, ilist.IList([qubit]))
+
+
+@kernel
+def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None:
+    """
+    Apply a correlated qubit loss channel to the given qubits.
+
+    All qubits are lost together with a probability `p`.
+
+    Args:
+        p (float): Probability of the qubits being lost.
+        qubits (IList[Qubit, Any]): The list of qubits to which the correlated noise channel is applied.
+    """
+    broadcast.correlated_qubit_loss(p, ilist.IList([qubits]))
+
+
+# NOTE: actual stdlib that doesn't wrap statements starts here
+
+
+@kernel
+def bit_flip(p: float, qubit: Qubit) -> None:
+    """
+    Apply a bit flip error channel to the qubit with probability `p`.
+
+    Args:
+        p (float): Probability of a bit flip error being applied.
+        qubit (Qubit): The qubit to which the noise channel is applied.
+    """
+    single_qubit_pauli_channel(p, 0, 0, qubit)

bloqade/stim/__init__.py

@@ -39,4 +39,5 @@ from ._wrappers import (
     pauli_channel1 as pauli_channel1,
     pauli_channel2 as pauli_channel2,
     observable_include as observable_include,
+    correlated_qubit_loss as correlated_qubit_loss,
 )

bloqade/stim/_wrappers.py

@@ -194,3 +194,9 @@ def z_error(p: float, targets: tuple[int, ...]) -> None: ...
 
 @wraps(noise.QubitLoss)
 def qubit_loss(probs: tuple[float, ...], targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.CorrelatedQubitLoss)
+def correlated_qubit_loss(
+    probs: tuple[float, ...], targets: tuple[int, ...]
+) -> None: ...

bloqade/stim/dialects/noise/emit.py

@@ -81,6 +81,7 @@ class EmitStimNoiseMethods(MethodTable):
         return ()
 
     @impl(stmts.TrivialCorrelatedError)
+    @impl(stmts.CorrelatedQubitLoss)
     def non_stim_corr_error(
         self,
         emit: EmitStimMain,
@@ -92,7 +93,11 @@ class EmitStimNoiseMethods(MethodTable):
         prob: tuple[str, ...] = frame.get_values(stmt.probs)
         prob_str: str = ", ".join(prob)
 
-        res = f"I_ERROR[{stmt.name}:{stmt.nonce}]({prob_str}) " + " ".join(targets)
+        res = (
+            f"I_ERROR[{stmt.name}:{emit.correlated_error_count}]({prob_str}) "
+            + " ".join(targets)
+        )
+        emit.correlated_error_count += 1
         emit.writeln(frame, res)
 
         return ()

bloqade/stim/dialects/noise/stmts.py

@@ -89,9 +89,6 @@ class NonStimError(ir.Statement):
 class NonStimCorrelatedError(ir.Statement):
     name = "NonStimCorrelatedError"
     traits = frozenset({lowering.FromPythonCall()})
-    nonce: int = (
-        info.attribute()
-    )  # Must be a unique value, otherwise stim might merge two correlated errors with equal probabilities
     probs: tuple[ir.SSAValue, ...] = info.argument(types.Float)
     targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
 
@@ -109,3 +106,8 @@ class TrivialError(NonStimError):
 @statement(dialect=dialect)
 class QubitLoss(NonStimError):
     name = "loss"
+
+
+@statement(dialect=dialect)
+class CorrelatedQubitLoss(NonStimCorrelatedError):
+    name = "correlated_loss"

bloqade/stim/emit/stim_str.py

@@ -20,11 +20,13 @@ class EmitStimMain(EmitStr):
     keys = ["emit.stim"]
     dialects: ir.DialectGroup = field(default_factory=_default_dialect_group)
     file: StringIO = field(default_factory=StringIO)
+    correlation_identifier_offset: int = 0
 
     def initialize(self):
         super().initialize()
         self.file.truncate(0)
         self.file.seek(0)
+        self.correlated_error_count = self.correlation_identifier_offset
         return self
 
     def eval_stmt_fallback(

bloqade/stim/parse/lowering.py

@@ -627,10 +627,13 @@ class Stim(lowering.LoweringABC[Node]):
         # Parse tag
         tag_parts = node.tag.split(";", maxsplit=1)[0].split(":", maxsplit=1)
         nonstim_name = tag_parts[0]
-        nonce = 0
         if len(tag_parts) == 2:
+            # This should be a correlated error of the form, e.g.,
+            # I_ERROR[correlated_loss:<identifier>](0.01) 0 1 2
+            # The identifier is a unique number that prevents stim from merging
+            # correlated errors. We discard the identifier, but verify it is an integer.
             try:
-                nonce = int(tag_parts[1])
+                _ = int(tag_parts[1])
             except ValueError:
                 # String was not an integer
                 if self.error_unknown_nonstim:
@@ -643,22 +646,14 @@ class Stim(lowering.LoweringABC[Node]):
                 f"Unknown non-stim statement name: {nonstim_name!r} ({node!r})"
             )
         statement_cls = self.nonstim_noise_ops.get(nonstim_name)
+        stmt = None
         if statement_cls is not None:
-            if issubclass(statement_cls, noise.NonStimCorrelatedError):
-                stmt = statement_cls(
-                    nonce=nonce,
-                    probs=self._get_float_args_ssa(state, node.gate_args_copy()),
-                    targets=self._get_multiple_qubit_or_rec_ssa(
-                        state, node, node.targets_copy()
-                    ),
-                )
-            else:
-                stmt = statement_cls(
-                    probs=self._get_float_args_ssa(state, node.gate_args_copy()),
-                    targets=self._get_multiple_qubit_or_rec_ssa(
-                        state, node, node.targets_copy()
-                    ),
-                )
+            stmt = statement_cls(
+                probs=self._get_float_args_ssa(state, node.gate_args_copy()),
+                targets=self._get_multiple_qubit_or_rec_ssa(
+                    state, node, node.targets_copy()
+                ),
+            )
         return stmt
 
     def visit_CircuitInstruction(

bloqade/stim/passes/__init__.py

@@ -1,4 +1,3 @@
 from .squin_to_stim import (
     SquinToStimPass as SquinToStimPass,
-    StimSimplifyIfs as StimSimplifyIfs,
 )

bloqade/stim/passes/flatten.py

@@ -0,0 +1,26 @@
+# Taken from Phillip Weinberg's bloqade-shuttle implementation
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite.abc import RewriteResult
+
+from bloqade.qasm2.passes.fold import AggressiveUnroll
+from bloqade.stim.passes.simplify_ifs import StimSimplifyIfs
+
+
+@dataclass
+class Flatten(Pass):
+
+    unroll: AggressiveUnroll = field(init=False)
+    simplify_if: StimSimplifyIfs = field(init=False)
+
+    def __post_init__(self):
+        self.unroll = AggressiveUnroll(self.dialects, no_raise=self.no_raise)
+        self.simplify_if = StimSimplifyIfs(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        rewrite_result = RewriteResult()
+        rewrite_result = self.simplify_if(mt).join(rewrite_result)
+        rewrite_result = self.unroll(mt).join(rewrite_result)
+        return rewrite_result

bloqade/stim/passes/simplify_ifs.py

@@ -7,8 +7,10 @@ from kirin.rewrite import (
     Chain,
     Fixpoint,
     ConstantFold,
+    DeadCodeElimination,
     CommonSubexpressionElimination,
 )
+from kirin.dialects.scf.trim import UnusedYield
 from kirin.dialects.ilist.passes import ConstList2IList
 
 from ..rewrite.ifs_to_stim import StimLiftThenBody, StimSplitIfStmts
@@ -20,7 +22,10 @@ class StimSimplifyIfs(Pass):
     def unsafe_run(self, mt: ir.Method):
 
         result = Chain(
-            Fixpoint(Walk(StimLiftThenBody())),
+            Walk(UnusedYield()),
+            Walk(StimLiftThenBody()),
+            # remove yields (if possible), then lift out as much stuff as possible
+            Walk(DeadCodeElimination()),
             Walk(StimSplitIfStmts()),
         ).rewrite(mt.code)
 

bloqade/stim/passes/squin_to_stim.py

@@ -1,29 +1,21 @@
 from dataclasses import dataclass
 
-from kirin.passes import Fold, TypeInfer
 from kirin.rewrite import (
     Walk,
     Chain,
     Fixpoint,
-    CFGCompactify,
     DeadCodeElimination,
     CommonSubexpressionElimination,
 )
-from kirin.dialects import ilist
 from kirin.ir.method import Method
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
-from kirin.passes.inline import InlinePass
-from kirin.rewrite.alias import InlineAlias
-from kirin.passes.aggressive import UnrollScf
 
 from bloqade.stim.rewrite import (
-    SquinWireToStim,
     PyConstantToStim,
     SquinNoiseToStim,
     SquinQubitToStim,
     SquinMeasureToStim,
-    SquinWireIdentityElimination,
 )
 from bloqade.squin.rewrite import (
     SquinU3ToClifford,
@@ -33,9 +25,8 @@ from bloqade.squin.rewrite import (
 from bloqade.rewrite.passes import CanonicalizeIList
 from bloqade.analysis.address import AddressAnalysis
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
-from bloqade.squin.rewrite.desugar import ApplyDesugarRule, MeasureDesugarRule
+from bloqade.stim.passes.flatten import Flatten
 
-from .simplify_ifs import StimSimplifyIfs
 from ..rewrite.ifs_to_stim import IfToStim
 
 
@@ -45,63 +36,8 @@ class SquinToStimPass(Pass):
     def unsafe_run(self, mt: Method) -> RewriteResult:
 
         # inline aggressively:
-        rewrite_result = InlinePass(
-            dialects=mt.dialects, no_raise=self.no_raise
-        ).unsafe_run(mt)
-
-        rewrite_result = Walk(ilist.rewrite.HintLen()).rewrite(mt.code)
-        rewrite_result = Fold(self.dialects).unsafe_run(mt).join(rewrite_result)
-
-        rewrite_result = (
-            UnrollScf(dialects=mt.dialects, no_raise=self.no_raise)
-            .fixpoint(mt)
-            .join(rewrite_result)
-        )
-
-        rewrite_result = (
-            Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(rewrite_result)
-        )
-
-        rewrite_result = Walk(InlineAlias()).rewrite(mt.code).join(rewrite_result)
-
-        rewrite_result = (
-            StimSimplifyIfs(mt.dialects, no_raise=self.no_raise)
-            .unsafe_run(mt)
-            .join(rewrite_result)
-        )
-
-        rewrite_result = (
-            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
-            .rewrite(mt.code)
-            .join(rewrite_result)
-        )
-        rewrite_result = Fold(mt.dialects, no_raise=self.no_raise)(mt)
-
-        rewrite_result = (
-            UnrollScf(mt.dialects, no_raise=self.no_raise)
-            .fixpoint(mt)
-            .join(rewrite_result)
-        )
-
-        rewrite_result = (
-            CanonicalizeIList(dialects=mt.dialects, no_raise=self.no_raise)
-            .unsafe_run(mt)
-            .join(rewrite_result)
-        )
-
-        rewrite_result = TypeInfer(
-            dialects=mt.dialects, no_raise=self.no_raise
-        ).unsafe_run(mt)
-
-        rewrite_result = (
-            Walk(
-                Chain(
-                    ApplyDesugarRule(),
-                    MeasureDesugarRule(),
-                )
-            )
-            .rewrite(mt.code)
-            .join(rewrite_result)
+        rewrite_result = Flatten(dialects=mt.dialects, no_raise=self.no_raise).fixpoint(
+            mt
         )
 
         # after this the program should be in a state where it is analyzable
@@ -145,8 +81,6 @@ class SquinToStimPass(Pass):
                 Chain(
                     SquinQubitToStim(),
                     SquinMeasureToStim(),
-                    SquinWireToStim(),
-                    SquinWireIdentityElimination(),
                 )
             )
             .rewrite(mt.code)
@@ -163,7 +97,7 @@ class SquinToStimPass(Pass):
         rewrite_result = Walk(PyConstantToStim()).rewrite(mt.code).join(rewrite_result)
 
         # clear up leftover stmts
-        # - remove any squin.qubit.new that's left around
+        # - remove any squin.qalloc that's left around
         rewrite_result = (
             Fixpoint(
                 Walk(

bloqade/stim/rewrite/__init__.py

@@ -1,9 +1,5 @@
 from .ifs_to_stim import IfToStim as IfToStim
 from .squin_noise import SquinNoiseToStim as SquinNoiseToStim
-from .wire_to_stim import SquinWireToStim as SquinWireToStim
 from .qubit_to_stim import SquinQubitToStim as SquinQubitToStim
 from .squin_measure import SquinMeasureToStim as SquinMeasureToStim
 from .py_constant_to_stim import PyConstantToStim as PyConstantToStim
-from .wire_identity_elimination import (
-    SquinWireIdentityElimination as SquinWireIdentityElimination,
-)

bloqade/stim/rewrite/ifs_to_stim.py

@@ -4,13 +4,13 @@ from kirin import ir
 from kirin.dialects import py, scf, func
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import op, qubit
+from bloqade.squin import gate
 from bloqade.rewrite.rules import LiftThenBody, SplitIfStmts
 from bloqade.squin.rewrite import AddressAttribute
 from bloqade.stim.rewrite.util import (
-    SQUIN_STIM_CONTROL_GATE_MAPPING,
     insert_qubit_idx_from_address,
 )
+from bloqade.stim.dialects.gate import CX as stim_CX, CY as stim_CY, CZ as stim_CZ
 from bloqade.analysis.measure_id import MeasureIDFrame
 from bloqade.stim.dialects.auxiliary import GetRecord
 from bloqade.analysis.measure_id.lattice import (
@@ -58,8 +58,7 @@ class IfElseSimplification:
         """Check if the IfElse statement has an else body."""
         if stmt.else_body.blocks and not (
             len(stmt.else_body.blocks[0].stmts) == 1
-            and isinstance(else_term := stmt.else_body.blocks[0].last_stmt, scf.Yield)
-            and not else_term.values  # empty yield
+            and isinstance(stmt.else_body.blocks[0].last_stmt, scf.Yield)
         ):
             return True
 
@@ -67,12 +66,13 @@ class IfElseSimplification:
 
 
 DontLiftType = (
-    qubit.Apply,
-    qubit.Broadcast,
-    scf.Yield,
+    gate.stmts.SingleQubitGate,
+    gate.stmts.RotationGate,
+    gate.stmts.ControlledGate,
     func.Return,
     func.Invoke,
     scf.IfElse,
+    scf.Yield,
 )
 
 
@@ -99,16 +99,16 @@ class StimSplitIfStmts(IfElseSimplification, SplitIfStmts):
     Given an IfElse with multiple valid statements in the then-body:
 
     if measure_result:
-        squin.qubit.apply(op.X, q0)
-        squin.qubit.apply(op.Y, q1)
+        squin.x(q0)
+        squin.y(q1)
 
     this should be rewritten to:
 
     if measure_result:
-        squin.qubit.apply(op.X, q0)
+        squin.x(q0)
 
     if measure_result:
-        squin.qubit.apply(op.Y, q1)
+        squin.y(q1)
     """
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
@@ -139,24 +139,23 @@ class IfToStim(IfElseSimplification, RewriteRule):
 
     def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
 
+        # Check the condition is a singular MeasurementIdBool
         if not isinstance(self.measure_frame.entries[stmt.cond], MeasureIdBool):
             return RewriteResult()
 
-        # check that there is only qubit.Apply in the then-body,
-        # if there's more than that, we can't do a valid rewrite.
-        # Can reuse logic from SplitIf
+        # Reusing code from SplitIf,
+        # there should only be one statement in the body and it should be a pauli X, Y, or Z
         *stmts, _ = stmt.then_body.stmts()
-        if len(stmts) != 1 or not isinstance(stmts[0], (qubit.Apply, qubit.Broadcast)):
+        if len(stmts) != 1:
             return RewriteResult()
 
-        apply_or_broadcast = stmts[0]
-        # Check that the gate being applied/broadcasted can be converted to a stim
-        # controlled gate.
-        ctrl_op_target_gate = apply_or_broadcast.operator.owner
-        assert isinstance(ctrl_op_target_gate, op.stmts.Operator)
-
-        stim_gate = SQUIN_STIM_CONTROL_GATE_MAPPING.get(type(ctrl_op_target_gate))
-        if stim_gate is None:
+        if isinstance(stmts[0], gate.stmts.X):
+            stim_gate = stim_CX
+        elif isinstance(stmts[0], gate.stmts.Y):
+            stim_gate = stim_CY
+        elif isinstance(stmts[0], gate.stmts.Z):
+            stim_gate = stim_CZ
+        else:
             return RewriteResult()
 
         # get necessary measurement ID type from analysis
@@ -169,12 +168,7 @@ class IfToStim(IfElseSimplification, RewriteRule):
         )
         get_record_stmt = GetRecord(id=measure_id_idx_stmt.result)  # noqa: F841
 
-        # get address attribute and generate qubit idx statements
-        if len(apply_or_broadcast.qubits) != 1:
-            # NOTE: this is actually invalid since we are dealing with single-qubit operators here
-            return RewriteResult()
-
-        address_attr = apply_or_broadcast.qubits[0].hints.get("address")
+        address_attr = stmts[0].qubits.hints.get("address")
 
         if address_attr is None:
             return RewriteResult()