almost there, still a problem with invariance checking

Author: johnzl-777

src/bloqade/analysis/record/analysis.py

@@ -1,36 +1,35 @@
-from typing import TypeVar
 from dataclasses import field, dataclass
 
 from kirin import ir
-from kirin.analysis import ForwardExtra, const
+from kirin.analysis import ForwardExtra
 from kirin.analysis.forward import ForwardFrame
 
 from .lattice import Record, RecordIdx
 
 
 @dataclass
 class GlobalRecordState:
-    stack: list[RecordIdx] = field(default_factory=list)
+    buffer: list[RecordIdx] = field(default_factory=list)
 
     # assume that this RecordIdx will always be -1
-    def increment_record_idx(self) -> RecordIdx:
+    def add_record_idxs(self, num_new_records: int) -> list[RecordIdx]:
         # adjust all previous indices
-        for record_idx in self.stack:
-            record_idx.idx -= 1
-        self.stack.append(RecordIdx(-1))
-        # Return for usage
-        return self.stack[-1]
+        for record_idx in self.buffer:
+            record_idx.idx -= num_new_records
 
-    def drop_record_idx(self, record_to_drop: RecordIdx):
-        # there is a chance now that the ordering is messed up but
-        # we can now update the indices to enforce consistency.
-        # We only have to update UP to the entry that was just removed
-        # everything else maintains ordering
-        dropped_idx = record_to_drop.idx
-        self.stack.remove(record_to_drop)
-        for record_idx in self.stack:
-            if record_idx.idx < dropped_idx:
-                record_idx.idx += 1
+        # generate new indices and add them to the buffer
+        new_record_idxs = [RecordIdx(-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
+
+    """
+    Might need a free after use! You can keep the size of the list small
+    but could be a premature optimization...
+    """
+    # def drop_record_idxs(self, record_tuple: RecordTuple):
+    #    for record_idx in record_tuple.members:
+    #        self.buffer.remove(record_idx)
 
 
 @dataclass
@@ -47,7 +46,7 @@ def initialize_frame(
     ) -> RecordFrame:
         return RecordFrame(node, has_parent_access=has_parent_access)
 
-    def eval_stmt_fallback(
+    def eval_fallback(
         self, frame: RecordFrame, node: ir.Statement
     ) -> tuple[Record, ...]:
         return tuple(self.lattice.bottom() for _ in node.results)
@@ -56,17 +55,5 @@ def run_method(self, method, args: tuple[Record, ...]):
         # NOTE: we do not support dynamic calls here, thus no need to propagate method object
         return self.run_method(method.code, (self.lattice.bottom(),) + args)
 
-    T = TypeVar("T")
-
-    def get_const_value(
-        self, input_type: type[T], value: ir.SSAValue
-    ) -> type[T] | None:
-        if isinstance(hint := value.hints.get("const"), const.Value):
-            data = hint.data
-            if isinstance(data, input_type):
-                return hint.data
-
-        return None
-
     def method_self(self, method: ir.Method) -> Record:
         return self.lattice.bottom()

src/bloqade/analysis/record/impls.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 
 from kirin import types as kirin_types, interp
+from kirin.ir import PyAttr
 from kirin.dialects import py, scf, ilist
 
 from bloqade import qubit, annotate
@@ -9,9 +10,9 @@
 from .lattice import (
     AnyRecord,
     NotRecord,
-    RecordIdx,
     RecordTuple,
     InvalidRecord,
+    ConstantCarrier,
     ImmutableRecords,
 )
 from .analysis import RecordFrame, RecordAnalysis
@@ -57,10 +58,7 @@ def measure_qubit_list(
         if not isinstance(num_qubits, kirin_types.Literal):
             return (AnyRecord(),)
 
-        record_idxs = []
-        for _ in range(num_qubits.data):
-            record_idx = frame.global_record_state.increment_record_idx()
-            record_idxs.append(record_idx)
+        record_idxs = frame.global_record_state.add_record_idxs(num_qubits.data)
 
         return (RecordTuple(members=tuple(record_idxs)),)
 
@@ -70,9 +68,22 @@ class PyIndexing(interp.MethodTable):
     @interp.impl(py.GetItem)
     def getitem(self, interp: RecordAnalysis, frame: RecordFrame, stmt: py.GetItem):
 
-        idx_or_slice = interp.get_const_value((int, slice), stmt.index)
-        if idx_or_slice is None:
-            return (InvalidRecord(),)
+        # maybe_const will work fine outside of any loops because
+        # constprop will put the expected data into a hint.
+
+        # if maybeconst fails, we fall back to getting the value from the frame
+        # (note that even outside loops, the constant impl will happily
+        # capture integer/slice constants so if THAT fails, then something
+        # has truly gone wrong).
+        possible_idx_or_slice = interp.maybe_const(stmt.index, (int, slice))
+        if possible_idx_or_slice is not None:
+            idx_or_slice = possible_idx_or_slice
+        else:
+            idx_or_slice = frame.get(stmt.index)
+            if not isinstance(idx_or_slice, ConstantCarrier):
+                return (InvalidRecord(),)
+            else:
+                idx_or_slice = idx_or_slice.value
 
         obj = frame.get(stmt.obj)
         if isinstance(obj, RecordTuple):
@@ -106,36 +117,68 @@ class PyAlias(interp.MethodTable):
     @interp.impl(py.Alias)
     def alias(
         self,
-        interp: RecordAnalysis,
+        interp_: RecordAnalysis,
         frame: RecordFrame,
         stmt: py.Alias,
     ):
-        value = frame.get(stmt.value)
-        if isinstance(value, RecordIdx):
-            frame.global_record_state.drop_record_idx(value)
-        elif isinstance(value, RecordTuple):
-            for member in value.members:
-                frame.global_record_state.drop_record_idx(member)
+        input = frame.get(stmt.value)  # expect this to be a RecordTuple
 
-        return (value,)
+        # two variables share the same references in the global state
+        return (input,)
 
 
 @scf.dialect.register(key="record")
-class LoopHandling(scf.absint.Methods):
+class LoopHandling(interp.MethodTable):
     @interp.impl(scf.stmts.For)
     def for_loop(
         self, interp_: RecordAnalysis, frame: RecordFrame, stmt: scf.stmts.For
     ):
 
-        # this will contain the in-loop measure variable declared outside the loop
         loop_vars = frame.get_values(stmt.initializers)
-        # NotRecord in the beginning just lets the sink  have some value
-        loop_vars = interp_.run_ssacfg_region(frame, stmt.body, loop_vars)
 
-        # need to update the information in the frame
-        if isinstance(loop_vars, interp.ReturnValue):
-            return loop_vars
-        elif loop_vars is None:
-            loop_vars = ()
+        for _ in range(2):
+            loop_vars = interp_.frame_call_region(
+                frame, stmt, stmt.body, InvalidRecord(), *loop_vars
+            )
+
+            if loop_vars is None:
+                loop_vars = ()
+
+            elif isinstance(loop_vars, interp.ReturnValue):
+                return loop_vars
 
         return loop_vars
+
+    @interp.impl(scf.stmts.Yield)
+    def for_yield(
+        self, interp_: RecordAnalysis, frame: RecordFrame, stmt: scf.stmts.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)
+    def constant(
+        self,
+        interp_: RecordAnalysis,
+        frame: RecordFrame,
+        stmt: py.Constant,
+    ):
+        # can't use interp_.maybe_const/expect_const because it assumes the data is already
+        # there to begin with...
+        if not isinstance(stmt.value, PyAttr):
+            return (InvalidRecord(),)
+
+        expected_int_or_slice = stmt.value.data
+
+        if not isinstance(expected_int_or_slice, (int, slice)):
+            return (InvalidRecord(),)
+
+        return (ConstantCarrier(value=expected_int_or_slice),)
+
+
+# outside_frame -> create new frame with context manager COPIED from outside frame
+# the frame and the stack are separate

src/bloqade/analysis/record/lattice.py

@@ -55,6 +55,19 @@ def is_subseteq(self, other: Record) -> bool:
         return isinstance(other, NotRecord)
 
 
+# For now I only care about propagating constant integers or slices,
+# things that can be used as indices to list of measurements
+@final
+@dataclass
+class ConstantCarrier(Record):
+    value: int | slice
+
+    def is_subseteq(self, other: Record) -> bool:
+        if isinstance(other, ConstantCarrier):
+            return self.value == other.value
+        return False
+
+
 @final
 @dataclass
 class RecordIdx(Record):

src/bloqade/stim/dialects/cf/__init__.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("stim.cf")

src/bloqade/stim/dialects/cf/_dialect.py

@@ -22,7 +22,7 @@ class StimSimplifyIfs(Pass):
     def unsafe_run(self, mt: ir.Method):
 
         result = Chain(
-            Walk(UnusedYield()),
+            Walk(UnusedYield()),  # this is being too aggressive, need to file an issue
             Walk(StimLiftThenBody()),
             # remove yields (if possible), then lift out as much stuff as possible
             Walk(DeadCodeElimination()),

src/bloqade/stim/dialects/cf/stmts.py

@@ -81,10 +81,12 @@ class SoftFlatten(Pass):
 
     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)
+
+        # DO NOT USE FOR NOW, TrimUnusedYield call messes up loop structure
+        # 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