still need to find a solution to proper IfElse handling

Author: johnzl-777

src/bloqade/analysis/measure_id/impls.py

@@ -14,7 +14,6 @@
     MeasureIdTuple,
     ConstantCarrier,
     InvalidMeasureId,
-    ImmutableMeasureIds,
     PredicatedMeasureId,
 )
 from .analysis import MeasureIDFrame, MeasurementIDAnalysis
@@ -80,7 +79,7 @@ def measurement_predicate(
             PredicatedMeasureId(measure_id.idx, predicate)
             for measure_id in original_measure_id_tuple.data
         ]
-        return (ImmutableMeasureIds(data=tuple(predicate_measure_ids)),)
+        return (MeasureIdTuple(data=tuple(predicate_measure_ids)),)
 
 
 @gemini.logical.dialect.register(key="measure_id")
@@ -107,7 +106,7 @@ def terminal_measurement(
         # Immutable usually desired for stim generation
         # but we can reuse it here to indicate
         # the measurement ids should not change anymore.
-        return (ImmutableMeasureIds(data=tuple(measure_id_bools)),)
+        return (MeasureIdTuple(data=tuple(measure_id_bools), immutable=True),)
 
 
 @annotate.dialect.register(key="measure_id")
@@ -124,17 +123,15 @@ def consumes_measurements(
 
         if not (
             isinstance(measure_id_tuple_at_stmt, MeasureIdTuple)
-            and kirin_types.is_tuple_of(
-                measure_id_tuple_at_stmt.data, PredicatedMeasureId
-            )
+            and kirin_types.is_tuple_of(measure_id_tuple_at_stmt.data, RawMeasureId)
         ):
             return (InvalidMeasureId(),)
 
         final_record_idxs = [
             deepcopy(record_idx) for record_idx in measure_id_tuple_at_stmt.data
         ]
 
-        return (ImmutableMeasureIds(data=tuple(final_record_idxs)),)
+        return (MeasureIdTuple(data=tuple(final_record_idxs), immutable=True),)
 
 
 @ilist.dialect.register(key="measure_id")
@@ -150,7 +147,7 @@ def new_ilist(
         stmt: ilist.New,
     ):
 
-        return (MeasureIdTuple(frame.get_values(stmt.values)),)
+        return (MeasureIdTuple(data=frame.get_values(stmt.values)),)
 
 
 @py.tuple.dialect.register(key="measure_id")
@@ -345,6 +342,21 @@ def for_yield(
     ):
         return interp.YieldValue(frame.get_values(stmt.values))
 
+    @interp.impl(scf.stmts.IfElse)
+    def if_else(
+        self,
+        interp_: MeasurementIDAnalysis,
+        frame: MeasureIDFrame,
+        stmt: scf.stmts.IfElse,
+    ):
+        cond_measure_id = frame.get(stmt.cond)
+        assert type(cond_measure_id) is PredicatedMeasureId
+        detached_cond_measure_id = PredicatedMeasureId(
+            idx=deepcopy(cond_measure_id.idx), predicate=cond_measure_id.predicate
+        )
+        # remove underlying reference to the frame
+        frame.set(stmt.cond, detached_cond_measure_id)
+
 
 @py.dialect.register(key="measure_id")
 class ConstantForwarding(interp.MethodTable):

src/bloqade/analysis/measure_id/lattice.py

@@ -93,27 +93,14 @@ def is_subseteq(self, other: MeasureId) -> bool:
 @dataclass
 class MeasureIdTuple(MeasureId):
     data: tuple[MeasureId, ...]
+    immutable: bool = False
 
     def is_subseteq(self, other: MeasureId) -> bool:
         if isinstance(other, MeasureIdTuple):
             return all(a.is_subseteq(b) for a, b in zip(self.data, other.data))
         return False
 
 
-@final
-@dataclass
-class ImmutableMeasureIds(MeasureId):
-    # SetDetector happily consumes RawMeasureIds, but
-    # for scf.IfElse rewrite with predicates I need to allow
-    # PredicatedMeasureIds as well.
-    data: tuple[PredicatedMeasureId | RawMeasureId, ...]
-
-    def is_subseteq(self, other: MeasureId) -> bool:
-        if isinstance(other, ImmutableMeasureIds):
-            return all(a.is_subseteq(b) for a, b in zip(self.data, other.data))
-        return False
-
-
 # For now I only care about propagating constant integers or slices,
 # things that can be used as indices to list of measurements
 @final

src/bloqade/stim/passes/soft_flatten.py

@@ -81,12 +81,10 @@ class SoftFlatten(Pass):
 
     def __post_init__(self):
         self.unroll = AggressiveUnroll(self.dialects, no_raise=self.no_raise)
-
-        # DO NOT USE FOR NOW, TrimUnusedYield call messes up loop structure
-        # self.simplify_if = StimSimplifyIfs(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.simplify_if(mt).join(rewrite_result)
         rewrite_result = self.unroll(mt).join(rewrite_result)
         return rewrite_result

src/bloqade/stim/rewrite/get_record_util.py

@@ -2,7 +2,7 @@
 from kirin.dialects import py
 
 from bloqade.stim.dialects import auxiliary
-from bloqade.analysis.measure_id.lattice import MeasureIdTuple, PredicatedMeasureId
+from bloqade.analysis.measure_id.lattice import RawMeasureId, MeasureIdTuple
 
 
 def insert_get_records(node: ir.Statement, measure_id_tuple: MeasureIdTuple):
@@ -11,7 +11,7 @@ def insert_get_records(node: ir.Statement, measure_id_tuple: MeasureIdTuple):
     """
     get_record_ssas = []
     for known_measure_id in measure_id_tuple.data:
-        assert isinstance(known_measure_id, PredicatedMeasureId)
+        assert isinstance(known_measure_id, RawMeasureId)
         idx_stmt = py.constant.Constant(known_measure_id.idx)
         idx_stmt.insert_before(node)
         get_record_stmt = auxiliary.GetRecord(idx_stmt.result)

test/analysis/measure_id/test_measure_id.py

@@ -4,13 +4,15 @@
 from bloqade import squin, gemini
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
 from bloqade.stim.passes.flatten import Flatten
+
+# from bloqade.stim.passes.soft_flatten import SoftFlatten
+from bloqade.stim.passes.squin_to_stim import SquinToStimPass
 from bloqade.analysis.measure_id.lattice import (
     Predicate,
     NotMeasureId,
     RawMeasureId,
     MeasureIdTuple,
     InvalidMeasureId,
-    ImmutableMeasureIds,
     PredicatedMeasureId,
 )
 
@@ -302,29 +304,29 @@ def test():
         test, ("is_zero_bools", "is_one_bools", "is_lost_bools")
     )
 
-    expected_is_zero_bools = ImmutableMeasureIds(
+    expected_is_zero_bools = MeasureIdTuple(
         data=tuple(
             [
                 PredicatedMeasureId(idx=i, predicate=Predicate.IS_ZERO)
                 for i in range(-3, 0)
             ]
-        )
+        ),
     )
-    expected_is_one_bools = ImmutableMeasureIds(
+    expected_is_one_bools = MeasureIdTuple(
         data=tuple(
             [
                 PredicatedMeasureId(idx=i, predicate=Predicate.IS_ONE)
                 for i in range(-3, 0)
             ]
-        )
+        ),
     )
-    expected_is_lost_bools = ImmutableMeasureIds(
+    expected_is_lost_bools = MeasureIdTuple(
         data=tuple(
             [
                 PredicatedMeasureId(idx=i, predicate=Predicate.IS_LOST)
                 for i in range(-3, 0)
             ]
-        )
+        ),
     )
 
     assert frame.get(results["is_zero_bools"]) == expected_is_zero_bools
@@ -347,9 +349,38 @@ def tm_logical_kernel():
     # basically a container of InvalidMeasureIds from the qubits that get allocated
     tm_logical_kernel.print(analysis=frame.entries)
     analysis_results = [
-        val for val in frame.entries.values() if isinstance(val, ImmutableMeasureIds)
+        val for val in frame.entries.values() if isinstance(val, MeasureIdTuple)
     ]
-    expected_result = ImmutableMeasureIds(
-        data=tuple([RawMeasureId(idx=-i) for i in range(1, 4)])
+    expected_result = MeasureIdTuple(
+        data=tuple([RawMeasureId(idx=-i) for i in range(1, 4)]),
+        immutable=True,
     )
     assert expected_result in analysis_results
+
+
+def test_if_else_happy_path():
+
+    @squin.kernel
+    def test():
+        qs = squin.qalloc(3)
+        ms = squin.broadcast.measure(qs)
+        # predicate
+        pred_ms = squin.broadcast.is_one(ms)
+        squin.broadcast.measure(qs)
+        squin.broadcast.measure(qs)
+        if pred_ms[0]:
+            squin.x(qs[1])
+
+        return
+
+    # Flatten(test.dialects).fixpoint(test)
+    # SoftFlatten(test.dialects).fixpoint(test)
+    test.print()
+    SquinToStimPass(test.dialects)(test)
+    test.print()
+    # test.print()
+    # frame, _ = MeasurementIDAnalysis(test.dialects).run(test)
+    # test.print(analysis=frame.entries)
+
+
+test_if_else_happy_path()

test/stim/passes/test_annotation_to_stim.py

@@ -148,10 +148,15 @@ def main():
 
         return
 
+    main.print()
     SquinToStimPass(main.dialects, no_raise=True)(main)
+    main.print()
     assert any(isinstance(stmt, scf.IfElse) for stmt in main.code.regions[0].stmts())
 
 
+test_missing_predicate()
+
+
 def test_incorrect_predicate():
 
     # You can only rewrite squin.is_one(...) predicates to