Fix rewrite from squin -> stim with measurement (#370)

This PR

1. unify liftIfs related rewrite rule btwn qasm2 and stim
2. refactor squin to stim rewrite.
3. Mark squin noise op as pure.
4. Add ilist canonicalize passes (with two rules added)

Co-authored-by: @johnzl-777

---------

Co-authored-by: John Long <[email protected]>

Author: kaihsin

src/bloqade/analysis/address/impls.py

@@ -6,13 +6,10 @@
 from kirin.analysis import ForwardFrame, const
 from kirin.dialects import cf, py, scf, func, ilist
 
-from bloqade import squin
-
 from .lattice import (
     Address,
     NotQubit,
     AddressReg,
-    AddressWire,
     AddressQubit,
     AddressTuple,
 )
@@ -163,63 +160,3 @@ def for_loop(
             return  # if terminate is Return, there is no result
 
         return loop_vars
-
-
-# Address lattice elements we can work with:
-## NotQubit (bottom), AnyAddress (top)
-
-## AddressTuple -> data: tuple[Address, ...]
-### Recursive type, could contain itself or other variants
-### This pops up in cases where you can have an IList/Tuple
-### That contains elements that could be other Address types
-
-## AddressReg -> data: Sequence[int]
-### specific to creation of a register of qubits
-
-## AddressQubit -> data: int
-### Base qubit address type
-
-
-@squin.wire.dialect.register(key="qubit.address")
-class SquinWireMethodTable(interp.MethodTable):
-
-    @interp.impl(squin.wire.Unwrap)
-    def unwrap(
-        self,
-        interp_: AddressAnalysis,
-        frame: ForwardFrame[Address],
-        stmt: squin.wire.Unwrap,
-    ):
-
-        origin_qubit = frame.get(stmt.qubit)
-
-        if isinstance(origin_qubit, AddressQubit):
-            return (AddressWire(origin_qubit=origin_qubit),)
-        else:
-            return (Address.top(),)
-
-    @interp.impl(squin.wire.Apply)
-    def apply(
-        self,
-        interp_: AddressAnalysis,
-        frame: ForwardFrame[Address],
-        stmt: squin.wire.Apply,
-    ):
-        return frame.get_values(stmt.inputs)
-
-
-@squin.qubit.dialect.register(key="qubit.address")
-class SquinQubitMethodTable(interp.MethodTable):
-
-    # This can be treated like a QRegNew impl
-    @interp.impl(squin.qubit.New)
-    def new(
-        self,
-        interp_: AddressAnalysis,
-        frame: ForwardFrame[Address],
-        stmt: squin.qubit.New,
-    ):
-        n_qubits = interp_.get_const_value(int, stmt.n_qubits)
-        addr = AddressReg(range(interp_.next_address, interp_.next_address + n_qubits))
-        interp_.next_address += n_qubits
-        return (addr,)

src/bloqade/analysis/measure_id/__init__.py

@@ -0,0 +1,2 @@
+from . import impls as impls
+from .analysis import MeasurementIDAnalysis as MeasurementIDAnalysis

src/bloqade/analysis/measure_id/analysis.py

@@ -0,0 +1,45 @@
+from typing import TypeVar
+
+from kirin import ir, interp
+from kirin.analysis import Forward, const
+from kirin.analysis.forward import ForwardFrame
+
+from .lattice import MeasureId, NotMeasureId
+
+
+class MeasurementIDAnalysis(Forward[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
+
+    # 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
+    ) -> 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")
+
+    # Xiu-zhe (Roger) Luo came up with this in the address analysis,
+    # reused here for convenience
+    # 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:
+        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}"
+        )

src/bloqade/analysis/measure_id/impls.py

@@ -0,0 +1,155 @@
+from kirin import types as kirin_types, interp
+from kirin.dialects import py, scf, func, ilist
+
+from bloqade.squin import wire, qubit
+
+from .lattice import (
+    AnyMeasureId,
+    NotMeasureId,
+    MeasureIdBool,
+    MeasureIdTuple,
+    InvalidMeasureId,
+)
+from .analysis import MeasurementIDAnalysis
+
+## Can't do wire right now because of
+## unresolved RFC on return type
+# from bloqade.squin import wire
+
+
+@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)
+    def measure_qubit_list(
+        self,
+        interp: MeasurementIDAnalysis,
+        frame: interp.Frame,
+        stmt: qubit.MeasureQubitList,
+    ):
+
+        # try to get the length of the list
+        ## "...safely assume the type inference will give you what you need"
+        qubits_type = stmt.qubits.type
+        # vars[0] is just the type of the elements in the ilist,
+        # vars[1] can contain a literal with length information
+        num_qubits = qubits_type.vars[1]
+        if not isinstance(num_qubits, kirin_types.Literal):
+            return (AnyMeasureId(),)
+
+        measure_id_bools = []
+        for _ in range(num_qubits.data):
+            interp.measure_count += 1
+            measure_id_bools.append(MeasureIdBool(interp.measure_count))
+
+        return (MeasureIdTuple(data=tuple(measure_id_bools)),)
+
+
+@wire.dialect.register(key="measure_id")
+class SquinWire(interp.MethodTable):
+
+    @interp.impl(wire.Measure)
+    def measure_qubit(
+        self,
+        interp: MeasurementIDAnalysis,
+        frame: interp.Frame,
+        stmt: wire.Measure,
+    ):
+        interp.measure_count += 1
+        return (MeasureIdBool(interp.measure_count),)
+
+
+@ilist.dialect.register(key="measure_id")
+class IList(interp.MethodTable):
+    @interp.impl(ilist.New)
+    # Because of the way GetItem works,
+    # A user could create an ilist of bools that
+    # ends up being a mixture of MeasureIdBool and NotMeasureId
+    def new_ilist(
+        self,
+        interp: MeasurementIDAnalysis,
+        frame: interp.Frame,
+        stmt: ilist.New,
+    ):
+
+        measure_ids_in_ilist = frame.get_values(stmt.values)
+        return (MeasureIdTuple(data=tuple(measure_ids_in_ilist)),)
+
+
+@py.tuple.dialect.register(key="measure_id")
+class PyTuple(interp.MethodTable):
+    @interp.impl(py.tuple.New)
+    def new_tuple(
+        self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.tuple.New
+    ):
+        measure_ids_in_tuple = frame.get_values(stmt.args)
+        return (MeasureIdTuple(data=tuple(measure_ids_in_tuple)),)
+
+
+@py.indexing.dialect.register(key="measure_id")
+class PyIndexing(interp.MethodTable):
+    @interp.impl(py.GetItem)
+    def getitem(
+        self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.GetItem
+    ):
+        idx = interp.get_const_value(int, stmt.index)
+        obj = frame.get(stmt.obj)
+        if isinstance(obj, MeasureIdTuple):
+            return (obj.data[idx],)
+        # just propagate these down the line
+        elif isinstance(obj, (AnyMeasureId, NotMeasureId)):
+            return (obj,)
+        else:
+            return (InvalidMeasureId(),)
+
+
+@py.binop.dialect.register(key="measure_id")
+class PyBinOp(interp.MethodTable):
+    @interp.impl(py.Add)
+    def add(self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.Add):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+
+        if isinstance(lhs, MeasureIdTuple) and isinstance(rhs, MeasureIdTuple):
+            return (MeasureIdTuple(data=lhs.data + rhs.data),)
+        else:
+            return (InvalidMeasureId(),)
+
+
+@func.dialect.register(key="measure_id")
+class Func(interp.MethodTable):
+    @interp.impl(func.Return)
+    def return_(self, _: MeasurementIDAnalysis, frame: interp.Frame, stmt: func.Return):
+        return interp.ReturnValue(frame.get(stmt.value))
+
+    # taken from Address Analysis implementation from Xiu-zhe (Roger) Luo
+    @interp.impl(
+        func.Invoke
+    )  # we know the callee already, func.Call would mean we don't know the callee @ compile time
+    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
+            ),
+        )
+        return (ret,)
+
+
+# Just let analysis propagate through
+# scf, particularly IfElse
+@scf.dialect.register(key="measure_id")
+class Scf(scf.absint.Methods):
+    pass

