bloqade-circuit 0.6.2__py3-none-any.whl → 0.9.1__py3-none-any.whl

bloqade/analysis/address/lattice.py

@@ -1,12 +1,18 @@
 from typing import Sequence, final
 from dataclasses import dataclass
 
+from kirin import ir, types
 from kirin.lattice import (
     SingletonMeta,
     BoundedLattice,
     SimpleJoinMixin,
     SimpleMeetMixin,
 )
+from kirin.analysis import const
+from kirin.dialects import ilist
+from kirin.ir.attrs.abc import LatticeAttributeMeta
+
+from bloqade.types import QubitType
 
 
 @dataclass
@@ -18,54 +24,94 @@ class Address(
 
     @classmethod
     def bottom(cls) -> "Address":
-        return NotQubit()
+        return Bottom()
 
     @classmethod
     def top(cls) -> "Address":
-        return AnyAddress()
+        return Unknown()
+
+    @classmethod
+    def from_type(cls, typ: types.TypeAttribute):
+        if typ.is_subseteq(ilist.IListType[QubitType]):
+            return UnknownReg()
+        elif typ.is_subseteq(QubitType):
+            return UnknownQubit()
+        else:
+            return Unknown()
 
 
 @final
-@dataclass
-class NotQubit(Address, metaclass=SingletonMeta):
+class Bottom(Address, metaclass=SingletonMeta):
+    """Error during interpretation"""
 
     def is_subseteq(self, other: Address) -> bool:
         return True
 
 
 @final
-@dataclass
-class AnyAddress(Address, metaclass=SingletonMeta):
+class Unknown(Address, metaclass=SingletonMeta):
+    """Can't determine if it is an address or constant."""
 
     def is_subseteq(self, other: Address) -> bool:
-        return isinstance(other, AnyAddress)
+        return isinstance(other, Unknown)
 
 
 @final
 @dataclass
-class AddressTuple(Address):
-    data: tuple[Address, ...]
+class ConstResult(Address):
+    """Stores a constant prop result in the lattice"""
+
+    result: const.Result
 
     def is_subseteq(self, other: Address) -> bool:
-        if isinstance(other, AddressTuple):
-            return all(a.is_subseteq(b) for a, b in zip(self.data, other.data))
-        return False
+        return isinstance(other, ConstResult) and self.result.is_subseteq(other.result)
+
+
+class QubitLike(Address):
+    def join(self, other: Address):
+        if isinstance(other, QubitLike):
+            return super().join(other)
+        return self.bottom()
+
+    def meet(self, other: Address):
+        if isinstance(other, QubitLike):
+            return super().meet(other)
+        return self.bottom()
 
 
 @final
-@dataclass
-class AddressReg(Address):
-    data: Sequence[int]
+class UnknownQubit(QubitLike, metaclass=SingletonMeta):
+    """A lattice element representing a single qubit with an unknown address."""
 
     def is_subseteq(self, other: Address) -> bool:
-        if isinstance(other, AddressReg):
-            return self.data == other.data
-        return False
+        return isinstance(other, QubitLike)
+
+
+class RegisterLike(Address):
+    def join(self, other: Address):
+        if isinstance(other, RegisterLike):
+            return super().join(other)
+        return self.bottom()
+
+    def meet(self, other: Address):
+        if isinstance(other, RegisterLike):
+            return super().meet(other)
+        return self.bottom()
+
+
+@final
+class UnknownReg(RegisterLike, metaclass=SingletonMeta):
+    """A lattice element representing a container of qubits with unknown indices."""
+
+    def is_subseteq(self, other: Address) -> bool:
+        return isinstance(other, RegisterLike)
 
 
 @final
 @dataclass
-class AddressQubit(Address):
+class AddressQubit(QubitLike):
+    """A lattice element representing a single qubit with a known address."""
+
     data: int
 
     def is_subseteq(self, other: Address) -> bool:
@@ -76,10 +122,149 @@ class AddressQubit(Address):
 
 @final
 @dataclass
-class AddressWire(Address):
-    origin_qubit: AddressQubit
+class AddressReg(RegisterLike):
+    """A lattice element representing a container of qubits with known indices."""
+
+    data: Sequence[int]
 
     def is_subseteq(self, other: Address) -> bool:
-        if isinstance(other, AddressWire):
-            return self.origin_qubit == other.origin_qubit
-        return False
+        return isinstance(other, AddressReg) and self.data == other.data
+
+    @property
+    def qubits(self) -> tuple[AddressQubit, ...]:
+        return tuple(AddressQubit(i) for i in self.data)
+
+
+@final
+@dataclass
+class PartialLambda(Address):
+    """Represents a partially known lambda function"""
+
+    argnames: list[str]
+    code: ir.Statement
+    captured: tuple[Address, ...]
+
+    def join(self, other: Address) -> Address:
+        if other is other.bottom():
+            return self
+
+        if not isinstance(other, PartialLambda):
+            return self.top().join(other)  # widen self
+
+        if self.code is not other.code:
+            return self.top()  # lambda stmt is pure
+
+        if len(self.captured) != len(other.captured):
+            return self.bottom()  # err
+
+        return PartialLambda(
+            self.argnames,
+            self.code,
+            tuple(x.join(y) for x, y in zip(self.captured, other.captured)),
+        )
+
+    def meet(self, other: Address) -> Address:
+        if not isinstance(other, PartialLambda):
+            return self.top().meet(other)
+
+        if self.code is not other.code:
+            return self.bottom()
+
+        if len(self.captured) != len(other.captured):
+            return self.top()
+
+        return PartialLambda(
+            self.argnames,
+            self.code,
+            tuple(x.meet(y) for x, y in zip(self.captured, other.captured)),
+        )
+
+    def is_subseteq(self, other: Address) -> bool:
+        return (
+            isinstance(other, PartialLambda)
+            and self.code is other.code
+            and self.argnames == other.argnames
+            and len(self.captured) == len(other.captured)
+            and all(
+                self_ele.is_subseteq(other_ele)
+                for self_ele, other_ele in zip(self.captured, other.captured)
+            )
+        )
+
+
+@dataclass
+class StaticContainer(Address):
+    """A lattice element representing the results of any static container, e. g. ilist or tuple."""
+
+    data: tuple[Address, ...]
+
+    @classmethod
+    def new(cls, data: tuple[Address, ...]):
+        return cls(data)
+
+    def join(self, other: "Address") -> "Address":
+        if isinstance(other, type(self)) and len(self.data) == len(other.data):
+            return self.new(tuple(x.join(y) for x, y in zip(self.data, other.data)))
+        return self.top()
+
+    def meet(self, other: "Address") -> "Address":
+        if isinstance(other, type(self)) and len(self.data) == len(other.data):
+            return self.new(tuple(x.meet(y) for x, y in zip(self.data, other.data)))
+        return self.bottom()
+
+    def is_subseteq(self, other: "Address") -> bool:
+        return (
+            isinstance(other, type(self))
+            and len(self.data) == len(other.data)
+            and all(x.is_subseteq(y) for x, y in zip(self.data, other.data))
+        )
+
+
+class PartialIListMeta(LatticeAttributeMeta):
+    """This metaclass assures that PartialILists of ConstResults or AddressQubits are canonicalized
+    to a single ConstResult or AddressReg respectively.
+
+    because AddressReg is a specialization of PartialIList, being a container of pure qubit
+    addresses. For Operations that act in generic containers (e.g., ilist.ForEach),
+    AddressReg is treated as PartialIList but for other types of analysis it is often
+    useful to distinguish between a generic IList and a pure qubit address list.
+
+    Inside the method tables the `GetValuesMixin` implements a method that effectively
+    undoes this canonicalization.
+
+    """
+
+    def __call__(cls, data: tuple[Address, ...]):
+        # TODO: when constant prop has PartialIList, make sure to canonicalize here.
+        if types.is_tuple_of(data, ConstResult) and types.is_tuple_of(
+            all_constants := tuple(ele.result for ele in data), const.Value
+        ):
+            # all constants, create constant list
+            return ConstResult(
+                const.Value(ilist.IList([ele.data for ele in all_constants]))
+            )
+        elif types.is_tuple_of(data, AddressQubit):
+            # all qubits create qubit register
+            return AddressReg(tuple(ele.data for ele in data))
+        else:
+            return super().__call__(data)
+
+
+@final
+class PartialIList(StaticContainer, metaclass=PartialIListMeta):
+    """A lattice element representing a partially known ilist."""
+
+
+class PartialTupleMeta(LatticeAttributeMeta):
+    """This metaclass assures that PartialTuples of ConstResults are canonicalized to a single ConstResult."""
+
+    def __call__(cls, data: tuple[Address, ...]):
+        if not types.is_tuple_of(data, ConstResult):
+            return super().__call__(data)
+
+        return ConstResult(const.PartialTuple(tuple(ele.result for ele in data)))
+
+
+@final
+class PartialTuple(StaticContainer, metaclass=PartialTupleMeta):
+    """A lattice element representing a partially known tuple."""

bloqade/analysis/fidelity/analysis.py

@@ -4,10 +4,9 @@ from dataclasses import field
 from kirin import ir
 from kirin.lattice import EmptyLattice
 from kirin.analysis import Forward
-from kirin.interp.value import Successor
 from kirin.analysis.forward import ForwardFrame
 
-from ..address import AddressAnalysis
+from ..address import Address, AddressAnalysis
 
 
 class FidelityAnalysis(Forward):
@@ -48,16 +47,12 @@ class FidelityAnalysis(Forward):
     The fidelity of the gate set described by the analysed program. It reduces whenever a noise channel is encountered.
     """
 
-    _current_gate_fidelity: float = field(init=False)
-
     atom_survival_probability: list[float] = field(init=False)
     """
     The probabilities that each of the atoms in the register survive the duration of the analysed program. The order of the list follows the order they are in the register.
     """
 
-    _current_atom_survival_probability: list[float] = field(init=False)
-
-    addr_frame: ForwardFrame = field(init=False)
+    addr_frame: ForwardFrame[Address] = field(init=False)
 
     def initialize(self):
         super().initialize()
@@ -67,28 +62,21 @@ class FidelityAnalysis(Forward):
         ]
         return self
 
-    def posthook_succ(self, frame: ForwardFrame, succ: Successor):
-        self.gate_fidelity *= self._current_gate_fidelity
-        for i, _current_survival in enumerate(self._current_atom_survival_probability):
-            self.atom_survival_probability[i] *= _current_survival
-
-    def eval_stmt_fallback(self, frame: ForwardFrame, stmt: ir.Statement):
+    def eval_fallback(self, frame: ForwardFrame, node: ir.Statement):
         # NOTE: default is to conserve fidelity, so do nothing here
         return
 
-    def run_method(self, method: ir.Method, args: tuple[EmptyLattice, ...]):
-        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+    def run(self, method: ir.Method, *args, **kwargs) -> tuple[ForwardFrame, Any]:
+        self._run_address_analysis(method)
+        return super().run(method, *args, **kwargs)
 
-    def run_analysis(
-        self, method: ir.Method, args: tuple | None = None, *, no_raise: bool = True
-    ) -> tuple[ForwardFrame, Any]:
-        self._run_address_analysis(method, no_raise=no_raise)
-        return super().run_analysis(method, args, no_raise=no_raise)
-
-    def _run_address_analysis(self, method: ir.Method, no_raise: bool):
+    def _run_address_analysis(self, method: ir.Method):
         addr_analysis = AddressAnalysis(self.dialects)
-        addr_frame, _ = addr_analysis.run_analysis(method=method, no_raise=no_raise)
+        addr_frame, _ = addr_analysis.run(method=method)
         self.addr_frame = addr_frame
 
         # NOTE: make sure we have as many probabilities as we have addresses
         self.atom_survival_probability = [1.0] * addr_analysis.qubit_count
+
+    def method_self(self, method: ir.Method) -> EmptyLattice:
+        return self.lattice.bottom()

bloqade/analysis/measure_id/__init__.py

@@ -1,2 +1,5 @@
 from . import impls as impls
-from .analysis import MeasurementIDAnalysis as MeasurementIDAnalysis
+from .analysis import (
+    MeasureIDFrame as MeasureIDFrame,
+    MeasurementIDAnalysis as MeasurementIDAnalysis,
+)

bloqade/analysis/measure_id/analysis.py

@@ -1,13 +1,19 @@
 from typing import TypeVar
+from dataclasses import field, dataclass
 
-from kirin import ir, interp
-from kirin.analysis import Forward, const
+from kirin import ir
+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
@@ -15,31 +21,34 @@ class MeasurementIDAnalysis(Forward[MeasureId]):
     # then use this to generate the negative values for target rec indices
     measure_count = 0
 
+    def initialize_frame(
+        self, node: ir.Statement, *, has_parent_access: bool = False
+    ) -> 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_stmt_fallback(
-        self, frame: ForwardFrame[MeasureId], stmt: ir.Statement
+    def eval_fallback(
+        self, frame: ForwardFrame[MeasureId], node: ir.Statement
     ) -> tuple[MeasureId, ...]:
-        return tuple(NotMeasureId() for _ in stmt.results)
-
-    def run_method(self, method: ir.Method, args: tuple[MeasureId, ...]):
-        # NOTE: we do not support dynamic calls here, thus no need to propagate method object
-        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
-
-    T = TypeVar("T")
+        return tuple(NotMeasureId() for _ in node.results)
 
     # Xiu-zhe (Roger) Luo came up with this in the address analysis,
-    # reused here for convenience
+    # reused here for convenience (now modified to be a bit more graceful)
     # TODO: Remove this function once upgrade to kirin 0.18 happens,
     #       method is built-in to interpreter then
-    def get_const_value(self, input_type: type[T], value: ir.SSAValue) -> T:
+
+    T = TypeVar("T")
+
+    def get_const_value(
+        self, input_type: type[T] | tuple[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
-            raise interp.InterpreterError(
-                f"Expected constant value <type = {input_type}>, got {data}"
-            )
-        raise interp.InterpreterError(
-            f"Expected constant value <type = {input_type}>, got {value}"
-        )
+
+        return None
+
+    def method_self(self, method: ir.Method) -> MeasureId:
+        return self.lattice.bottom()

bloqade/analysis/measure_id/impls.py

@@ -1,7 +1,8 @@
 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
+from bloqade import qubit, annotate
 
 from .lattice import (
     AnyMeasureId,
@@ -10,7 +11,7 @@ from .lattice import (
     MeasureIdTuple,
     InvalidMeasureId,
 )
-from .analysis import MeasurementIDAnalysis
+from .analysis import MeasureIDFrame, MeasurementIDAnalysis
 
 ## Can't do wire right now because of
 ## unresolved RFC on return type
@@ -20,22 +21,12 @@ from .analysis import MeasurementIDAnalysis
 @qubit.dialect.register(key="measure_id")
 class SquinQubit(interp.MethodTable):
 
-    @interp.impl(qubit.MeasureQubit)
-    def measure_qubit(
-        self,
-        interp: MeasurementIDAnalysis,
-        frame: interp.Frame,
-        stmt: qubit.MeasureQubit,
-    ):
-        interp.measure_count += 1
-        return (MeasureIdBool(interp.measure_count),)
-
-    @interp.impl(qubit.MeasureQubitList)
+    @interp.impl(qubit.stmts.Measure)
     def measure_qubit_list(
         self,
         interp: MeasurementIDAnalysis,
         frame: interp.Frame,
-        stmt: qubit.MeasureQubitList,
+        stmt: qubit.stmts.Measure,
     ):
 
         # try to get the length of the list
@@ -55,18 +46,18 @@ class SquinQubit(interp.MethodTable):
         return (MeasureIdTuple(data=tuple(measure_id_bools)),)
 
 
-@wire.dialect.register(key="measure_id")
-class SquinWire(interp.MethodTable):
-
-    @interp.impl(wire.Measure)
-    def measure_qubit(
+@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: interp.Frame,
-        stmt: wire.Measure,
+        frame: MeasureIDFrame,
+        stmt: annotate.stmts.SetObservable | annotate.stmts.SetDetector,
     ):
-        interp.measure_count += 1
-        return (MeasureIdBool(interp.measure_count),)
+        frame.num_measures_at_stmt[stmt] = interp.measure_count
+        return (NotMeasureId(),)
 
 
 @ilist.dialect.register(key="measure_id")
@@ -102,10 +93,23 @@ class PyIndexing(interp.MethodTable):
     def getitem(
         self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.GetItem
     ):
-        idx = interp.get_const_value(int, stmt.index)
+
+        idx_or_slice = interp.get_const_value((int, slice), stmt.index)
+        if idx_or_slice is None:
+            return (InvalidMeasureId(),)
+
+        # hint = stmt.index.hints.get("const")
+        # if hint is None or not isinstance(hint, const.Value):
+        #    return (InvalidMeasureId(),)
+
         obj = frame.get(stmt.obj)
         if isinstance(obj, MeasureIdTuple):
-            return (obj.data[idx],)
+            if isinstance(idx_or_slice, slice):
+                return (MeasureIdTuple(data=obj.data[idx_or_slice]),)
+            elif isinstance(idx_or_slice, int):
+                return (obj.data[idx_or_slice],)
+            else:
+                return (InvalidMeasureId(),)
         # just propagate these down the line
         elif isinstance(obj, (AnyMeasureId, NotMeasureId)):
             return (obj,)
@@ -113,6 +117,15 @@ class PyIndexing(interp.MethodTable):
             return (InvalidMeasureId(),)
 
 
+@py.assign.dialect.register(key="measure_id")
+class PyAssign(interp.MethodTable):
+    @interp.impl(py.Alias)
+    def alias(
+        self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.assign.Alias
+    ):
+        return (frame.get(stmt.value),)
+
+
 @py.binop.dialect.register(key="measure_id")
 class PyBinOp(interp.MethodTable):
     @interp.impl(py.Add)
@@ -139,11 +152,10 @@ class Func(interp.MethodTable):
     def invoke(
         self, interp_: MeasurementIDAnalysis, frame: interp.Frame, stmt: func.Invoke
     ):
-        _, ret = interp_.run_method(
-            stmt.callee,
-            interp_.permute_values(
-                stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
-            ),
+        _, ret = interp_.call(
+            stmt.callee.code,
+            interp_.method_self(stmt.callee),
+            *frame.get_values(stmt.inputs),
         )
         return (ret,)
 
@@ -152,4 +164,33 @@ class Func(interp.MethodTable):
 # 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)

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,
+)

bloqade/annotate/_dialect.py

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

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: ...

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)

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)