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

bloqade/stim/rewrite/qubit_to_stim.py

@@ -1,13 +1,11 @@
 from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import op, noise, qubit
+from bloqade import qubit
+from bloqade.squin import gate
 from bloqade.squin.rewrite import AddressAttribute
-from bloqade.stim.dialects import gate
+from bloqade.stim.dialects import gate as stim_gate, collapse as stim_collapse
 from bloqade.stim.rewrite.util import (
-    SQUIN_STIM_OP_MAPPING,
-    rewrite_Control,
-    rewrite_QubitLoss,
     insert_qubit_idx_from_address,
 )
 
@@ -20,64 +18,115 @@ class SquinQubitToStim(RewriteRule):
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
         match node:
-            case qubit.Apply() | qubit.Broadcast():
-                return self.rewrite_Apply_and_Broadcast(node)
+            # not supported by Stim
+            case gate.stmts.T() | gate.stmts.RotationGate():
+                return RewriteResult()
+            # If you've reached this point all gates have stim equivalents
+            case qubit.stmts.Reset():
+                return self.rewrite_Reset(node)
+            case gate.stmts.SingleQubitGate():
+                return self.rewrite_SingleQubitGate(node)
+            case gate.stmts.ControlledGate():
+                return self.rewrite_ControlledGate(node)
             case _:
                 return RewriteResult()
 
-    def rewrite_Apply_and_Broadcast(
-        self, stmt: qubit.Apply | qubit.Broadcast
-    ) -> RewriteResult:
-        """
-        Rewrite Apply and Broadcast nodes to their stim equivalent statements.
-        """
-
-        # this is an SSAValue, need it to be the actual operator
-        applied_op = stmt.operator.owner
+    def rewrite_Reset(self, stmt: qubit.stmts.Reset) -> RewriteResult:
 
-        if isinstance(applied_op, noise.stmts.QubitLoss):
-            return rewrite_QubitLoss(stmt)
+        qubit_addr_attr = stmt.qubits.hints.get("address", None)
 
-        assert isinstance(applied_op, op.stmts.Operator)
+        if qubit_addr_attr is None:
+            return RewriteResult()
 
-        if isinstance(applied_op, op.stmts.Control):
-            return rewrite_Control(stmt)
+        assert isinstance(qubit_addr_attr, AddressAttribute)
 
-        # need to handle Control through separate means
+        qubit_idx_ssas = insert_qubit_idx_from_address(
+            address=qubit_addr_attr, stmt_to_insert_before=stmt
+        )
 
-        # check if its adjoint, assume its canonicalized so no nested adjoints.
-        is_conj = False
-        if isinstance(applied_op, op.stmts.Adjoint):
-            if not applied_op.is_unitary:
-                return RewriteResult()
+        if qubit_idx_ssas is None:
+            return RewriteResult()
 
-            is_conj = True
-            applied_op = applied_op.op.owner
+        stim_stmt = stim_collapse.RZ(targets=tuple(qubit_idx_ssas))
+        stmt.replace_by(stim_stmt)
 
-        stim_1q_op = SQUIN_STIM_OP_MAPPING.get(type(applied_op))
-        if stim_1q_op is None:
-            return RewriteResult()
+        return RewriteResult(has_done_something=True)
 
-        address_attr = stmt.qubits[0].hints.get("address")
+    def rewrite_SingleQubitGate(
+        self, stmt: gate.stmts.SingleQubitGate
+    ) -> RewriteResult:
+        """
+        Rewrite single qubit gate nodes to their stim equivalent statements.
+        Address Analysis should have been run along with Wrap Analysis before this rewrite is applied.
+        """
 
-        if address_attr is None:
+        qubit_addr_attr = stmt.qubits.hints.get("address", None)
+        if qubit_addr_attr is None:
             return RewriteResult()
 
-        assert isinstance(address_attr, AddressAttribute)
+        assert isinstance(qubit_addr_attr, AddressAttribute)
+
         qubit_idx_ssas = insert_qubit_idx_from_address(
-            address=address_attr, stmt_to_insert_before=stmt
+            address=qubit_addr_attr, stmt_to_insert_before=stmt
         )
 
         if qubit_idx_ssas is None:
             return RewriteResult()
 
