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

bloqade/stim/parse/lowering.py

@@ -98,6 +98,8 @@ def loads(
         signature=func.Signature((), return_node.value.type),
         body=body,
     )
+    self_arg = ir.BlockArgument(body.blocks[0], 0)  # Self argument
+    body.blocks[0]._args = (self_arg,)
     return ir.Method(
         mod=None,
         py_func=None,
@@ -627,10 +629,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 +648,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 +1,3 @@
-from .squin_to_stim import SquinToStimPass as SquinToStimPass
+from .squin_to_stim import (
+    SquinToStimPass as SquinToStimPass,
+)

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.rewrite.passes 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,11 @@ 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
 
@@ -19,12 +22,23 @@ 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)
 
+        # because nested python lists don't have their
+        # member lists converted to ILists, ConstantFold
+        # can add python lists that can't be hashed, causing
+        # issues with CSE. ConstList2IList remedies that problem here.
         result = (
-            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            Chain(
+                Fixpoint(Walk(ConstantFold())),
+                Walk(ConstList2IList()),
+                Walk(CommonSubexpressionElimination()),
+            )
             .rewrite(mt.code)
             .join(result)
         )

bloqade/stim/passes/squin_to_stim.py

@@ -1,29 +1,21 @@
 from dataclasses import dataclass
 
-from kirin.passes import Fold
 from kirin.rewrite import (
     Walk,
     Chain,
     Fixpoint,
-    CFGCompactify,
-    InlineGetItem,
-    InlineGetField,
     DeadCodeElimination,
     CommonSubexpressionElimination,
 )
