bloqade-circuit 0.6.2__py3-none-any.whl → 0.9.1__py3-none-any.whl

bloqade/stim/rewrite/set_detector_to_stim.py

@@ -0,0 +1,68 @@
+from typing import Iterable
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.dialects.py import Constant
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.stim.dialects import auxiliary
+from bloqade.annotate.stmts import SetDetector
+from bloqade.analysis.measure_id import MeasureIDFrame
+from bloqade.stim.dialects.auxiliary import Detector
+from bloqade.analysis.measure_id.lattice import MeasureIdTuple
+
+from ..rewrite.get_record_util import insert_get_records
+
+
+@dataclass
+class SetDetectorToStim(RewriteRule):
+    """
+    Rewrite SetDetector to GetRecord and Detector in the stim dialect
+    """
+
+    measure_id_frame: MeasureIDFrame
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        match node:
+            case SetDetector():
+                return self.rewrite_SetDetector(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_SetDetector(self, node: SetDetector) -> RewriteResult:
+
+        # get coordinates and generate correct consts
+        coord_ssas = []
+        if not isinstance(node.coordinates.owner, Constant):
+            return RewriteResult()
+
+        coord_values = node.coordinates.owner.value.unwrap()
+
+        if not isinstance(coord_values, Iterable):
+            return RewriteResult()
+
+        if any(not isinstance(value, (int, float)) for value in coord_values):
+            return RewriteResult()
+
+        for coord_value in coord_values:
+            if isinstance(coord_value, float):
+                coord_stmt = auxiliary.ConstFloat(value=coord_value)
+            else:  # int
+                coord_stmt = auxiliary.ConstInt(value=coord_value)
+            coord_ssas.append(coord_stmt.result)
+            coord_stmt.insert_before(node)
+
+        measure_ids = self.measure_id_frame.entries[node.measurements]
+        assert isinstance(measure_ids, MeasureIdTuple)
+
+        get_record_list = insert_get_records(
+            node, measure_ids, self.measure_id_frame.num_measures_at_stmt[node]
+        )
+
+        detector_stmt = Detector(
+            coord=tuple(coord_ssas), targets=tuple(get_record_list)
+        )
+
+        node.replace_by(detector_stmt)
+
+        return RewriteResult(has_done_something=True)

bloqade/stim/rewrite/set_observable_to_stim.py

@@ -0,0 +1,52 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.stim.dialects import auxiliary
+from bloqade.annotate.stmts import SetObservable
+from bloqade.analysis.measure_id import MeasureIDFrame
+from bloqade.stim.dialects.auxiliary import ObservableInclude
+from bloqade.analysis.measure_id.lattice import MeasureIdTuple
+
+from ..rewrite.get_record_util import insert_get_records
+
+
+@dataclass
+class SetObservableToStim(RewriteRule):
+    """
+    Rewrite SetObservable to GetRecord and ObservableInclude in the stim dialect
+    """
+
+    measure_id_frame: MeasureIDFrame
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        match node:
+            case SetObservable():
+                return self.rewrite_SetObservable(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_SetObservable(self, node: SetObservable) -> RewriteResult:
+
+        # set idx to 0 for now, but this
+        # should be something that a user can set on their own.
+        # SetObservable needs to accept an int.
+
+        idx_stmt = auxiliary.ConstInt(value=0)
+        idx_stmt.insert_before(node)
+
+        measure_ids = self.measure_id_frame.entries[node.measurements]
+        assert isinstance(measure_ids, MeasureIdTuple)
+
+        get_record_list = insert_get_records(
+            node, measure_ids, self.measure_id_frame.num_measures_at_stmt[node]
+        )
+
+        observable_include_stmt = ObservableInclude(
+            idx=idx_stmt.result, targets=tuple(get_record_list)
+        )
+
+        node.replace_by(observable_include_stmt)
+
+        return RewriteResult(has_done_something=True)

bloqade/stim/rewrite/squin_measure.py

@@ -2,47 +2,15 @@
 from dataclasses import dataclass
 
 from kirin import ir
-from kirin.dialects import py, ilist
+from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import wire, qubit
+from bloqade import qubit
 from bloqade.squin.rewrite import AddressAttribute
-from bloqade.stim.dialects import collapse, auxiliary
+from bloqade.stim.dialects import collapse
 from bloqade.stim.rewrite.util import (
-    is_measure_result_used,
     insert_qubit_idx_from_address,
 )
-from bloqade.analysis.measure_id.lattice import MeasureId, MeasureIdBool, MeasureIdTuple
-
-
-def replace_get_record(
-    node: ir.Statement, measure_id_bool: MeasureIdBool, meas_count: int
-):
-    assert isinstance(measure_id_bool, MeasureIdBool)
-    target_rec_idx = (measure_id_bool.idx - 1) - meas_count
-    idx_stmt = py.constant.Constant(target_rec_idx)
-    idx_stmt.insert_before(node)
-    get_record_stmt = auxiliary.GetRecord(idx_stmt.result)
-    node.replace_by(get_record_stmt)
-
-
-def insert_get_record_list(
-    node: ir.Statement, measure_id_tuple: MeasureIdTuple, meas_count: int
-):
-    """
-    Insert GetRecord statements before the given node
-    """
-    get_record_ssas = []
-    for measure_id_bool in measure_id_tuple.data:
-        assert isinstance(measure_id_bool, MeasureIdBool)
-        target_rec_idx = (measure_id_bool.idx - 1) - meas_count
-        idx_stmt = py.constant.Constant(target_rec_idx)
-        idx_stmt.insert_before(node)
-        get_record_stmt = auxiliary.GetRecord(idx_stmt.result)
-        get_record_stmt.insert_before(node)
-        get_record_ssas.append(get_record_stmt.result)
-
-    node.replace_by(ilist.New(values=get_record_ssas))
 
 
 @dataclass
@@ -51,29 +19,20 @@ class SquinMeasureToStim(RewriteRule):
     Rewrite squin measure-related statements to stim statements.
     """
 
-    measure_id_result: dict[ir.SSAValue, MeasureId]
-    total_measure_count: int
-
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
         match node:
-            case qubit.MeasureQubit() | qubit.MeasureQubitList() | wire.Measure():
+            case qubit.stmts.Measure():
                 return self.rewrite_Measure(node)
             case _:
                 return RewriteResult()
 
-    def rewrite_Measure(
-        self, measure_stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure
-    ) -> RewriteResult:
+    def rewrite_Measure(self, measure_stmt: qubit.stmts.Measure) -> RewriteResult:
 
         qubit_idx_ssas = self.get_qubit_idx_ssas(measure_stmt)
         if qubit_idx_ssas is None:
             return RewriteResult()
 
-        measure_id = self.measure_id_result[measure_stmt.result]
-        if not isinstance(measure_id, (MeasureIdBool, MeasureIdTuple)):
-            return RewriteResult()
-
         prob_noise_stmt = py.constant.Constant(0.0)
         stim_measure_stmt = collapse.MZ(
             p=prob_noise_stmt.result,
@@ -82,48 +41,21 @@ class SquinMeasureToStim(RewriteRule):
         prob_noise_stmt.insert_before(measure_stmt)
         stim_measure_stmt.insert_before(measure_stmt)
 
-        if not is_measure_result_used(measure_stmt):
+        # if the measurement is not being used anywhere
+        # we can safely get rid of it. Measure cannot be DCE'd because
+        # it is not pure.
+        if not bool(measure_stmt.result.uses):
             measure_stmt.delete()
-            return RewriteResult(has_done_something=True)
-
-        # replace dataflow with new stmt!
-        measure_id = self.measure_id_result[measure_stmt.result]
-        if isinstance(measure_id, MeasureIdBool):
-            replace_get_record(
-                node=measure_stmt,
-                measure_id_bool=measure_id,
-                meas_count=self.total_measure_count,
-            )
-        elif isinstance(measure_id, MeasureIdTuple):
-            insert_get_record_list(
-                node=measure_stmt,
-                measure_id_tuple=measure_id,
-                meas_count=self.total_measure_count,
-            )
-        else:
-            # already checked before, so this should not happen
-            raise ValueError(
-                f"Unexpected measure ID type: {type(measure_id)} for measure statement {measure_stmt}"
-            )
 
         return RewriteResult(has_done_something=True)
 
     def get_qubit_idx_ssas(
-        self, measure_stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure
+        self, measure_stmt: qubit.stmts.Measure
     ) -> tuple[ir.SSAValue, ...] | None:
         """
         Extract the address attribute and insert qubit indices for the given measure statement.
         """
-        match measure_stmt:
-            case qubit.MeasureQubit():
-                address_attr = measure_stmt.qubit.hints.get("address")
-            case qubit.MeasureQubitList():
-                address_attr = measure_stmt.qubits.hints.get("address")
-            case wire.Measure():
-                address_attr = measure_stmt.wire.hints.get("address")
-            case _:
-                return None
-
+        address_attr = measure_stmt.qubits.hints.get("address")
         if address_attr is None:
             return None
 

bloqade/stim/rewrite/squin_noise.py

@@ -1,17 +1,17 @@
+import itertools
 from typing import Tuple
 from dataclasses import dataclass
 
+from kirin import types
 from kirin.ir import SSAValue, Statement
-from kirin.dialects import py, ilist
+from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import op, wire, noise as squin_noise, qubit
+from bloqade.squin import noise as squin_noise
 from bloqade.stim.dialects import noise as stim_noise
-from bloqade.stim.rewrite.util import (
-    get_const_value,
-    create_wire_passthrough,
-    insert_qubit_idx_after_apply,
-)
+from bloqade.stim.rewrite.util import insert_qubit_idx_from_address
+from bloqade.analysis.address.lattice import AddressReg, PartialIList
+from bloqade.squin.rewrite.wrap_analysis import AddressAttribute
 
 
 @dataclass
@@ -19,157 +19,183 @@ class SquinNoiseToStim(RewriteRule):
 
     def rewrite_Statement(self, node: Statement) -> RewriteResult:
         match node:
-            case qubit.Apply() | qubit.Broadcast():
-                return self.rewrite_Apply_and_Broadcast(node)
+            case squin_noise.stmts.NoiseChannel():
+                return self.rewrite_NoiseChannel(node)
             case _:
                 return RewriteResult()
 
-    def rewrite_Apply_and_Broadcast(
-        self, stmt: qubit.Apply | qubit.Broadcast
+    def rewrite_NoiseChannel(
+        self, stmt: squin_noise.stmts.NoiseChannel
     ) -> RewriteResult:
-        """Rewrite Apply and Broadcast to their stim statements."""
+        """Rewrite NoiseChannel statements to their stim equivalents."""
 
-        # this is an SSAValue, need it to be the actual operator
-        applied_op = stmt.operator.owner
+        rewrite_method = getattr(self, f"rewrite_{type(stmt).__name__}", None)
 
-        if isinstance(applied_op, squin_noise.stmts.QubitLoss):
+        # No rewrite method exists and the rewrite should stop
+        if rewrite_method is None:
             return RewriteResult()
+        if isinstance(stmt, squin_noise.stmts.CorrelatedQubitLoss):
+            # CorrelatedQubitLoss represents a broadcast operation, but Stim does not
+            # support broadcasting for multi-qubit noise channels.
+            # Therefore, we must expand the broadcast into individual stim statements.
+            qubit_address_attr = stmt.qubits.hints.get("address", None)
 
-        if isinstance(applied_op, squin_noise.stmts.NoiseChannel):
+            if not isinstance(qubit_address_attr, AddressAttribute):
+                return RewriteResult()
+
+            if not isinstance(address := qubit_address_attr.address, PartialIList):
+                return RewriteResult()
 
-            qubit_idx_ssas = insert_qubit_idx_after_apply(stmt=stmt)
-            if qubit_idx_ssas is None:
+            if not types.is_tuple_of(data := address.data, AddressReg):
                 return RewriteResult()
 
-            rewrite_method = getattr(self, f"rewrite_{type(applied_op).__name__}")
-            stim_stmt = rewrite_method(stmt, qubit_idx_ssas)
+            for address_reg in data:
 
-            if isinstance(stmt, (wire.Apply, wire.Broadcast)):
-                create_wire_passthrough(stmt)
+                qubit_idx_ssas = insert_qubit_idx_from_address(
+                    AddressAttribute(address_reg), stmt
+                )
 
-            if stim_stmt is not None:
-                stmt.replace_by(stim_stmt)
-            if len(stmt.operator.owner.result.uses) == 0:
-                stmt.operator.owner.delete()
+                stim_stmt = rewrite_method(stmt, qubit_idx_ssas)
+                stim_stmt.insert_before(stmt)
+
+            stmt.delete()
 
             return RewriteResult(has_done_something=True)
-        return RewriteResult()
 
-    def rewrite_PauliError(
-        self,
-        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
-        qubit_idx_ssas: Tuple[SSAValue],
-    ) -> Statement:
-        """Rewrite squin.noise.PauliError to XError, YError, ZError."""
-        squin_channel = stmt.operator.owner
-        assert isinstance(squin_channel, squin_noise.stmts.PauliError)
-        basis = squin_channel.basis.owner
-        assert isinstance(basis, op.stmts.PauliOp)
-        p = get_const_value(float, squin_channel.p)
-
-        p_stmt = py.Constant(p)
-        p_stmt.insert_before(stmt)
-
-        if isinstance(basis, op.stmts.X):
-            stim_stmt = stim_noise.XError(targets=qubit_idx_ssas, p=p_stmt.result)
-        elif isinstance(basis, op.stmts.Y):
-            stim_stmt = stim_noise.YError(targets=qubit_idx_ssas, p=p_stmt.result)
+        if isinstance(stmt, squin_noise.stmts.SingleQubitNoiseChannel):
+            qubit_address_attr = stmt.qubits.hints.get("address", None)
+            if qubit_address_attr is None:
+                return RewriteResult()
+            qubit_idx_ssas = insert_qubit_idx_from_address(qubit_address_attr, stmt)
+
+        elif isinstance(stmt, squin_noise.stmts.TwoQubitNoiseChannel):
+            control_address_attr = stmt.controls.hints.get("address", None)
+            target_address_attr = stmt.targets.hints.get("address", None)
+            if control_address_attr is None or target_address_attr is None:
+                return RewriteResult()
+            control_qubit_idx_ssas = insert_qubit_idx_from_address(
+                control_address_attr, stmt
+            )
+            target_qubit_idx_ssas = insert_qubit_idx_from_address(
+                target_address_attr, stmt
+            )
+            if control_qubit_idx_ssas is None or target_qubit_idx_ssas is None:
+                return RewriteResult()
+
+            # For stim statements you want to interleave the control and target qubit indices:
+            # ex: CX controls = (0,1) targets = (2,3) in stim is: CX 0 2 1 3
+            qubit_idx_ssas = list(
+                itertools.chain.from_iterable(
+                    zip(control_qubit_idx_ssas, target_qubit_idx_ssas)
+                )
+            )
         else:
-            stim_stmt = stim_noise.ZError(targets=qubit_idx_ssas, p=p_stmt.result)
-        return stim_stmt
+            return RewriteResult()
+
+        # guaranteed that you have a valid stim_stmt to plug in
+        stim_stmt = rewrite_method(stmt, tuple(qubit_idx_ssas))
+        stmt.replace_by(stim_stmt)
+
+        return RewriteResult(has_done_something=True)
 
     def rewrite_SingleQubitPauliChannel(
         self,
-        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+        stmt: squin_noise.stmts.SingleQubitPauliChannel,
         qubit_idx_ssas: Tuple[SSAValue],
     ) -> Statement:
         """Rewrite squin.noise.SingleQubitPauliChannel to stim.PauliChannel1."""
 
-        squin_channel = stmt.operator.owner
-        assert isinstance(squin_channel, squin_noise.stmts.SingleQubitPauliChannel)
-
-        params = get_const_value(ilist.IList, squin_channel.params)
-        new_stmts = [
-            p_x := py.Constant(params[0]),
-            p_y := py.Constant(params[1]),
-            p_z := py.Constant(params[2]),
-        ]
-        for new_stmt in new_stmts:
-            new_stmt.insert_before(stmt)
-
         stim_stmt = stim_noise.PauliChannel1(
             targets=qubit_idx_ssas,
-            px=p_x.result,
-            py=p_y.result,
-            pz=p_z.result,
+            px=stmt.px,
+            py=stmt.py,
+            pz=stmt.pz,
         )
         return stim_stmt
 
-    def rewrite_TwoQubitPauliChannel(
+    def rewrite_QubitLoss(
         self,
-        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+        stmt: squin_noise.stmts.QubitLoss,
         qubit_idx_ssas: Tuple[SSAValue],
     ) -> Statement:
-        """Rewrite squin.noise.SingleQubitPauliChannel to stim.PauliChannel1."""
+        """Rewrite squin.noise.QubitLoss to stim.TrivialError."""
 
-        squin_channel = stmt.operator.owner
-        assert isinstance(squin_channel, squin_noise.stmts.TwoQubitPauliChannel)
+        stim_stmt = stim_noise.QubitLoss(
+            targets=qubit_idx_ssas,
+            probs=(stmt.p,),
+        )
 
-        params = get_const_value(ilist.IList, squin_channel.params)
-        param_stmts = [py.Constant(p) for p in params]
-        for param_stmt in param_stmts:
-            param_stmt.insert_before(stmt)
+        return stim_stmt
 
-        stim_stmt = stim_noise.PauliChannel2(
+    def rewrite_CorrelatedQubitLoss(
+        self,
+        stmt: squin_noise.stmts.CorrelatedQubitLoss,
+        qubit_idx_ssas: Tuple[SSAValue],
+    ) -> Statement:
+        """Rewrite squin.noise.CorrelatedQubitLoss to stim.CorrelatedQubitLoss."""
+        stim_stmt = stim_noise.CorrelatedQubitLoss(
             targets=qubit_idx_ssas,
-            pix=param_stmts[0].result,
-            piy=param_stmts[1].result,
-            piz=param_stmts[2].result,
-            pxi=param_stmts[3].result,
-            pxx=param_stmts[4].result,
-            pxy=param_stmts[5].result,
-            pxz=param_stmts[6].result,
-            pyi=param_stmts[7].result,
-            pyx=param_stmts[8].result,
-            pyy=param_stmts[9].result,
-            pyz=param_stmts[10].result,
-            pzi=param_stmts[11].result,
-            pzx=param_stmts[12].result,
-            pzy=param_stmts[13].result,
-            pzz=param_stmts[14].result,
+            probs=(stmt.p,),
         )
+
         return stim_stmt
 
-    def rewrite_Depolarize2(
+    def rewrite_Depolarize(
         self,
-        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+        stmt: squin_noise.stmts.Depolarize,
         qubit_idx_ssas: Tuple[SSAValue],
     ) -> Statement:
-        """Rewrite squin.noise.Depolarize2 to stim.Depolarize2."""
-
-        squin_channel = stmt.operator.owner
-        assert isinstance(squin_channel, squin_noise.stmts.Depolarize2)
+        """Rewrite squin.noise.Depolarize to stim.Depolarize1."""
 
-        p = get_const_value(float, squin_channel.p)
-        p_stmt = py.Constant(p)
-        p_stmt.insert_before(stmt)
+        stim_stmt = stim_noise.Depolarize1(
+            targets=qubit_idx_ssas,
+            p=stmt.p,
+        )
 
-        stim_stmt = stim_noise.Depolarize2(targets=qubit_idx_ssas, p=p_stmt.result)
         return stim_stmt
 
-    def rewrite_Depolarize(
+    def rewrite_TwoQubitPauliChannel(
         self,
-        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+        stmt: squin_noise.stmts.TwoQubitPauliChannel,
         qubit_idx_ssas: Tuple[SSAValue],
     ) -> Statement:
-        """Rewrite squin.noise.Depolarize to stim.Depolarize1."""
+        """Rewrite squin.noise.TwoQubitPauliChannel to stim.PauliChannel2."""
+
+        params = stmt.probabilities
+        prob_ssas = []
+        for idx in range(15):
+            idx_stmt = py.Constant(value=idx)
+            idx_stmt.insert_before(stmt)
+            getitem_stmt = py.GetItem(obj=params, index=idx_stmt.result)
+            getitem_stmt.insert_before(stmt)
+            prob_ssas.append(getitem_stmt.result)
 
-        squin_channel = stmt.operator.owner
-        assert isinstance(squin_channel, squin_noise.stmts.Depolarize)
+        stim_stmt = stim_noise.PauliChannel2(
+            targets=qubit_idx_ssas,
+            pix=prob_ssas[0],
+            piy=prob_ssas[1],
+            piz=prob_ssas[2],
+            pxi=prob_ssas[3],
+            pxx=prob_ssas[4],
+            pxy=prob_ssas[5],
+            pxz=prob_ssas[6],
+            pyi=prob_ssas[7],
+            pyx=prob_ssas[8],
+            pyy=prob_ssas[9],
+            pyz=prob_ssas[10],
+            pzi=prob_ssas[11],
+            pzx=prob_ssas[12],
+            pzy=prob_ssas[13],
+            pzz=prob_ssas[14],
+        )
+        return stim_stmt
 
-        p = get_const_value(float, squin_channel.p)
-        p_stmt = py.Constant(p)
-        p_stmt.insert_before(stmt)
+    def rewrite_Depolarize2(
+        self,
+        stmt: squin_noise.stmts.Depolarize2,
+        qubit_idx_ssas: Tuple[SSAValue],
+    ) -> Statement:
+        """Rewrite squin.noise.Depolarize2 to stim.Depolarize2."""
 
-        stim_stmt = stim_noise.Depolarize1(targets=qubit_idx_ssas, p=p_stmt.result)
+        stim_stmt = stim_noise.Depolarize2(targets=qubit_idx_ssas, p=stmt.p)
         return stim_stmt