-        if isinstance(stim_1q_op, gate.stmts.Gate):
-            stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas), dagger=is_conj)
+        # Get the name of the inputted stmt and see if there is an
+        # equivalently named statement in stim,
+        # then create an instance of that stim statement
+        stmt_name = type(stmt).__name__
+        stim_stmt_cls = getattr(stim_gate.stmts, stmt_name, None)
+        if stim_stmt_cls is None:
+            return RewriteResult()
+
+        if isinstance(stmt, gate.stmts.SingleQubitNonHermitianGate):
+            stim_stmt = stim_stmt_cls(
+                targets=tuple(qubit_idx_ssas), dagger=stmt.adjoint
+            )
         else:
-            stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
-        stmt.replace_by(stim_1q_stmt)
+            stim_stmt = stim_stmt_cls(targets=tuple(qubit_idx_ssas))
+        stmt.replace_by(stim_stmt)
 
         return RewriteResult(has_done_something=True)
 
+    def rewrite_ControlledGate(self, stmt: gate.stmts.ControlledGate) -> RewriteResult:
+        """
+        Rewrite controlled gate nodes to their stim equivalent statements.
+        Address Analysis should have been run along with Wrap Analysis before this rewrite is applied.
+        """
 
-# put rewrites for measure statements in separate rule, then just have to dispatch
+        controls_addr_attr = stmt.controls.hints.get("address", None)
+        targets_addr_attr = stmt.targets.hints.get("address", None)
+
+        if controls_addr_attr is None or targets_addr_attr is None:
+            return RewriteResult()
+
+        assert isinstance(controls_addr_attr, AddressAttribute)
+        assert isinstance(targets_addr_attr, AddressAttribute)
+
+        controls_idx_ssas = insert_qubit_idx_from_address(
+            address=controls_addr_attr, stmt_to_insert_before=stmt
+        )
+        targets_idx_ssas = insert_qubit_idx_from_address(
+            address=targets_addr_attr, stmt_to_insert_before=stmt
+        )
+
+        if controls_idx_ssas is None or targets_idx_ssas is None:
+            return RewriteResult()
+
+        # Get the name of the inputted stmt and see if there is an
+        # equivalently named statement in stim,
+        # then create an instance of that stim statement
+        stmt_name = type(stmt).__name__
+        stim_stmt_cls = getattr(stim_gate.stmts, stmt_name, None)
+        if stim_stmt_cls is None:
+            return RewriteResult()
+
+        stim_stmt = stim_stmt_cls(
+            targets=tuple(targets_idx_ssas), controls=tuple(controls_idx_ssas)
+        )
+        stmt.replace_by(stim_stmt)
+
+        return RewriteResult(has_done_something=True)

bloqade/stim/rewrite/squin_measure.py

@@ -5,11 +5,10 @@ from kirin import ir
 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
 from bloqade.stim.rewrite.util import (
-    is_measure_result_used,
     insert_qubit_idx_from_address,
 )
 
@@ -23,14 +22,12 @@ class SquinMeasureToStim(RewriteRule):
     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:
@@ -44,27 +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)
 
     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,159 +19,183 @@ class SquinNoiseToStim(RewriteRule):
 
     def rewrite_Statement(self, node: Statement) -> RewriteResult:
         match node:
-            case qubit.Apply() | qubit.Broadcast() | wire.Apply() | wire.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 | wire.Apply | wire.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()
 
-            rewrite_method = getattr(self, f"rewrite_{type(applied_op).__name__}", None)
-            # No rewrite method exists and the rewrite should stop
-            if rewrite_method is None:
+            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()
 
-            stim_stmt = rewrite_method(stmt, qubit_idx_ssas)
+            for address_reg in data:
+
+                qubit_idx_ssas = insert_qubit_idx_from_address(
+                    AddressAttribute(address_reg), stmt
+                )
 
-            if isinstance(stmt, (wire.Apply, wire.Broadcast)):
-                create_wire_passthrough(stmt)
+                stim_stmt = rewrite_method(stmt, qubit_idx_ssas)
+                stim_stmt.insert_before(stmt)
 
-            # guaranteed that you have a valid stim_stmt to plug in
-            stmt.replace_by(stim_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