add ability to track number of measurements at statements of interest (#413)

Every time an scf.IfElse is encountered the total number of measurements
that have occurred will be saved and accessible in later rewrites

---------

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

Author: johnzl-777

src/bloqade/analysis/measure_id/analysis.py

@@ -1,20 +1,31 @@
 from typing import TypeVar
+from dataclasses import field, dataclass
 
 from kirin import ir, interp
-from kirin.analysis import Forward, const
+from kirin.analysis import ForwardExtra, const
 from kirin.analysis.forward import ForwardFrame
 
 from .lattice import MeasureId, NotMeasureId
 
 
-class MeasurementIDAnalysis(Forward[MeasureId]):
+@dataclass
+class MeasureIDFrame(ForwardFrame[MeasureId]):
+    num_measures_at_stmt: dict[ir.Statement, int] = field(default_factory=dict)
+
+
+class MeasurementIDAnalysis(ForwardExtra[MeasureIDFrame, MeasureId]):
 
     keys = ["measure_id"]
     lattice = MeasureId
     # for every kind of measurement encountered, increment this
     # then use this to generate the negative values for target rec indices
     measure_count = 0
 
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> MeasureIDFrame:
+        return MeasureIDFrame(code, has_parent_access=has_parent_access)
+
     # Still default to bottom,
     # but let constants return the softer "NoMeasureId" type from impl
     def eval_stmt_fallback(

src/bloqade/analysis/measure_id/impls.py

@@ -1,4 +1,5 @@
 from kirin import types as kirin_types, interp
+from kirin.analysis import const
 from kirin.dialects import py, scf, func, ilist
 
 from bloqade.squin import wire, qubit
@@ -10,7 +11,7 @@
     MeasureIdTuple,
     InvalidMeasureId,
 )
-from .analysis import MeasurementIDAnalysis
+from .analysis import MeasureIDFrame, MeasurementIDAnalysis
 
 ## Can't do wire right now because of
 ## unresolved RFC on return type
@@ -152,4 +153,33 @@ def invoke(
 # scf, particularly IfElse
 @scf.dialect.register(key="measure_id")
 class Scf(scf.absint.Methods):
-    pass
+
+    @interp.impl(scf.IfElse)
+    def if_else(
+        self,
+        interp_: MeasurementIDAnalysis,
+        frame: MeasureIDFrame,
+        stmt: scf.IfElse,
+    ):
+
+        frame.num_measures_at_stmt[stmt] = interp_.measure_count
+
+        # rest of the code taken directly from scf.absint.Methods base implementation
+
+        if isinstance(hint := stmt.cond.hints.get("const"), const.Value):
+            if hint.data:
+                return self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
+            else:
+                return self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
+        then_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
+        else_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
+
+        match (then_results, else_results):
+            case (interp.ReturnValue(then_value), interp.ReturnValue(else_value)):
+                return interp.ReturnValue(then_value.join(else_value))
+            case (interp.ReturnValue(then_value), _):
+                return then_results
+            case (_, interp.ReturnValue(else_value)):
+                return else_results
+            case _:
+                return interp_.join_results(then_results, else_results)

test/analysis/measure_id/test_measure_id.py

@@ -1,9 +1,17 @@
-# from kirin.types import AnyType
-# from kirin.dialects.ilist import IList
+from kirin.passes import HintConst
+from kirin.dialects import scf
 
 from bloqade.squin import op, qubit, kernel
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
-from bloqade.analysis.measure_id.lattice import MeasureIdBool, MeasureIdTuple
+from bloqade.analysis.measure_id.lattice import (
+    NotMeasureId,
+    MeasureIdBool,
+    MeasureIdTuple,
+)
+
+
+def results_at(kern, block_id, stmt_id):
+    return kern.code.body.blocks[block_id].stmts.at(stmt_id).results  # type: ignore
 
 
 def test_add():
@@ -29,3 +37,74 @@ def test():
         data=tuple([MeasureIdBool(idx=i) for i in range(1, 11)])
     )
     assert measure_id_tuples[-1] == expected_measure_id_tuple
+
+
+def test_measure_count_at_if_else():
+
+    @kernel
+    def test():
+        q = qubit.new(5)
+        qubit.apply(op.x(), q[2])
+        ms = qubit.measure(q)
+
+        if ms[1]:
+            qubit.apply(op.x(), q[0])
+
+        if ms[3]:
+            qubit.apply(op.y(), q[1])
+
+    frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
+
+    assert all(
+        isinstance(stmt, scf.IfElse) and measures_accumulated == 5
+        for stmt, measures_accumulated in frame.num_measures_at_stmt.items()
+    )
+
+
+def test_scf_cond_true():
+    @kernel
+    def test():
+        q = qubit.new(1)
+        qubit.apply(op.x(), q[2])
+
+        ms = None
+        cond = True
+        if cond:
+            ms = qubit.measure(q)
+        else:
+            ms = qubit.measure(q[0])
+
+        return ms
+
+    HintConst(dialects=test.dialects).unsafe_run(test)
+    frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
+
+    assert [frame.entries[result] for result in results_at(test, 0, 7)] == [
+        NotMeasureId(),
+        MeasureIdTuple((MeasureIdBool(idx=1),)),
+    ]
+
+
+def test_scf_cond_false():
+
+    @kernel
+    def test():
+        q = qubit.new(5)
+        qubit.apply(op.x(), q[2])
+
+        ms = None
+        cond = False
+        if cond:
+            ms = qubit.measure(q)
+        else:
+            ms = qubit.measure(q[0])
+
+        return ms
+
+    HintConst(dialects=test.dialects).unsafe_run(test)
+    frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
+
+    assert [frame.entries[result] for result in results_at(test, 0, 7)] == [
+        NotMeasureId(),
+        MeasureIdBool(idx=1),
+    ]