src/bloqade/analysis/measure_id/lattice.py

@@ -0,0 +1,82 @@
+from typing import final
+from dataclasses import dataclass
+
+from kirin.lattice import (
+    SingletonMeta,
+    BoundedLattice,
+    SimpleJoinMixin,
+    SimpleMeetMixin,
+)
+
+# Taken directly from Kai-Hsin Wu's implementation
+# with minor changes to names and addition of CanMeasureId type
+
+
+@dataclass
+class MeasureId(
+    SimpleJoinMixin["MeasureId"],
+    SimpleMeetMixin["MeasureId"],
+    BoundedLattice["MeasureId"],
+):
+
+    @classmethod
+    def bottom(cls) -> "MeasureId":
+        return InvalidMeasureId()
+
+    @classmethod
+    def top(cls) -> "MeasureId":
+        return AnyMeasureId()
+
+
+# Can pop up if user constructs some list containing a mixture
+# of bools from measure results and other places,
+# in which case the whole list is invalid
+@final
+@dataclass
+class InvalidMeasureId(MeasureId, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: MeasureId) -> bool:
+        return True
+
+
+@final
+@dataclass
+class AnyMeasureId(MeasureId, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: MeasureId) -> bool:
+        return isinstance(other, AnyMeasureId)
+
+
+@final
+@dataclass
+class NotMeasureId(MeasureId, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: MeasureId) -> bool:
+        return isinstance(other, NotMeasureId)
+
+
+@final
+@dataclass
+class MeasureIdBool(MeasureId):
+    idx: int
+
+    def is_subseteq(self, other: MeasureId) -> bool:
+        if isinstance(other, MeasureIdBool):
+            return self.idx == other.idx
+        return False
+
+
+# Might be nice to have some print override
+# here so all the CanMeasureId's/other types are consolidated for
+# readability
+
+
+@final
+@dataclass
+class MeasureIdTuple(MeasureId):
+    data: tuple[MeasureId, ...]
+
+    def is_subseteq(self, other: MeasureId) -> bool:
+        if isinstance(other, MeasureIdTuple):
+            return all(a.is_subseteq(b) for a, b in zip(self.data, other.data))
+        return False

