get initial implementation of new intermediate dialect approach

Author: johnzl-777

src/bloqade/analysis/measure_id/__init__.py

@@ -1,4 +1,5 @@
 from . import impls as impls
+from .lattice import RecId as RecId
 from .analysis import (
     MeasureIDFrame as MeasureIDFrame,
     MeasurementIDAnalysis as MeasurementIDAnalysis,

src/bloqade/analysis/measure_id/analysis.py

@@ -1,6 +1,7 @@
 from dataclasses import field, dataclass
 
 from kirin import ir
+from kirin.interp import StatementResult
 from kirin.analysis import ForwardExtra
 from typing_extensions import Self
 from kirin.analysis.forward import ForwardFrame
@@ -17,8 +18,6 @@ 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
     detector_count = 0
     observable_count = 0
@@ -34,12 +33,18 @@ def initialize_frame(
     ) -> MeasureIDFrame:
         return MeasureIDFrame(node, has_parent_access=has_parent_access)
 
-    # Still default to bottom,
-    # but let constants return the softer "NoMeasureId" type from impl
     def eval_fallback(
         self, frame: ForwardFrame[MeasureId], node: ir.Statement
     ) -> tuple[MeasureId, ...]:
         return tuple(NotMeasureId() for _ in node.results)
 
+    def frame_eval(
+        self, frame: MeasureIDFrame, node: ir.Statement
+    ) -> StatementResult[MeasureId]:
+        method = self.lookup_registry(frame, node)
+        if method is not None:
+            return method(self, frame, node)
+        return self.eval_fallback(frame, node)
+
     def method_self(self, method: ir.Method) -> MeasureId:
         return self.lattice.bottom()

src/bloqade/analysis/measure_id/impls.py

@@ -4,9 +4,15 @@
 
 from bloqade import qubit
 from bloqade.decoders.dialects import annotate
+from bloqade.record_idx_helper import (
+    GetRecIdxFromPredicate,
+    GetRecIdxFromMeasurement,
+    dialect as record_idx_helper_dialect,
+)
 from bloqade.gemini.logical.dialects import operations
 
 from .lattice import (
+    RecId,
     MeasureId,
     Predicate,
     DetectorId,
@@ -316,3 +322,38 @@ def if_else(
                 return else_results
             case _:
                 return interp_.join_results(then_results, else_results)
+
+
+@record_idx_helper_dialect.register(key="measure_id")
+class RecordIdxHelperAnalysis(interp.MethodTable):
+
+    @interp.impl(GetRecIdxFromMeasurement)
+    def get_rec_idx_from_measurement(
+        self,
+        interp_: MeasurementIDAnalysis,
+        frame: MeasureIDFrame,
+        stmt: GetRecIdxFromMeasurement,
+    ):
+        measurement_id = frame.get(stmt.measurement)
+        if not isinstance(measurement_id, (RawMeasureId, MeasureIdBool)):
+            return (InvalidMeasureId(),)
+        computed_idx = (measurement_id.idx - 1) - interp_.measure_count
+        predicate = (
+            measurement_id.predicate
+            if isinstance(measurement_id, MeasureIdBool)
+            else None
+        )
+        return (RecId(idx=computed_idx, predicate=predicate),)
+
+    @interp.impl(GetRecIdxFromPredicate)
+    def get_rec_idx_from_predicate(
+        self,
+        interp_: MeasurementIDAnalysis,
+        frame: MeasureIDFrame,
+        stmt: GetRecIdxFromPredicate,
+    ):
+        measurement_id = frame.get(stmt.predicate_result)
+        if not isinstance(measurement_id, MeasureIdBool):
+            return (InvalidMeasureId(),)
+        computed_idx = (measurement_id.idx - 1) - interp_.measure_count
+        return (RecId(idx=computed_idx, predicate=measurement_id.predicate),)

src/bloqade/analysis/measure_id/lattice.py

@@ -91,6 +91,13 @@ class MeasureIdBool(ConcreteMeasureId):
     predicate: Predicate
 
 
+@final
+@dataclass
+class RecId(ConcreteMeasureId):
+    idx: int
+    predicate: Predicate | None
+
+
 @final
 @dataclass
 class DetectorId(MeasureId):

src/bloqade/record_idx_helper/__init__.py

@@ -0,0 +1,5 @@
+from .stmts import (
+    GetRecIdxFromPredicate as GetRecIdxFromPredicate,
+    GetRecIdxFromMeasurement as GetRecIdxFromMeasurement,
+)
+from ._dialect import dialect as dialect

src/bloqade/record_idx_helper/_dialect.py

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

src/bloqade/record_idx_helper/stmts.py

@@ -0,0 +1,22 @@
+from kirin import ir, types
+from kirin.decl import info, statement
+
+from bloqade.types import MeasurementResultType
+
+from ._dialect import dialect
+
+
+@statement(dialect=dialect)
+class GetRecIdxFromMeasurement(ir.Statement):
+    name = "get_rec_idx_from_measurement"
+    traits = frozenset({ir.Pure()})
+    measurement: ir.SSAValue = info.argument(type=MeasurementResultType)
+    result: ir.ResultValue = info.result(type=types.Int)
+
+
+@statement(dialect=dialect)
+class GetRecIdxFromPredicate(ir.Statement):
+    name = "get_rec_idx_from_predicate"
+    traits = frozenset({ir.Pure()})
+    predicate_result: ir.SSAValue = info.argument(type=types.Bool)
+    result: ir.ResultValue = info.result(type=types.Int)

src/bloqade/stim/passes/squin_to_stim.py

@@ -10,12 +10,17 @@
 from kirin.ir.method import Method
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
+from kirin.passes.hint_const import HintConst
 
 from bloqade.stim.rewrite import (
+    IfToStimPartial,
     PyConstantToStim,
+    ResolveGetRecIdx,
     SquinNoiseToStim,
     SquinQubitToStim,
+    SetDetectorPartial,
     SquinMeasureToStim,
+    SetObservablePartial,
 )
 from bloqade.squin.rewrite import (
     SquinU3ToClifford,
@@ -24,82 +29,79 @@
 )
 from bloqade.rewrite.passes import CanonicalizeIList
 from bloqade.analysis.address import AddressAnalysis
+from bloqade.record_idx_helper import dialect as record_idx_helper_dialect
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
 from bloqade.stim.passes.flatten import Flatten
 
-from ..rewrite import IfToStim, SetDetectorToStim, SetObservableToStim
-
 
 @dataclass
 class SquinToStimPass(Pass):
 
     def unsafe_run(self, mt: Method) -> RewriteResult:
 
-        # inline aggressively:
         rewrite_result = Flatten(dialects=mt.dialects, no_raise=self.no_raise).fixpoint(
             mt
         )
 
-        # after this the program should be in a state where it is analyzable
-        # -------------------------------------------------------------------
-
-        mia = MeasurementIDAnalysis(dialects=mt.dialects)
-        meas_analysis_frame, _ = mia.run(mt)
-
         aa = AddressAnalysis(dialects=mt.dialects)
         address_analysis_frame, _ = aa.run(mt)
 
-        # wrap the address analysis result
         rewrite_result = (
             Walk(WrapAddressAnalysis(address_analysis=address_analysis_frame.entries))
             .rewrite(mt.code)
             .join(rewrite_result)
         )
 
-        # 2. rewrite
-        ## Invoke DCE afterwards to eliminate any GetItems
-        ## that are no longer being used. This allows for
-        ## SquinMeasureToStim to safely eliminate
-        ## unused measure statements.
+        # --- partial rewrite (before analysis) ---
         rewrite_result = (
-            Chain(
-                Walk(IfToStim(measure_frame=meas_analysis_frame)),
-                Walk(SetDetectorToStim(measure_id_frame=meas_analysis_frame)),
-                Walk(SetObservableToStim(measure_id_frame=meas_analysis_frame)),
-                Fixpoint(Walk(DeadCodeElimination())),
+            Walk(
+                Chain(
+                    SetDetectorPartial(),
+                    SetObservablePartial(),
+                    IfToStimPartial(),
+                )
             )
             .rewrite(mt.code)
             .join(rewrite_result)
         )
 
-        # Rewrite the noise statements first.
         rewrite_result = Walk(SquinNoiseToStim()).rewrite(mt.code).join(rewrite_result)
-
-        # Wrap Rewrite + SquinToStim can happen w/ standard walk
         rewrite_result = Walk(SquinU3ToClifford()).rewrite(mt.code).join(rewrite_result)
+        rewrite_result = Walk(SquinQubitToStim()).rewrite(mt.code).join(rewrite_result)
 
+        # --- analysis (produces RecId for GetRecIdxFromMeasurement / GetRecIdxFromPredicate) ---
+        analysis_dialects = mt.dialects.add(record_idx_helper_dialect)
         rewrite_result = (
-            Walk(
-                Chain(
-                    SquinQubitToStim(),
-                    SquinMeasureToStim(),
-                )
+            HintConst(analysis_dialects, no_raise=self.no_raise)
+            .unsafe_run(mt)
+            .join(rewrite_result)
+        )
+        mia = MeasurementIDAnalysis(dialects=analysis_dialects)
+        meas_analysis_frame, _ = mia.run(mt)
+
+        # --- post-analysis: resolve helper stmts into direct integer constants ---
+        rewrite_result = (
+            Chain(
+                Walk(ResolveGetRecIdx(measure_id_frame=meas_analysis_frame)),
+                Fixpoint(Walk(DeadCodeElimination())),
             )
             .rewrite(mt.code)
             .join(rewrite_result)
         )
 
+        # --- rewrite measures (must stay until after analysis) ---
+        rewrite_result = (
+            Walk(SquinMeasureToStim()).rewrite(mt.code).join(rewrite_result)
+        )
+
         rewrite_result = (
             CanonicalizeIList(dialects=mt.dialects, no_raise=self.no_raise)
             .unsafe_run(mt)
             .join(rewrite_result)
         )
 
-        # Convert all PyConsts to Stim Constants
         rewrite_result = Walk(PyConstantToStim()).rewrite(mt.code).join(rewrite_result)
 
-        # clear up leftover stmts
-        # - remove any squin.qalloc that's left around
         rewrite_result = (
             Fixpoint(
                 Walk(

src/bloqade/stim/rewrite/__init__.py

@@ -2,6 +2,8 @@
 from .squin_noise import SquinNoiseToStim as SquinNoiseToStim
 from .qubit_to_stim import SquinQubitToStim as SquinQubitToStim
 from .squin_measure import SquinMeasureToStim as SquinMeasureToStim
+from .ifs_to_stim_partial import IfToStimPartial as IfToStimPartial
 from .py_constant_to_stim import PyConstantToStim as PyConstantToStim
-from .set_detector_to_stim import SetDetectorToStim as SetDetectorToStim
-from .set_observable_to_stim import SetObservableToStim as SetObservableToStim
+from .resolve_get_rec_idx import ResolveGetRecIdx as ResolveGetRecIdx
+from .set_detector_partial import SetDetectorPartial as SetDetectorPartial
+from .set_observable_partial import SetObservablePartial as SetObservablePartial

src/bloqade/stim/rewrite/ifs_to_stim_partial.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.dialects import scf
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin import gate
+from bloqade.squin.rewrite import AddressAttribute
+from bloqade.record_idx_helper import GetRecIdxFromPredicate
+from bloqade.stim.rewrite.util import insert_qubit_idx_from_address
+from bloqade.stim.dialects.gate import CX as stim_CX, CY as stim_CY, CZ as stim_CZ
+from bloqade.stim.dialects.auxiliary import GetRecord
+from bloqade.stim.rewrite.ifs_to_stim import IfElseSimplification
+
+PAULI_TO_CONTROLLED = {
+    gate.stmts.X: stim_CX,
+    gate.stmts.Y: stim_CY,
+    gate.stmts.Z: stim_CZ,
+}
+
+
+@dataclass
+class IfToStimPartial(IfElseSimplification, RewriteRule):
+    """Rewrite measurement-conditioned IfElse using GetRecIdxFromPredicate.
+
+    Accepts the Bool condition directly (result of IsOne/IsZero) and creates
+    GetRecIdxFromPredicate -> GetRecord -> CX/CY/CZ. If the body contains
+    multiple Pauli gates, splits into multiple controlled gates sharing the
+    same GetRecord result.
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        match node:
+            case scf.IfElse():
+                return self.rewrite_IfElse(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
+        if not self.is_rewriteable(stmt):
+            return RewriteResult()
+
+        *body_stmts, _ = stmt.then_body.stmts()
+        if not body_stmts:
+            return RewriteResult()
+
+        idx_from_predicate_calc = GetRecIdxFromPredicate(predicate_result=stmt.cond)
+        idx_from_predicate_calc.insert_before(stmt)
+
+        get_record_stmt = GetRecord(id=idx_from_predicate_calc.result)
+        get_record_stmt.insert_before(stmt)
+
+        for body_stmt in body_stmts:
+            address_attr = body_stmt.qubits.hints.get("address")
+            if address_attr is None:
+                return RewriteResult()
+            assert isinstance(address_attr, AddressAttribute)
+
+            qubit_idx_ssas = insert_qubit_idx_from_address(
+                address=address_attr, stmt_to_insert_before=stmt
+            )
+            if qubit_idx_ssas is None:
+                return RewriteResult()
+
+            stim_gate_cls = PAULI_TO_CONTROLLED[type(body_stmt)]
+            stim_stmt = stim_gate_cls(
+                targets=tuple(qubit_idx_ssas),
+                controls=(get_record_stmt.result,) * len(qubit_idx_ssas),
+            )
+            stim_stmt.insert_before(stmt)
+
+        stmt.delete()
+
+        return RewriteResult(has_done_something=True)