bloqade-circuit 0.7.12__py3-none-any.whl → 0.8.0__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

@@ -7,7 +7,7 @@ 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):
@@ -57,7 +57,7 @@ class FidelityAnalysis(Forward):
 
     _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()

bloqade/analysis/measure_id/impls.py

@@ -2,7 +2,7 @@ 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
 
 from .lattice import (
     AnyMeasureId,
@@ -21,22 +21,12 @@ from .analysis import MeasureIDFrame, 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
@@ -56,20 +46,6 @@ 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(
-        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)

bloqade/cirq_utils/__init__.py

@@ -1,4 +1,6 @@
-from . import noise as noise
+from . import emit as emit, noise as noise, lowering as lowering
+from .emit import emit_circuit as emit_circuit
+from .lowering import load_circuit as load_circuit
 from .parallelize import (
     transpile as transpile,
     parallelize as parallelize,

bloqade/cirq_utils/emit/__init__.py

@@ -0,0 +1,3 @@
+# NOTE: just to register methods
+from . import gate as gate, noise as noise, qubit as qubit
+from .base import emit_circuit as emit_circuit

bloqade/cirq_utils/emit/base.py

@@ -0,0 +1,243 @@
+from typing import Sequence
+from warnings import warn
+from dataclasses import field, dataclass
+
+import cirq
+from kirin import ir, types, interp
+from kirin.emit import EmitABC, EmitError, EmitFrame
+from kirin.interp import MethodTable, impl
+from kirin.dialects import py, func
+from typing_extensions import Self
+
+from bloqade.squin import kernel
+from bloqade.rewrite.passes import AggressiveUnroll
+
+
+def emit_circuit(
+    mt: ir.Method,
+    qubits: Sequence[cirq.Qid] | None = None,
+    circuit_qubits: Sequence[cirq.Qid] | None = None,
+    args: tuple = (),
+    ignore_returns: bool = False,
+) -> cirq.Circuit:
+    """Converts a squin.kernel method to a cirq.Circuit object.
+
+    Args:
+        mt (ir.Method): The kernel method from which to construct the circuit.
+
+    Keyword Args:
+        circuit_qubits (Sequence[cirq.Qid] | None):
+            A list of qubits to use as the qubits in the circuit. Defaults to None.
+            If this is None, then `cirq.LineQubit`s are inserted for every `squin.qalloc`
+            statement in the order they appear inside the kernel.
+            **Note**: If a list of qubits is provided, make sure that there is a sufficient
+            number of qubits for the resulting circuit.
+        args (tuple):
+            The arguments of the kernel function from which to emit a circuit.
+        ignore_returns (bool):
+            If `False`, emitting a circuit from a kernel that returns a value will error.
+            Set it to `True` in order to ignore the return value(s). Defaults to `False`.
+
+    ## Examples:
+
+    Here's a very basic example:
+
+    ```python
+    from bloqade import squin
+    from bloqade.cirq_utils import emit_circuit
+
+    @squin.kernel
+    def main():
+        q = squin.qalloc(2)
+        squin.h(q[0])
+        squin.cx(q[0], q[1])
+
+    circuit = emit_circuit(main)
+
+    print(circuit)
+    ```
+
+    You can also compose multiple kernels. Those are emitted as subcircuits within the "main" circuit.
+    Subkernels can accept arguments and return a value.
+
+    ```python
+    from bloqade import squin
+    from bloqade.cirq_utils import emit_circuit
+    from kirin.dialects import ilist
+    from typing import Literal
+    import cirq
+
+    @squin.kernel
+    def entangle(q: ilist.IList[squin.qubit.Qubit, Literal[2]]):
+        squin.h(q[0])
+        squin.cx(q[0], q[1])
+
+    @squin.kernel
+    def main():
+        q = squin.qalloc(2)
+        q2 = squin.qalloc(3)
+        squin.cx(q[1], q2[2])
+
+
+    # custom list of qubits on grid
+    qubits = [cirq.GridQubit(i, i+1) for i in range(5)]
+
+    circuit = emit_circuit(main, circuit_qubits=qubits)
+    print(circuit)
+
+    ```
+
+    We also passed in a custom list of qubits above. This allows you to provide a custom geometry
+    and manipulate the qubits in other circuits directly written in cirq as well.
+    """
+
+    if circuit_qubits is None and qubits is not None:
+        circuit_qubits = qubits
+        warn(
+            "The keyword argument `qubits` is deprecated. Use `circuit_qubits` instead."
+        )
+
+    if (
+        not ignore_returns
+        and isinstance(mt.code, func.Function)
+        and not mt.code.signature.output.is_subseteq(types.NoneType)
+    ):
+        raise EmitError(
+            "The method you are trying to convert to a circuit has a return value, but returning from a circuit is not supported."
+            " Set `ignore_returns = True` in order to simply ignore the return values and emit a circuit."
+        )
+
+    if len(args) != len(mt.args):
+        raise ValueError(
+            f"The method from which you're trying to emit a circuit takes {len(mt.args)} as input, but you passed in {len(args)} via the `args` keyword!"
+        )
+
+    emitter = EmitCirq(qubits=circuit_qubits)
+
+    symbol_op_trait = mt.code.get_trait(ir.SymbolOpInterface)
+    if (symbol_op_trait := mt.code.get_trait(ir.SymbolOpInterface)) is None:
+        raise EmitError("The method is not a symbol, cannot emit circuit!")
+
+    sym_name = symbol_op_trait.get_sym_name(mt.code).unwrap()
+
+    if (signature_trait := mt.code.get_trait(ir.HasSignature)) is None:
+        raise EmitError(
+            f"The method {sym_name} does not have a signature, cannot emit circuit!"
+        )
+
+    signature = signature_trait.get_signature(mt.code)
+    new_signature = func.Signature(inputs=(), output=signature.output)
+
+    callable_region = mt.callable_region.clone()
+    entry_block = callable_region.blocks[0]
+    args_ssa = list(entry_block.args)
+    first_stmt = entry_block.first_stmt
+
+    assert first_stmt is not None, "Method has no statements!"
+    if len(args_ssa) - 1 != len(args):
+        raise EmitError(
+            f"The method {sym_name} takes {len(args_ssa) - 1} arguments, but you passed in {len(args)} via the `args` keyword!"
+        )
+
+    for arg, arg_ssa in zip(args, args_ssa[1:], strict=True):
+        (value := py.Constant(arg)).insert_before(first_stmt)
+        arg_ssa.replace_by(value.result)
+        entry_block.args.delete(arg_ssa)
+
+    new_func = func.Function(
+        sym_name=sym_name, body=callable_region, signature=new_signature
+    )
+    mt_ = ir.Method(None, None, sym_name, [], mt.dialects, new_func)
+
+    AggressiveUnroll(mt_.dialects).fixpoint(mt_)
+    return emitter.run(mt_, args=())
+
+
+@dataclass
+class EmitCirqFrame(EmitFrame):
+    qubit_index: int = 0
+    qubits: Sequence[cirq.Qid] | None = None
+    circuit: cirq.Circuit = field(default_factory=cirq.Circuit)
+
+
+def _default_kernel():
+    return kernel
+
+
+@dataclass
+class EmitCirq(EmitABC[EmitCirqFrame, cirq.Circuit]):
+    keys = ["emit.cirq", "main"]
+    dialects: ir.DialectGroup = field(default_factory=_default_kernel)
+    void = cirq.Circuit()
+    qubits: Sequence[cirq.Qid] | None = None
+
+    def initialize(self) -> Self:
+        return super().initialize()
+
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> EmitCirqFrame:
+        return EmitCirqFrame(
+            code, has_parent_access=has_parent_access, qubits=self.qubits
+        )
+
+    def run_method(self, method: ir.Method, args: tuple[cirq.Circuit, ...]):
+        return self.run_callable(method.code, args)
+
+    def run_callable_region(
+        self,
+        frame: EmitCirqFrame,
+        code: ir.Statement,
+        region: ir.Region,
+        args: tuple,
+    ):
+        if len(region.blocks) > 0:
+            block_args = list(region.blocks[0].args)
+            # NOTE: skip self arg
+            frame.set_values(block_args[1:], args)
+
+        results = self.eval_stmt(frame, code)
+        if isinstance(results, tuple):
+            if len(results) == 0:
+                return self.void
+            elif len(results) == 1:
+                return results[0]
+        raise interp.InterpreterError(f"Unexpected results {results}")
+
+    def emit_block(self, frame: EmitCirqFrame, block: ir.Block) -> cirq.Circuit:
+        for stmt in block.stmts:
+            result = self.eval_stmt(frame, stmt)
+            if isinstance(result, tuple):
+                frame.set_values(stmt.results, result)
+
+        return frame.circuit
+
+
+@func.dialect.register(key="emit.cirq")
+class __FuncEmit(MethodTable):
+
+    @impl(func.Function)
+    def emit_func(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Function):
+        emit.run_ssacfg_region(frame, stmt.body, ())
+        return (frame.circuit,)
+
+    @impl(func.Invoke)
+    def emit_invoke(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Invoke):
+        raise EmitError(
+            "Function invokes should need to be inlined! "
+            "If you called the emit_circuit method, that should have happened, please report this issue."
+        )
+
+    @impl(func.Return)
+    def return_(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Return):
+        # NOTE: should only be hit if ignore_returns == True
+        return ()
+
+
+@py.indexing.dialect.register(key="emit.cirq")
+class __Concrete(interp.MethodTable):
+
+    @interp.impl(py.indexing.GetItem)
+    def getindex(self, interp, frame: interp.Frame, stmt: py.indexing.GetItem):
+        # NOTE: no support for indexing into single statements in cirq
+        return ()