src/bloqade/qasm2/passes/unroll_if.py

@@ -7,15 +7,22 @@
     ConstantFold,
     CommonSubexpressionElimination,
 )
+from kirin.dialects import scf, func
 
-from ..rewrite.split_ifs import LiftThenBody, SplitIfStmts
+from bloqade.rewrite.rules import LiftThenBody, SplitIfStmts
+
+from ..dialects.uop.stmts import SingleQubitGate, TwoQubitCtrlGate
+from ..dialects.core.stmts import Reset, Measure
+
+AllowedThenType = (SingleQubitGate, TwoQubitCtrlGate, Measure, Reset)
+DontLiftType = AllowedThenType + (scf.Yield, func.Return, func.Invoke)
 
 
 class UnrollIfs(Pass):
     """This pass lifts statements that are not UOP out of the if body and then splits whatever is left into multiple if statements so you obtain valid QASM2"""
 
     def unsafe_run(self, mt: ir.Method):
-        result = Walk(LiftThenBody()).rewrite(mt.code)
+        result = Walk(LiftThenBody(exclude_stmts=DontLiftType)).rewrite(mt.code)
         result = Walk(SplitIfStmts()).rewrite(mt.code).join(result)
         result = (
             Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))

src/bloqade/rewrite/__init__.py

@@ -0,0 +1 @@
+from .canonicalize_ilist import CanonicalizeIList as CanonicalizeIList