bloqade-circuit 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

bloqade/analysis/address/impls.py

@@ -6,13 +6,10 @@ from kirin import interp
 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,
 )
@@ -73,8 +70,19 @@ class PyList(interp.MethodTable):
 class PyIndexing(interp.MethodTable):
     @interp.impl(py.GetItem)
     def getitem(self, interp: AddressAnalysis, frame: interp.Frame, stmt: py.GetItem):
-        # Integer index into the thing being indexed
-        idx = interp.get_const_value(int, stmt.index)
+
+        # determine if the index is an int constant
+        # or a slice
+        hint = stmt.index.hints.get("const")
+        if hint is None:
+            return (NotQubit(),)
+
+        if isinstance(hint, const.Value):
+            idx = hint.data
+        elif isinstance(hint, slice):
+            idx = hint
+        else:
+            return (NotQubit(),)
 
         # The object being indexed into
         obj = frame.get(stmt.obj)
@@ -82,10 +90,15 @@ class PyIndexing(interp.MethodTable):
         # so we just extract that here
         if isinstance(obj, AddressTuple):
             return (obj.data[idx],)
-        # an AddressReg is guaranteed to just have some sequence
-        # of integers which is directly pluggable to AddressQubit
+        # If idx is an integer index into an AddressReg,
+        # then it's safe to assume a single qubit is being accessed.
+        # On the other hand, if it's a slice, we return
+        # a new AddressReg to preserve the new sequence.
         elif isinstance(obj, AddressReg):
-            return (AddressQubit(obj.data[idx]),)
+            if isinstance(idx, slice):
+                return (AddressReg(data=obj.data[idx]),)
+            if isinstance(idx, int):
+                return (AddressQubit(obj.data[idx]),)
         else:
             return (NotQubit(),)
 
@@ -163,63 +176,3 @@ class Scf(scf.absint.Methods):
             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,)

bloqade/analysis/measure_id/__init__.py

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

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}"
+        )

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

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

bloqade/qasm2/passes/unroll_if.py

@@ -7,15 +7,22 @@ from kirin.rewrite import (
     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())))

bloqade/rewrite/passes/__init__.py

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

bloqade/rewrite/passes/canonicalize_ilist.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+)
+from kirin.analysis import const
+
+from ..rules.flatten_ilist import FlattenAddOpIList
+from ..rules.inline_getitem_ilist import InlineGetItemFromIList
+
+
+@dataclass
+class CanonicalizeIList(Pass):
+
+    def unsafe_run(self, mt: ir.Method):
+
+        cp_result_frame, _ = const.Propagate(dialects=mt.dialects).run_analysis(mt)
+
+        return Fixpoint(
+            Chain(
+                Walk(InlineGetItemFromIList(constprop_result=cp_result_frame.entries)),
+                Walk(FlattenAddOpIList()),
+            )
+        ).rewrite(mt.code)

bloqade/rewrite/rules/__init__.py

@@ -0,0 +1 @@
+from .split_ifs import LiftThenBody as LiftThenBody, SplitIfStmts as SplitIfStmts

bloqade/rewrite/rules/flatten_ilist.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.dialects import py, ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class FlattenAddOpIList(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, py.binop.Add):
+            return RewriteResult()
+
+        # check if we are adding two ilist.New objects
+        new_data = ()
+
+        # lhs:
+        if not isinstance(node.lhs.owner, ilist.New):
+            if not (
+                isinstance(node.lhs.owner, py.Constant)
+                and isinstance(
+                    const_ilist := node.lhs.owner.value.unwrap(), ilist.IList
+                )
+                and len(const_ilist.data) == 0
+            ):
+                return RewriteResult()
+
+        else:
+            new_data += node.lhs.owner.values
+
+        # rhs:
+        if not isinstance(node.rhs.owner, ilist.New):
+            if not (
+                isinstance(node.rhs.owner, py.Constant)
+                and isinstance(
+                    const_ilist := node.rhs.owner.value.unwrap(), ilist.IList
+                )
+                and len(const_ilist.data) == 0
+            ):
+                return RewriteResult()
+
+        else:
+            new_data += node.rhs.owner.values
+
+        new_stmt = ilist.New(values=new_data)
+        node.replace_by(new_stmt)
+
+        return RewriteResult(
+            has_done_something=True,
+        )

