detector/observable annotation dialect (#603)

Author: johnzl-777

.pre-commit-config.yaml

@@ -7,6 +7,7 @@ repos:
     - id: check-yaml
       args: ['--unsafe']
     - id: end-of-file-fixer
+      exclude: .*\.stim$
     - id: trailing-whitespace
 - repo: https://github.com/pycqa/isort
   rev: 6.0.1

src/bloqade/analysis/measure_id/impls.py

@@ -2,7 +2,7 @@
 from kirin.analysis import const
 from kirin.dialects import py, scf, func, ilist
 
-from bloqade import qubit
+from bloqade import qubit, annotate
 
 from .lattice import (
     AnyMeasureId,
@@ -46,6 +46,20 @@ def measure_qubit_list(
         return (MeasureIdTuple(data=tuple(measure_id_bools)),)
 
 
+@annotate.dialect.register(key="measure_id")
+class Annotate(interp.MethodTable):
+    @interp.impl(annotate.stmts.SetObservable)
+    @interp.impl(annotate.stmts.SetDetector)
+    def consumes_measurement_results(
+        self,
+        interp: MeasurementIDAnalysis,
+        frame: MeasureIDFrame,
+        stmt: annotate.stmts.SetObservable | annotate.stmts.SetDetector,
+    ):
+        frame.num_measures_at_stmt[stmt] = interp.measure_count
+        return (NotMeasureId(),)
+
+
 @ilist.dialect.register(key="measure_id")
 class IList(interp.MethodTable):
     @interp.impl(ilist.New)

src/bloqade/annotate/__init__.py

@@ -0,0 +1,6 @@
+from . import stmts as stmts
+from ._dialect import dialect as dialect
+from ._interface import (
+    set_detector as set_detector,
+    set_observable as set_observable,
+)

src/bloqade/annotate/_dialect.py

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

src/bloqade/annotate/_interface.py

@@ -0,0 +1,22 @@
+from typing import Any
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.types import MeasurementResult
+
+from .stmts import SetDetector, SetObservable
+from .types import Detector, Observable
+
+
+@wraps(SetDetector)
+def set_detector(
+    measurements: ilist.IList[MeasurementResult, Any] | list[MeasurementResult],
+    coordinates: ilist.IList[float | int, Any] | list[float | int],
+) -> Detector: ...
+
+
+@wraps(SetObservable)
+def set_observable(
+    measurements: ilist.IList[MeasurementResult, Any] | list[MeasurementResult],
+) -> Observable: ...

src/bloqade/annotate/stmts.py

@@ -0,0 +1,29 @@
+from kirin import ir, types as kirin_types, lowering
+from kirin.decl import info, statement
+from kirin.dialects import ilist
+
+from bloqade.types import MeasurementResultType
+from bloqade.annotate.types import DetectorType, ObservableType
+
+from ._dialect import dialect
+
+
+@statement
+class ConsumesMeasurementResults(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    measurements: ir.SSAValue = info.argument(
+        ilist.IListType[MeasurementResultType, kirin_types.Any]
+    )
+
+
+@statement(dialect=dialect)
+class SetDetector(ConsumesMeasurementResults):
+    coordinates: ir.SSAValue = info.argument(
+        type=ilist.IListType[kirin_types.Int | kirin_types.Float, kirin_types.Any]
+    )
+    result: ir.ResultValue = info.result(DetectorType)
+
+
+@statement(dialect=dialect)
+class SetObservable(ConsumesMeasurementResults):
+    result: ir.ResultValue = info.result(ObservableType)

src/bloqade/annotate/types.py

@@ -0,0 +1,13 @@
+from kirin import types
+
+
+class Detector:
+    pass
+
+
+class Observable:
+    pass
+
+
+DetectorType = types.PyClass(Detector)
+ObservableType = types.PyClass(Observable)

src/bloqade/squin/__init__.py

@@ -3,7 +3,7 @@
     noise as noise,
     analysis as analysis,
 )
-from .. import qubit as qubit
+from .. import qubit as qubit, annotate as annotate
 from ..qubit import (
     reset as reset,
     qalloc as qalloc,
@@ -12,6 +12,7 @@
     get_measurement_id as get_measurement_id,
 )
 from .groups import kernel as kernel
+from ..annotate import set_detector as set_detector, set_observable as set_observable
 from .stdlib.simple import (
     h as h,
     s as s,

src/bloqade/squin/groups.py

@@ -3,10 +3,10 @@
 from kirin.dialects import debug, ilist
 
 from . import gate, noise
-from .. import qubit
+from .. import qubit, annotate
 
 
-@ir.dialect_group(structural_no_opt.union([qubit, noise, gate, debug]))
+@ir.dialect_group(structural_no_opt.union([qubit, noise, gate, debug, annotate]))
 def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)

src/bloqade/stim/passes/squin_to_stim.py

@@ -27,7 +27,7 @@
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
 from bloqade.stim.passes.flatten import Flatten
 
-from ..rewrite.ifs_to_stim import IfToStim
+from ..rewrite import IfToStim, SetDetectorToStim, SetObservableToStim
 
 
 @dataclass
@@ -64,6 +64,8 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
         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())),
             )
             .rewrite(mt.code)