remove more record analysis debug prints and move everything into measurement_id

Author: johnzl-777

src/bloqade/analysis/measure_id/analysis.py

@@ -4,15 +4,15 @@
 from kirin.analysis import ForwardExtra
 from kirin.analysis.forward import ForwardFrame
 
-from .lattice import MeasureId, NotMeasureId, KnownMeasureId
+from .lattice import MeasureId, NotMeasureId, KnownMeasureId, MeasureIdTuple
 
 
 @dataclass
 class GlobalRecordState:
     buffer: list[KnownMeasureId] = field(default_factory=list)
 
     # assume that this KnownMeasureId will always be -1
-    def add_record_idxs(self, num_new_records: int) -> list[KnownMeasureId]:
+    def add_record_idxs(self, num_new_records: int) -> MeasureIdTuple:
         # adjust all previous indices
         for record_idx in self.buffer:
             record_idx.idx -= num_new_records
@@ -21,12 +21,33 @@ def add_record_idxs(self, num_new_records: int) -> list[KnownMeasureId]:
         new_record_idxs = [KnownMeasureId(-i) for i in range(num_new_records, 0, -1)]
         self.buffer += new_record_idxs
         # Return for usage, idxs linked to the global state
-        return new_record_idxs
+        return MeasureIdTuple(data=tuple(new_record_idxs))
+
+    # Need for loop invariance, especially when you
+    # run the loop twice "behind the scenes". Then
+    # it isn't sufficient to just have two
+    # copies of a lattice element point to one entry on the
+    # buffer
+    def clone_record_idxs(self, measure_id_tuple: MeasureIdTuple) -> MeasureIdTuple:
+        cloned_members = []
+        for known_measure_id in measure_id_tuple.data:
+            assert isinstance(known_measure_id, KnownMeasureId)
+            cloned_known_measure_id = KnownMeasureId(known_measure_id.idx)
+            # put into the global buffer but also
+            # return an analysis-facing copy
+            self.buffer.append(cloned_known_measure_id)
+            cloned_members.append(cloned_known_measure_id)
+        return MeasureIdTuple(data=tuple(cloned_members))
+
+    def offset_existing_records(self, offset: int):
+        for record_idx in self.buffer:
+            record_idx.idx -= offset
 
 
 @dataclass
 class MeasureIDFrame(ForwardFrame[MeasureId]):
     global_record_state: GlobalRecordState = field(default_factory=GlobalRecordState)
+    measure_count_offset: int = 0
 
 
 class MeasurementIDAnalysis(ForwardExtra[MeasureIDFrame, MeasureId]):

src/bloqade/analysis/measure_id/impls.py

@@ -38,9 +38,15 @@ def measure_qubit_list(
         if not isinstance(num_qubits, kirin_types.Literal):
             return (AnyMeasureId(),)
 
-        record_idxs = frame.global_record_state.add_record_idxs(num_qubits.data)
+        # increment the parent frame measure count offset.
+        # Loop analysis relies on local state tracking
+        # so we use this data after exiting a loop to
+        # readjust the previous global measure count.
+        frame.measure_count_offset += num_qubits.data
 
-        return (MeasureIdTuple(data=tuple(record_idxs)),)
+        measure_id_tuple = frame.global_record_state.add_record_idxs(num_qubits.data)
+
+        return (measure_id_tuple,)
 
 
 @annotate.dialect.register(key="measure_id")
@@ -130,11 +136,16 @@ def getitem(
 class PyAssign(interp.MethodTable):
     @interp.impl(py.Alias)
     def alias(
-        self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.assign.Alias
+        self,
+        interp: MeasurementIDAnalysis,
+        frame: MeasureIDFrame,
+        stmt: py.assign.Alias,
     ):
 
         input = frame.get(stmt.value)
-        return (input,)
+
+        new_input = frame.global_record_state.clone_record_idxs(input)
+        return (new_input,)
 
 
 @py.binop.dialect.register(key="measure_id")
@@ -183,9 +194,11 @@ def for_loop(
         # You go through the loops twice to verify the loop invariant.
         # we need to freeze the frame entries right after exiting the loop
 
+        local_state = deepcopy(frame.global_record_state)
+
         first_loop_frame = MeasureIDFrame(
             stmt,
-            global_record_state=frame.global_record_state,
+            global_record_state=local_state,
             parent=frame,
             has_parent_access=True,
         )
@@ -206,7 +219,7 @@ def for_loop(
 
         second_loop_frame = MeasureIDFrame(
             stmt,
-            global_record_state=frame.global_record_state,
+            global_record_state=local_state,
             parent=frame,
             has_parent_access=True,
         )
@@ -231,6 +244,9 @@ def for_loop(
             unified_frame_buffer[ssa_val] = verified_latticed_element
 
         frame.entries.update(unified_frame_buffer)
+        frame.global_record_state.offset_existing_records(
+            first_loop_frame.measure_count_offset
+        )
 
         if captured_first_loop_vars is None or second_loop_vars is None:
             return ()
@@ -241,6 +257,20 @@ def for_loop(
         ):
             joined_loop_vars.append(first_loop_var.join(second_loop_var))
 
+        # TrimYield is currently disabled meaning that the same RecordIdx
+        # can get copied into the parent frame twice! As a result
+        # we need to be careful to only add unique RecordIdx entries
+        witnessed_record_idxs = set()
+        for var in joined_loop_vars:
+            if isinstance(var, MeasureIdTuple):
+                for member in var.data:
+                    if (
+                        isinstance(member, KnownMeasureId)
+                        and member.idx not in witnessed_record_idxs
+                    ):
+                        witnessed_record_idxs.add(member.idx)
+                        frame.global_record_state.buffer.append(member)
+
         return tuple(joined_loop_vars)
 
     @interp.impl(scf.stmts.Yield)

src/bloqade/analysis/record/analysis.py

@@ -42,9 +42,6 @@ def clone_record_idxs(self, record_tuple: RecordTuple, id: int) -> RecordTuple:
     def offset_existing_records(self, offset: int):
         for record_idx in self.buffer:
             record_idx.idx -= offset
-            print("offset is now:", offset)
-            print("The record idx is now:", record_idx.idx)
-        # print the record_idx after offsetting
 
     """
     Might need a free after use! You can keep the size of the list small

src/bloqade/analysis/record/impls.py

@@ -245,7 +245,6 @@ def for_loop_double_pass(
     def for_yield(
         self, interp_: RecordAnalysis, frame: RecordFrame, stmt: scf.stmts.Yield
     ):
-        print("yield encountered, yielding values:", frame.get_values(stmt.values))
         return interp.YieldValue(frame.get_values(stmt.values))