Merge branch 'main' into john/measurement-result-predicate

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 (
     Predicate,
@@ -83,6 +83,20 @@ def measurement_predicate(
         return (MeasureIdTuple(data=tuple(predicate_measure_ids)),)
 
 
+@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/rewrite/passes/aggressive_unroll.py

@@ -42,6 +42,7 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             ilist.rewrite.InlineGetItem(),
             ilist.rewrite.FlattenAdd(),
             ilist.rewrite.HintLen(),
+            DeadCodeElimination(),
         )
         result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
 

src/bloqade/rewrite/passes/canonicalize_ilist.py

@@ -5,6 +5,7 @@
     Walk,
     Chain,
     Fixpoint,
+    DeadCodeElimination,
 )
 from kirin.dialects.ilist import rewrite
 
@@ -24,6 +25,7 @@ def unsafe_run(self, mt: ir.Method):
                     rewrite.InlineGetItem(),
                     rewrite.FlattenAdd(),
                     rewrite.HintLen(),
+                    DeadCodeElimination(),
                 )
             )
         ).rewrite(mt.code)

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,
     is_one as is_one,
@@ -15,6 +15,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,