Adapt SquinToStim to use new MeasureIDAnalysis (#414)

This is dependent on PR #413 (Please review that one first) but the new
SquinToStim pass should now be able to generate the correct record
indices. A couple of comments in terms of the changes:

- I don't see a need to generate a GetRecord unless it NEEDS to be
there. The only time this is a reality is when IfsToStim sees a
compatible boolean condition, in which case a GetRecord should be
generated.
- There was one pre-existing test with regards to SquinToStim that
failed with the new changes, but this was because the record indexes
were historically incorrect. The test in question is `simple_if_rewrite`
in `test/stim/passes/test_squin_meas_to_stim.py`

```
# The test's structure (in pseudocode) is:

q = qubits(4)
ms = measure(q)

if ms[0]
    # apply some gates that can be translated to feedforward

if ms[1]
    # same as above
   
measure(q) # another measurement
```

Because of #398 the global number of measures was used and the indexes
generated were `-7` and `-8`. I think this is still technically correct
(they are first - and oldest - measurements that occur) but the more
correct behavior is that at the time of the scf.IfElse you only have 4
measurements total so the correct index generated should be -4 and -3.

---------

Co-authored-by: Kai-Hsin Wu <[email protected]>

Author: johnzl-777

src/bloqade/analysis/measure_id/__init__.py

@@ -1,2 +1,5 @@
 from . import impls as impls
-from .analysis import MeasurementIDAnalysis as MeasurementIDAnalysis
+from .analysis import (
+    MeasureIDFrame as MeasureIDFrame,
+    MeasurementIDAnalysis as MeasurementIDAnalysis,
+)

src/bloqade/stim/passes/__init__.py

@@ -1 +1,5 @@
-from .squin_to_stim import SquinToStimPass as SquinToStimPass
+from .squin_to_stim import (
+    SquinToStimPass as SquinToStimPass,
+    StimSimplifyIfs as StimSimplifyIfs,
+    AggressiveForLoopUnroll as AggressiveForLoopUnroll,
+)

src/bloqade/stim/passes/squin_to_stim.py

@@ -128,12 +128,14 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
         )
 
         # 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)),
+                Fixpoint(Walk(DeadCodeElimination())),
             )
             .rewrite(mt.code)
             .join(rewrite_result)
@@ -149,17 +151,15 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             Walk(
                 Chain(
                     SquinQubitToStim(),
+                    SquinMeasureToStim(),
                     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(),
                 )
             )
             .rewrite(mt.code)
             .join(rewrite_result)
         )
+
         rewrite_result = (
             CanonicalizeIList(dialects=mt.dialects, no_raise=self.no_raise)
             .unsafe_run(mt)

src/bloqade/stim/rewrite/__init__.py

@@ -1,3 +1,4 @@
+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

src/bloqade/stim/rewrite/ifs_to_stim.py

@@ -11,9 +11,9 @@
     SQUIN_STIM_CONTROL_GATE_MAPPING,
     insert_qubit_idx_from_address,
 )
+from bloqade.analysis.measure_id import MeasureIDFrame
 from bloqade.stim.dialects.auxiliary import GetRecord
 from bloqade.analysis.measure_id.lattice import (
-    MeasureId,
     MeasureIdBool,
 )
 
@@ -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,7 +139,7 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
     def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
 
-        if not isinstance(self.measure_analysis[stmt.cond], MeasureIdBool):
+        if not isinstance(self.measure_frame.entries[stmt.cond], MeasureIdBool):
             return RewriteResult()
 
         # check that there is only qubit.Apply in the then-body,
@@ -161,12 +160,12 @@ def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
             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
 

src/bloqade/stim/rewrite/squin_measure.py

@@ -2,47 +2,16 @@
 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.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,9 +20,6 @@ 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:
@@ -70,10 +36,6 @@ def rewrite_Measure(
         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,
@@ -84,27 +46,6 @@ def rewrite_Measure(
 
         if not is_measure_result_used(measure_stmt):
             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)
 

test/stim/passes/stim_reference_programs/qubit/complex_storage_index_order.stim

@@ -0,0 +1,8 @@
+
+MZ(0.00000000) 0 1 2 3
+CZ rec[-4] 0
+MZ(0.00000000) 0 1 2 3
+CX rec[-4] 1
+MZ(0.00000000) 0 1 2 3
+CY rec[-4] 2
+CZ rec[-1] 3

test/stim/passes/stim_reference_programs/qubit/record_index_order.stim

@@ -0,0 +1,7 @@
+
+MZ(0.00000000) 0 1 2 3
+CZ rec[-4] 0
+MZ(0.00000000) 0 1 2 3
+CX rec[-4] 0
+MZ(0.00000000) 0 1 2 3
+CY rec[-12] 1

test/stim/passes/stim_reference_programs/qubit/simple_if_rewrite.stim

@@ -0,0 +1,8 @@
+
+MZ(0.00000000) 0 1 2 3
+CZ rec[-4] 0
+CX rec[-4] 1 rec[-4] 2 rec[-4] 3
+CZ rec[-4] 0 rec[-4] 1 rec[-4] 2 rec[-4] 3
+CX rec[-3] 0
+CY rec[-3] 1
+MZ(0.00000000) 0 1 2 3