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

bloqade/squin/rewrite/qubit_to_stim.py

@@ -1,7 +1,6 @@
 from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade import stim
 from bloqade.squin import op, qubit
 from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
 from bloqade.squin.rewrite.stim_rewrite_util import (
@@ -18,8 +17,6 @@ class SquinQubitToStim(RewriteRule):
         match node:
             case qubit.Apply() | qubit.Broadcast():
                 return self.rewrite_Apply_and_Broadcast(node)
-            case qubit.Reset():
-                return self.rewrite_Reset(node)
             case _:
                 return RewriteResult()
 
@@ -60,25 +57,5 @@ class SquinQubitToStim(RewriteRule):
 
         return RewriteResult(has_done_something=True)
 
-    def rewrite_Reset(self, reset_stmt: qubit.Reset) -> RewriteResult:
-        qubit_ilist_ssa = reset_stmt.qubits
-        # qubits are in an ilist which makes up an AddressTuple
-        address_attr = qubit_ilist_ssa.hints.get("address")
-        if address_attr is None:
-            return RewriteResult()
-
-        assert isinstance(address_attr, AddressAttribute)
-        qubit_idx_ssas = insert_qubit_idx_from_address(
-            address=address_attr, stmt_to_insert_before=reset_stmt
-        )
-
-        if qubit_idx_ssas is None:
-            return RewriteResult()
-
-        stim_rz_stmt = stim.collapse.stmts.RZ(targets=qubit_idx_ssas)
-        reset_stmt.replace_by(stim_rz_stmt)
-
-        return RewriteResult(has_done_something=True)
-
 
 # put rewrites for measure statements in separate rule, then just have to dispatch

bloqade/squin/rewrite/squin_measure.py

@@ -3,8 +3,8 @@ from kirin import ir
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade import stim
 from bloqade.squin import wire, qubit
+from bloqade.stim.dialects import collapse
 from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
 from bloqade.squin.rewrite.stim_rewrite_util import (
     is_measure_result_used,
@@ -22,8 +22,6 @@ class SquinMeasureToStim(RewriteRule):
         match node:
             case qubit.MeasureQubit() | qubit.MeasureQubitList() | wire.Measure():
                 return self.rewrite_Measure(node)
-            case qubit.MeasureAndReset() | wire.MeasureAndReset():
-                return self.rewrite_MeasureAndReset(node)
             case _:
                 return RewriteResult()
 
@@ -38,7 +36,7 @@ class SquinMeasureToStim(RewriteRule):
             return RewriteResult()
 
         prob_noise_stmt = py.constant.Constant(0.0)
-        stim_measure_stmt = stim.collapse.MZ(
+        stim_measure_stmt = collapse.MZ(
             p=prob_noise_stmt.result,
             targets=qubit_idx_ssas,
         )
@@ -47,29 +45,6 @@ class SquinMeasureToStim(RewriteRule):
 
         return RewriteResult(has_done_something=True)
 
-    def rewrite_MeasureAndReset(
-        self, meas_and_reset_stmt: qubit.MeasureAndReset | wire.MeasureAndReset
-    ) -> RewriteResult:
-        if not is_measure_result_used(meas_and_reset_stmt):
-            return RewriteResult()
-
-        qubit_idx_ssas = self.get_qubit_idx_ssas(meas_and_reset_stmt)
-
-        if qubit_idx_ssas is None:
-            return RewriteResult()
-
-        error_p_stmt = py.Constant(0.0)
-        stim_mz_stmt = stim.collapse.MZ(targets=qubit_idx_ssas, p=error_p_stmt.result)
-        stim_rz_stmt = stim.collapse.RZ(
-            targets=qubit_idx_ssas,
-        )
-
-        error_p_stmt.insert_before(meas_and_reset_stmt)
-        stim_mz_stmt.insert_before(meas_and_reset_stmt)
-        meas_and_reset_stmt.replace_by(stim_rz_stmt)
-
-        return RewriteResult(has_done_something=True)
-
     def get_qubit_idx_ssas(
         self, measure_stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure
     ) -> tuple[ir.SSAValue, ...] | None:

bloqade/squin/rewrite/stim_rewrite_util.py

@@ -3,7 +3,7 @@ from kirin.dialects import py
 from kirin.rewrite.abc import RewriteResult
 
 from bloqade.squin import op, wire, qubit
-from bloqade.stim.dialects import gate
+from bloqade.stim.dialects import gate, collapse
 from bloqade.analysis.address import AddressWire, AddressQubit, AddressTuple
 from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
 
@@ -14,6 +14,7 @@ SQUIN_STIM_GATE_MAPPING = {
     op.stmts.H: gate.H,
     op.stmts.S: gate.S,
     op.stmts.Identity: gate.Identity,
+    op.stmts.Reset: collapse.RZ,
 }
 
 
@@ -144,13 +145,7 @@ def rewrite_Control(
 
 
 def is_measure_result_used(
-    stmt: (
-        qubit.MeasureAndReset
-        | qubit.MeasureQubit
-        | qubit.MeasureQubitList
-        | wire.MeasureAndReset
-        | wire.Measure
-    ),
+    stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure,
 ) -> bool:
     """
     Check if the result of a measure statement is used in the program.

bloqade/squin/rewrite/wire_to_stim.py

@@ -1,13 +1,10 @@
 from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade import stim
 from bloqade.squin import op, wire
-from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
 from bloqade.squin.rewrite.stim_rewrite_util import (
     SQUIN_STIM_GATE_MAPPING,
     rewrite_Control,
-    insert_qubit_idx_from_address,
     insert_qubit_idx_from_wire_ssa,
 )
 
@@ -18,8 +15,6 @@ class SquinWireToStim(RewriteRule):
         match node:
             case wire.Apply() | wire.Broadcast():
                 return self.rewrite_Apply_and_Broadcast(node)
-            case wire.Reset():
-                return self.rewrite_Reset(node)
             case _:
                 return RewriteResult()
 
@@ -55,19 +50,3 @@ class SquinWireToStim(RewriteRule):
         stmt.replace_by(stim_1q_stmt)
 
         return RewriteResult(has_done_something=True)
-
-    def rewrite_Reset(self, reset_stmt: wire.Reset) -> RewriteResult:
-        address_attr = reset_stmt.wire.hints.get("address")
-        if address_attr is None:
-            return RewriteResult()
-        assert isinstance(address_attr, AddressAttribute)
-        qubit_idx_ssas = insert_qubit_idx_from_address(
-            address=address_attr, stmt_to_insert_before=reset_stmt
-        )
-        if qubit_idx_ssas is None:
-            return RewriteResult()
-
-        stim_rz_stmt = stim.collapse.stmts.RZ(targets=qubit_idx_ssas)
-        reset_stmt.replace_by(stim_rz_stmt)
-
-        return RewriteResult(has_done_something=True)

bloqade/squin/wire.py

@@ -95,23 +95,18 @@ class Broadcast(ir.Statement):
 class Measure(ir.Statement):
     traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
     wire: ir.SSAValue = info.argument(WireType)
+    qubit: ir.SSAValue = info.argument(QubitType)
     result: ir.ResultValue = info.result(types.Int)
 
 
 @statement(dialect=dialect)
-class MeasureAndReset(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
-    wire: ir.SSAValue = info.argument(WireType)
+class NonDestructiveMeasure(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    input_wire: ir.SSAValue = info.argument(WireType)
     result: ir.ResultValue = info.result(types.Int)
     out_wire: ir.ResultValue = info.result(WireType)
 
 
-@statement(dialect=dialect)
-class Reset(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
-    wire: ir.SSAValue = info.argument(WireType)
-
-
 @wraps(Unwrap)
 def unwrap(qubit: Qubit) -> Wire: ...
 

bloqade/stim/__init__.py

@@ -1,4 +1,4 @@
-from . import emit as emit, dialects as dialects
+from . import emit as emit, parse as parse, dialects as dialects
 from .groups import main as main
 from ._wrappers import (
     h as h,
@@ -34,6 +34,7 @@ from ._wrappers import (
     depolarize1 as depolarize1,
     depolarize2 as depolarize2,
     pauli_string as pauli_string,
+    qubit_coords as qubit_coords,
     pauli_channel1 as pauli_channel1,
     pauli_channel2 as pauli_channel2,
     observable_include as observable_include,

bloqade/stim/_wrappers.py

@@ -99,6 +99,10 @@ def pauli_string(
 ) -> auxiliary.PauliString: ...
 
 
+@wraps(auxiliary.QubitCoordinates)
+def qubit_coords(coord: tuple[Union[int, float], ...], target: int) -> None: ...
+
+
 # dialect:: collapse
 @wraps(collapse.MZ)
 def mz(p: float, targets: tuple[int, ...]) -> None: ...

bloqade/stim/dialects/auxiliary/__init__.py

@@ -10,6 +10,7 @@ from .stmts import (
     GetRecord as GetRecord,
     ConstFloat as ConstFloat,
     NewPauliString as NewPauliString,
+    QubitCoordinates as QubitCoordinates,
     ObservableInclude as ObservableInclude,
 )
 from .types import (

bloqade/stim/dialects/auxiliary/emit.py

@@ -69,8 +69,10 @@ class EmitStimAuxMethods(MethodTable):
 
         coord_str: str = ", ".join(coords)
         target_str: str = " ".join(targets)
-        emit.writeln(frame, f"DETECTOR({coord_str}) {target_str}")
-
+        if len(coords):
+            emit.writeln(frame, f"DETECTOR({coord_str}) {target_str}")
+        else:
+            emit.writeln(frame, f"DETECTOR {target_str}")
         return ()
 
     @impl(stmts.ObservableInclude)
@@ -100,3 +102,16 @@ class EmitStimAuxMethods(MethodTable):
         )
 
         return (out,)
+
+    @impl(stmts.QubitCoordinates)
+    def qubit_coordinates(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.QubitCoordinates
+    ):
+
+        coords: tuple[str, ...] = frame.get_values(stmt.coord)
+        target: str = frame.get(stmt.target)
+
+        coord_str: str = ", ".join(coords)
+        emit.writeln(frame, f"QUBIT_COORDS({coord_str}) {target}")
+
+        return ()

bloqade/stim/dialects/auxiliary/stmts/__init__.py

@@ -10,5 +10,6 @@ from .annotate import (
     Detector as Detector,
     GetRecord as GetRecord,
     NewPauliString as NewPauliString,
+    QubitCoordinates as QubitCoordinates,
     ObservableInclude as ObservableInclude,
 )

bloqade/stim/dialects/auxiliary/stmts/annotate.py

@@ -45,3 +45,11 @@ class NewPauliString(ir.Statement):
     flipped: tuple[ir.SSAValue, ...] = info.argument(types.Bool)
     targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
     result: ir.ResultValue = info.result(type=PauliStringType)
+
+
+@statement(dialect=dialect)
+class QubitCoordinates(ir.Statement):
+    name = "qubit_coordinates"
+    traits = frozenset({lowering.FromPythonCall()})
+    coord: tuple[ir.SSAValue, ...] = info.argument(PyNum)
+    target: ir.SSAValue = info.argument(types.Int)

bloqade/stim/dialects/collapse/emit_str.py

@@ -60,7 +60,9 @@ class EmitStimCollapseMethods(MethodTable):
         self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.PPMeasurement
     ):
         probability: str = frame.get(stmt.p)
-        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+        targets: tuple[str, ...] = tuple(
+            targ.upper() for targ in frame.get_values(stmt.targets)
+        )
 
         out = f"MPP({probability}) " + " ".join(targets)
         emit.writeln(frame, out)

bloqade/stim/dialects/gate/emit.py

@@ -12,6 +12,7 @@ from .stmts.base import SingleQubitGate, ControlledTwoQubitGate
 class EmitStimGateMethods(MethodTable):
 
     gate_1q_map: dict[str, tuple[str, str]] = {
+        stmts.Identity.name: ("I", "I"),
         stmts.X.name: ("X", "X"),
         stmts.Y.name: ("Y", "Y"),
         stmts.Z.name: ("Z", "Z"),
@@ -22,6 +23,7 @@ class EmitStimGateMethods(MethodTable):
         stmts.SqrtZ.name: ("SQRT_Z", "SQRT_Z_DAG"),
     }
 
+    @impl(stmts.Identity)
     @impl(stmts.X)
     @impl(stmts.Y)
     @impl(stmts.Z)
@@ -80,8 +82,13 @@ class EmitStimGateMethods(MethodTable):
     @impl(stmts.SPP)
     def spp(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.SPP):
 
-        targets: tuple[str, ...] = frame.get_values(stmt.targets)
-        res = "SPP " + " ".join(targets)
+        targets: tuple[str, ...] = tuple(
+            targ.upper() for targ in frame.get_values(stmt.targets)
+        )
+        if stmt.dagger:
+            res = "SPP_DAG " + " ".join(targets)
+        else:
+            res = "SPP " + " ".join(targets)
         emit.writeln(frame, res)
 
         return ()

bloqade/stim/dialects/noise/emit.py

@@ -44,7 +44,7 @@ class EmitStimNoiseMethods(MethodTable):
         px: str = frame.get(stmt.px)
         py: str = frame.get(stmt.py)
         pz: str = frame.get(stmt.pz)
-        res = f"PAULI_CHANNEL_1({px},{py},{pz}) " + " ".join(targets)
+        res = f"PAULI_CHANNEL_1({px}, {py}, {pz}) " + " ".join(targets)
         emit.writeln(frame, res)
 
         return ()
@@ -64,3 +64,34 @@ class EmitStimNoiseMethods(MethodTable):
         emit.writeln(frame, res)
 
         return ()
+
+    @impl(stmts.TrivialError)
+    def non_stim_error(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.TrivialError
+    ):
+
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+        prob: tuple[str, ...] = frame.get_values(stmt.probs)
+        prob_str: str = ", ".join(prob)
+
+        res = f"I_ERROR[{stmt.name}]({prob_str}) " + " ".join(targets)
+        emit.writeln(frame, res)
+
+        return ()
+
+    @impl(stmts.TrivialCorrelatedError)
+    def non_stim_corr_error(
+        self,
+        emit: EmitStimMain,
+        frame: EmitStrFrame,
+        stmt: stmts.TrivialCorrelatedError,
+    ):
+
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+        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)
+        emit.writeln(frame, res)
+
+        return ()

bloqade/stim/dialects/noise/stmts.py

@@ -75,3 +75,32 @@ class ZError(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     p: ir.SSAValue = info.argument(types.Float)
     targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
+
+
+@statement
+class NonStimError(ir.Statement):
+    name = "NonStimError"
+    traits = frozenset({lowering.FromPythonCall()})
+    probs: tuple[ir.SSAValue, ...] = info.argument(types.Float)
+    targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
+
+
+@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)
+
+
+@statement(dialect=dialect)
+class TrivialCorrelatedError(NonStimCorrelatedError):
+    name = "TRIV_CORR_ERROR"
+
+
+@statement(dialect=dialect)
+class TrivialError(NonStimError):
+    name = "TRIV_ERROR"

bloqade/stim/parse/__init__.py

@@ -0,0 +1 @@
+from .lowering import loads as loads, loadfile as loadfile