loop invariance support complete

Author: johnzl-777

src/bloqade/analysis/record/analysis.py

@@ -23,6 +23,19 @@ def add_record_idxs(self, num_new_records: int) -> list[RecordIdx]:
         # Return for usage, idxs linked to the global state
         return new_record_idxs
 
+    """
+    def clone_record_idxs(self, record_tuple: RecordTuple) -> RecordTuple:
+        cloned_members = []
+        for record_idx in record_tuple.members:
+            cloned_record_idx = RecordIdx(record_idx.idx)
+            # put into the global buffer but also
+            # return an analysis-facing copy
+            self.buffer.append(cloned_record_idx)
+            cloned_members.append(cloned_record_idx)
+
+        return RecordTuple(members=tuple(cloned_members))
+    """
+
     """
     Might need a free after use! You can keep the size of the list small
     but could be a premature optimization...

src/bloqade/analysis/record/impls.py

@@ -10,6 +10,7 @@
 from .lattice import (
     AnyRecord,
     NotRecord,
+    RecordIdx,
     RecordTuple,
     InvalidRecord,
     ConstantCarrier,
@@ -31,6 +32,12 @@ def consumes_measurements(
         # Get the measurement results being consumed
         record_tuple_at_stmt = frame.get(stmt.measurements)
 
+        if not (
+            isinstance(record_tuple_at_stmt, RecordTuple)
+            and kirin_types.is_tuple_of(record_tuple_at_stmt.members, RecordIdx)
+        ):
+            return (InvalidRecord(),)
+
         final_record_idxs = [
             deepcopy(record_idx) for record_idx in record_tuple_at_stmt.members
         ]
@@ -122,7 +129,7 @@ def alias(
         stmt: py.Alias,
     ):
         input = frame.get(stmt.value)  # expect this to be a RecordTuple
-
+        # frame.global_record_state.clone_record_idxs(input)
         # two variables share the same references in the global state
         return (input,)
 
@@ -134,20 +141,70 @@ def for_loop(
         self, interp_: RecordAnalysis, frame: RecordFrame, stmt: scf.stmts.For
     ):
 
-        loop_vars = frame.get_values(stmt.initializers)
-
-        for _ in range(2):
-            loop_vars = interp_.frame_call_region(
-                frame, stmt, stmt.body, InvalidRecord(), *loop_vars
+        init_loop_vars = frame.get_values(stmt.initializers)
+
+        # You go through the loops twice to verify the loop invariant.
+        # we need to freeze the frame entries right after exiting the loop
+
+        first_loop_frame = RecordFrame(
+            stmt,
+            global_record_state=frame.global_record_state,
+            parent=frame,
+            has_parent_access=True,
+        )
+        first_loop_vars = interp_.frame_call_region(
+            first_loop_frame, stmt, stmt.body, InvalidRecord(), *init_loop_vars
+        )
+
+        if first_loop_vars is None:
+            first_loop_vars = ()
+        elif isinstance(first_loop_vars, interp.ReturnValue):
+            return first_loop_vars
+
+        captured_first_loop_entries = {}
+        captured_first_loop_vars = deepcopy(first_loop_vars)
+
+        for ssa_val, lattice_element in first_loop_frame.entries.items():
+            captured_first_loop_entries[ssa_val] = deepcopy(lattice_element)
+
+        second_loop_frame = RecordFrame(
+            stmt,
+            global_record_state=frame.global_record_state,
+            parent=frame,
+            has_parent_access=True,
+        )
+        second_loop_vars = interp_.frame_call_region(
+            second_loop_frame, stmt, stmt.body, InvalidRecord(), *first_loop_vars
+        )
+
+        if second_loop_vars is None:
+            second_loop_vars = ()
+        elif isinstance(second_loop_vars, interp.ReturnValue):
+            return second_loop_vars
+
+        # take the entries in the first and second loops
+        # update the parent frame
+
+        unified_frame_buffer = {}
+        for ssa_val, lattice_element in captured_first_loop_entries.items():
+            verified_latticed_element = second_loop_frame.entries[ssa_val].join(
+                lattice_element
             )
+            # print(f"Joining {lattice_element} and {second_loop_frame.entries[ssa_val]} to get {verified_latticed_element}")
+            unified_frame_buffer[ssa_val] = verified_latticed_element
+
+        frame.entries.update(unified_frame_buffer)
 
-            if loop_vars is None:
-                loop_vars = ()
+        if captured_first_loop_vars is None or second_loop_vars is None:
+            return ()
 
-            elif isinstance(loop_vars, interp.ReturnValue):
-                return loop_vars
+        joined_loop_vars = []
+        for first_loop_var, second_loop_var in zip(
+            captured_first_loop_vars, second_loop_vars
+        ):
+            joined_loop_vars.append(first_loop_var.join(second_loop_var))
 
-        return loop_vars
+        return tuple(joined_loop_vars)
 
     @interp.impl(scf.stmts.Yield)
     def for_yield(
@@ -156,8 +213,6 @@ def for_yield(
         return interp.YieldValue(frame.get_values(stmt.values))
 
 
-# Only carry about carrying integers for now because
-# the current issue is that
 @py.dialect.register(key="record")
 class ConstantForwarding(interp.MethodTable):
     @interp.impl(py.Constant)

test/analysis/record/test_record_analysis.py

@@ -162,13 +162,11 @@ def test():
             squin.annotate.set_detector([prev_ms[0], curr_ms[0]], coordinates=[0, 0])
             squin.annotate.set_detector([prev_ms[1], curr_ms[1]], coordinates=[0, 1])
 
-        squin.annotate.set_detector(curr_ms, coordinates=[0, 0])
         data_ms = squin.broadcast.measure(data_qs)
 
         squin.set_detector([data_ms[0], data_ms[1], curr_ms[0]], coordinates=[2, 0])
         squin.set_detector([data_ms[2], data_ms[1], curr_ms[1]], coordinates=[2, 1])
 
-    test.print()
     SoftFlatten(dialects=test.dialects).fixpoint(test)
     test.print()
     frame, _ = RecordAnalysis(dialects=test.dialects).run(test)