bloqade/rewrite/rules/inline_getitem_ilist.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.analysis import const
+from kirin.dialects import py, ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class InlineGetItemFromIList(RewriteRule):
+    constprop_result: dict[ir.SSAValue, const.Result]
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, py.indexing.GetItem):
+            return RewriteResult()
+
+        if not isinstance(node.obj.owner, ilist.New):
+            return RewriteResult()
+
+        if not isinstance(
+            index_value := self.constprop_result.get(node.index), const.Value
+        ):
+            return RewriteResult()
+
+        elem_ssa = node.obj.owner.values[index_value.data]
+
+        node.result.replace_by(elem_ssa)
+
+        return RewriteResult(
+            has_done_something=True,
+        )

bloqade/{qasm2/rewrite → rewrite/rules}/split_ifs.py

@@ -1,18 +1,23 @@
+from dataclasses import field, dataclass
+
 from kirin import ir
 from kirin.dialects import scf, func
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from ..dialects.uop.stmts import SingleQubitGate, TwoQubitCtrlGate
-from ..dialects.core.stmts import Reset, Measure
 
-# TODO: unify with PR #248
-AllowedThenType = SingleQubitGate | TwoQubitCtrlGate | Measure | Reset
+@dataclass
+class LiftThenBody(RewriteRule):
+    """
+    Lifts anything that's not in the `exclude_stmts` in the *then* body
+
 
-DontLiftType = AllowedThenType | scf.Yield | func.Return | func.Invoke
+    Args:
+        exclude_stmts: A tuple of statement types that should not be lifted from the then body.
+            Defaults to an empty tuple, meaning all statements are lifted.
 
+    """
 
-class LiftThenBody(RewriteRule):
-    """Lifts anything that's not a UOP or a yield/return out of the then body"""
+    exclude_stmts: tuple[type[ir.Statement], ...] = field(default_factory=tuple)
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         if not isinstance(node, scf.IfElse):
@@ -20,7 +25,9 @@ class LiftThenBody(RewriteRule):
 
         then_stmts = node.then_body.stmts()
 
-        lift_stmts = [stmt for stmt in then_stmts if not isinstance(stmt, DontLiftType)]
+        lift_stmts = [
+            stmt for stmt in then_stmts if not isinstance(stmt, self.exclude_stmts)
+        ]
 
         if len(lift_stmts) == 0:
             return RewriteResult()

bloqade/squin/__init__.py

@@ -3,6 +3,7 @@ from . import (
     wire as wire,
     noise as noise,
     qubit as qubit,
+    analysis as analysis,
     lowering as lowering,
     _typeinfer as _typeinfer,
 )

bloqade/squin/analysis/__init__.py

@@ -0,0 +1 @@
+from . import address_impl as address_impl

bloqade/squin/analysis/address_impl.py

@@ -0,0 +1,71 @@
+from kirin import interp
+from kirin.analysis import ForwardFrame
+
+from bloqade.analysis.address.lattice import (
+    Address,
+    AddressReg,
+    AddressWire,
+    AddressQubit,
+)
+from bloqade.analysis.address.analysis import AddressAnalysis
+
+from .. import wire, qubit
+
+# 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
+
+
+@wire.dialect.register(key="qubit.address")
+class SquinWireMethodTable(interp.MethodTable):
+
+    @interp.impl(wire.Unwrap)
+    def unwrap(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: 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(wire.Apply)
+    def apply(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: wire.Apply,
+    ):
+        return frame.get_values(stmt.inputs)
+
+
+@qubit.dialect.register(key="qubit.address")
+class SquinQubitMethodTable(interp.MethodTable):
+
+    # This can be treated like a QRegNew impl
+    @interp.impl(qubit.New)
+    def new(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: 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,)

bloqade/squin/cirq/lowering.py

@@ -281,7 +281,8 @@ class Squin(lowering.LoweringABC[CirqNode]):
             return state.current_frame.push(op.stmts.Z())
 
         # NOTE: just for the Z gate, an arbitrary exponent is equivalent to the ShiftOp
-        t = node.exponent
+        # up to a minus sign!
+        t = -node.exponent
         theta = state.current_frame.push(py.Constant(math.pi * t))
         return state.current_frame.push(op.stmts.ShiftOp(theta=theta.result))