-from kirin.dialects import scf, 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 bloqade.stim.rewrite import (
-    SquinWireToStim,
     PyConstantToStim,
     SquinNoiseToStim,
     SquinQubitToStim,
     SquinMeasureToStim,
-    SquinWireIdentityElimination,
 )
 from bloqade.squin.rewrite import (
     SquinU3ToClifford,
@@ -33,9 +25,9 @@ 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.stim.passes.flatten import Flatten
 
-from .simplify_ifs import StimSimplifyIfs
-from ..rewrite.ifs_to_stim import IfToStim
+from ..rewrite import IfToStim, SetDetectorToStim, SetObservableToStim
 
 
 @dataclass
@@ -44,52 +36,18 @@ 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)
-
-        rule = Chain(
-            InlineGetField(),
-            InlineGetItem(),
-            scf.unroll.ForLoop(),
-            scf.trim.UnusedYield(),
-        )
-        rewrite_result = Fixpoint(Walk(rule)).rewrite(mt.code).join(rewrite_result)
-        # fold_pass = Fold(mt.dialects, no_raise=self.no_raise)
-        # rewrite_result = fold_pass(mt)
-        rewrite_result = (
-            Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(rewrite_result)
-        )
-        rewrite_result = (
-            StimSimplifyIfs(mt.dialects, no_raise=self.no_raise)
-            .unsafe_run(mt)
-            .join(rewrite_result)
-        )
-
-        # run typeinfer again after unroll etc. because we now insert
-        # a lot of new nodes, which might have more precise types
-        # self.typeinfer.unsafe_run(mt)
-        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 = (
-            CanonicalizeIList(dialects=mt.dialects, no_raise=self.no_raise)
-            .unsafe_run(mt)
-            .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
         # -------------------------------------------------------------------
 
         mia = MeasurementIDAnalysis(dialects=mt.dialects)
-        meas_analysis_frame, _ = mia.run_analysis(mt, no_raise=self.no_raise)
+        meas_analysis_frame, _ = mia.run(mt)
 
         aa = AddressAnalysis(dialects=mt.dialects)
-        address_analysis_frame, _ = aa.run_analysis(mt, no_raise=self.no_raise)
+        address_analysis_frame, _ = aa.run(mt)
 
         # wrap the address analysis result
         rewrite_result = (
@@ -99,12 +57,16 @@ class SquinToStimPass(Pass):
         )
 
         # 2. rewrite
+        ## Invoke DCE afterwards to eliminate any GetItems
+        ## that are no longer being used. This allows for
+        ## SquinMeasureToStim to safely eliminate
+        ## unused measure statements.
         rewrite_result = (
-            Walk(
-                IfToStim(
-                    measure_analysis=meas_analysis_frame.entries,
-                    measure_count=mia.measure_count,
-                )
+            Chain(
+                Walk(IfToStim(measure_frame=meas_analysis_frame)),
+                Walk(SetDetectorToStim(measure_id_frame=meas_analysis_frame)),
+                Walk(SetObservableToStim(measure_id_frame=meas_analysis_frame)),
+                Fixpoint(Walk(DeadCodeElimination())),
             )
             .rewrite(mt.code)
             .join(rewrite_result)
@@ -120,17 +82,13 @@ class SquinToStimPass(Pass):
             Walk(
                 Chain(
                     SquinQubitToStim(),
-                    SquinWireToStim(),
-                    SquinMeasureToStim(
-                        measure_id_result=meas_analysis_frame.entries,
-                        total_measure_count=mia.measure_count,
-                    ),  # reduce duplicated logic, can split out even more rules later
-                    SquinWireIdentityElimination(),
+                    SquinMeasureToStim(),
                 )
             )
             .rewrite(mt.code)
             .join(rewrite_result)
         )
+
         rewrite_result = (
             CanonicalizeIList(dialects=mt.dialects, no_raise=self.no_raise)
             .unsafe_run(mt)
@@ -141,7 +99,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,8 +1,7 @@
+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,
-)
+from .set_detector_to_stim import SetDetectorToStim as SetDetectorToStim
+from .set_observable_to_stim import SetObservableToStim as SetObservableToStim

bloqade/stim/rewrite/get_record_util.py

@@ -0,0 +1,24 @@
+from kirin import ir
+from kirin.dialects import py
+
+from bloqade.stim.dialects import auxiliary
+from bloqade.analysis.measure_id.lattice import MeasureIdBool, MeasureIdTuple
+
+
+def insert_get_records(
+    node: ir.Statement, measure_id_tuple: MeasureIdTuple, meas_count_at_stmt: 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_at_stmt
+        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)
+
+    return get_record_ssas

bloqade/stim/rewrite/ifs_to_stim.py

@@ -4,16 +4,16 @@ 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 (
-    MeasureId,
     MeasureIdBool,
 )
 
@@ -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:
@@ -127,8 +127,7 @@ class IfToStim(IfElseSimplification, RewriteRule):
     Rewrite if statements to stim equivalent statements.
     """
 
-    measure_analysis: dict[ir.SSAValue, MeasureId]
-    measure_count: int
+    measure_frame: MeasureIDFrame
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
@@ -140,38 +139,37 @@ class IfToStim(IfElseSimplification, RewriteRule):
 
     def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
 
-        if not isinstance(self.measure_analysis[stmt.cond], MeasureIdBool):
+        # 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
-        measure_id_bool = self.measure_analysis[stmt.cond]
+        measure_id_bool = self.measure_frame.entries[stmt.cond]
         assert isinstance(measure_id_bool, MeasureIdBool)
 
         # generate get record statement
         measure_id_idx_stmt = py.Constant(
-            (measure_id_bool.idx - 1) - self.measure_count
+            (measure_id_bool.idx - 1) - self.measure_frame.num_measures_at_stmt[stmt]
         )
         get_record_stmt = GetRecord(id=measure_id_idx_stmt.result)  # noqa: F841
 
-        # get address attribute and generate qubit idx statements
-        address_attr = apply_or_broadcast.qubits.hints.get("address")
+        address_attr = stmts[0].qubits.hints.get("address")
+
         if address_attr is None:
             return RewriteResult()
         assert isinstance(address_attr, AddressAttribute)

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