bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/analysis/address/__init__.py

@@ -0,0 +1,11 @@
+from . import impls as impls
+from .lattice import (
+    Address as Address,
+    NotQubit as NotQubit,
+    AddressReg as AddressReg,
+    AnyAddress as AnyAddress,
+    AddressWire as AddressWire,
+    AddressQubit as AddressQubit,
+    AddressTuple as AddressTuple,
+)
+from .analysis import AddressAnalysis as AddressAnalysis

bloqade/analysis/address/analysis.py

@@ -0,0 +1,60 @@
+from typing import TypeVar
+from dataclasses import field
+
+from kirin import ir, interp
+from kirin.analysis import Forward, const
+from kirin.analysis.forward import ForwardFrame
+
+from bloqade.types import QubitType
+
+from .lattice import Address
+
+
+class AddressAnalysis(Forward[Address]):
+    """
+    This analysis pass can be used to track the global addresses of qubits and wires.
+    """
+
+    keys = ["qubit.address"]
+    lattice = Address
+    next_address: int = field(init=False)
+
+    def initialize(self):
+        super().initialize()
+        self.next_address: int = 0
+        return self
+
+    @property
+    def qubit_count(self) -> int:
+        """Total number of qubits found by the analysis."""
+        return self.next_address
+
+    T = TypeVar("T")
+
+    def get_const_value(self, typ: type[T], value: ir.SSAValue) -> T:
+        if isinstance(hint := value.hints.get("const"), const.Value):
+            data = hint.data
+            if isinstance(data, typ):
+                return hint.data
+            raise interp.InterpreterError(
+                f"Expected constant value <type = {typ}>, got {data}"
+            )
+        raise interp.InterpreterError(
+            f"Expected constant value <type = {typ}>, got {value}"
+        )
+
+    def eval_stmt_fallback(
+        self, frame: ForwardFrame[Address], stmt: ir.Statement
+    ) -> tuple[Address, ...] | interp.SpecialValue[Address]:
+        return tuple(
+            (
+                self.lattice.top()
+                if result.type.is_subseteq(QubitType)
+                else self.lattice.bottom()
+            )
+            for result in stmt.results
+        )
+
+    def run_method(self, method: ir.Method, args: tuple[Address, ...]):
+        # 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)

bloqade/analysis/address/impls.py

@@ -0,0 +1,228 @@
+"""
+qubit.address method table for a few builtin dialects.
+"""
+
+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,
+)
+from .analysis import AddressAnalysis
+
+
+@py.binop.dialect.register(key="qubit.address")
+class PyBinOp(interp.MethodTable):
+
+    @interp.impl(py.Add)
+    def add(self, interp: AddressAnalysis, frame: interp.Frame, stmt: py.Add):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+
+        if isinstance(lhs, AddressTuple) and isinstance(rhs, AddressTuple):
+            return (AddressTuple(data=lhs.data + rhs.data),)
+        else:
+            return (NotQubit(),)
+
+
+@py.tuple.dialect.register(key="qubit.address")
+class PyTuple(interp.MethodTable):
+    @interp.impl(py.tuple.New)
+    def new_tuple(
+        self,
+        interp: AddressAnalysis,
+        frame: interp.Frame,
+        stmt: py.tuple.New,
+    ):
+        return (AddressTuple(frame.get_values(stmt.args)),)
+
+
+@ilist.dialect.register(key="qubit.address")
+class IList(interp.MethodTable):
+    @interp.impl(ilist.New)
+    def new_ilist(
+        self,
+        interp: AddressAnalysis,
+        frame: interp.Frame,
+        stmt: ilist.New,
+    ):
+        return (AddressTuple(frame.get_values(stmt.values)),)
+
+
+@py.list.dialect.register(key="qubit.address")
+class PyList(interp.MethodTable):
+    @interp.impl(py.list.New)
+    def new_ilist(
+        self,
+        interp: AddressAnalysis,
+        frame: interp.Frame,
+        stmt: py.list.New,
+    ):
+        return (AddressTuple(frame.get_values(stmt.args)),)
+
+
+@py.indexing.dialect.register(key="qubit.address")
+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)
+        # The object being indexed into
+        obj = frame.get(stmt.obj)
+        # The `data` attributes holds onto other Address types
+        # 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
+        elif isinstance(obj, AddressReg):
+            return (AddressQubit(obj.data[idx]),)
+        else:
+            return (NotQubit(),)
+
+
+@py.assign.dialect.register(key="qubit.address")
+class PyAssign(interp.MethodTable):
+    @interp.impl(py.Alias)
+    def alias(self, interp: AddressAnalysis, frame: interp.Frame, stmt: py.Alias):
+        return (frame.get(stmt.value),)
+
+
+@func.dialect.register(key="qubit.address")
+class Func(interp.MethodTable):
+    @interp.impl(func.Return)
+    def return_(self, _: AddressAnalysis, frame: interp.Frame, stmt: func.Return):
+        return interp.ReturnValue(frame.get(stmt.value))
+
+    # TODO: replace with the generic implementation
+    @interp.impl(func.Invoke)
+    def invoke(self, interp_: AddressAnalysis, 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,)
+
+    # TODO: support lambda?
+
+
+@cf.dialect.register(key="qubit.address")
+class Cf(cf.typeinfer.TypeInfer):
+    # NOTE: cf just re-use the type infer method table
+    # it's the same process as type infer.
+    pass
+
+
+@scf.dialect.register(key="qubit.address")
+class Scf(scf.absint.Methods):
+
+    @interp.impl(scf.For)
+    def for_loop(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: scf.For,
+    ):
+        if not isinstance(hint := stmt.iterable.hints.get("const"), const.Value):
+            return interp_.eval_stmt_fallback(frame, stmt)
+
+        iterable = hint.data
+        loop_vars = frame.get_values(stmt.initializers)
+        body_block = stmt.body.blocks[0]
+        block_args = body_block.args
+
+        # NOTE: we need to actually run iteration in case there are
+        # new allocations/re-assign in the loop body.
+        for _ in iterable:
+            with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
+                body_frame.entries.update(frame.entries)
+                body_frame.set_values(
+                    block_args,
+                    (NotQubit(),) + loop_vars,
+                )
+                loop_vars = interp_.run_ssacfg_region(body_frame, stmt.body)
+
+            if loop_vars is None:
+                loop_vars = ()
+            elif isinstance(loop_vars, interp.ReturnValue):
+                return loop_vars
+
+        if isinstance(body_block.last_stmt, func.Return):
+            frame.worklist.append(interp.Successor(body_block, NotQubit(), *loop_vars))
+            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)
+
+        return (AddressWire(origin_qubit=origin_qubit),)
+
+    @interp.impl(squin.wire.Apply)
+    def apply(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: squin.wire.Apply,
+    ):
+
+        origin_qubits = tuple(
+            [frame.get(input_elem).origin_qubit for input_elem in stmt.inputs]
+        )
+        new_address_wires = tuple(
+            [AddressWire(origin_qubit=origin_qubit) for origin_qubit in origin_qubits]
+        )
+        return new_address_wires
+
+
+@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/address/lattice.py

@@ -0,0 +1,85 @@
+from typing import Sequence, final
+from dataclasses import dataclass
+
+from kirin.lattice import (
+    SingletonMeta,
+    BoundedLattice,
+    SimpleJoinMixin,
+    SimpleMeetMixin,
+)
+
+
+@dataclass
+class Address(
+    SimpleJoinMixin["Address"],
+    SimpleMeetMixin["Address"],
+    BoundedLattice["Address"],
+):
+
+    @classmethod
+    def bottom(cls) -> "Address":
+        return NotQubit()
+
+    @classmethod
+    def top(cls) -> "Address":
+        return AnyAddress()
+
+
+@final
+@dataclass
+class NotQubit(Address, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: Address) -> bool:
+        return True
+
+
+@final
+@dataclass
+class AnyAddress(Address, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: Address) -> bool:
+        return isinstance(other, AnyAddress)
+
+
+@final
+@dataclass
+class AddressTuple(Address):
+    data: tuple[Address, ...]
+
+    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
+
+
+@final
+@dataclass
+class AddressReg(Address):
+    data: Sequence[int]
+
+    def is_subseteq(self, other: Address) -> bool:
+        if isinstance(other, AddressReg):
+            return self.data == other.data
+        return False
+
+
+@final
+@dataclass
+class AddressQubit(Address):
+    data: int
+
+    def is_subseteq(self, other: Address) -> bool:
+        if isinstance(other, AddressQubit):
+            return self.data == other.data
+        return False
+
+
+@final
+@dataclass
+class AddressWire(Address):
+    origin_qubit: AddressQubit
+
+    def is_subseteq(self, other: Address) -> bool:
+        if isinstance(other, AddressWire):
+            return self.origin_qubit == self.origin_qubit
+        return False

bloqade/noise/__init__.py

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

bloqade/noise/native/__init__.py

@@ -0,0 +1,20 @@
+"""NOTE: This module is not guaranteed to be supported long-term in bloqade. We will be
+moving towards a more general approach to noise modeling in the future."""
+
+from .model import (
+    GateNoiseParams as GateNoiseParams,
+    TwoRowZoneModel as TwoRowZoneModel,
+    MoveNoiseModelABC as MoveNoiseModelABC,
+)
+from .stmts import (
+    PauliChannel as PauliChannel,
+    CZPauliChannel as CZPauliChannel,
+    AtomLossChannel as AtomLossChannel,
+)
+from .rewrite import RemoveNoisePass as RemoveNoisePass
+from ._dialect import dialect as dialect
+from ._wrappers import (
+    pauli_channel as pauli_channel,
+    cz_pauli_channel as cz_pauli_channel,
+    atom_loss_channel as atom_loss_channel,
+)

bloqade/noise/native/_dialect.py

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

bloqade/noise/native/_wrappers.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.noise import native
+from bloqade.qasm2.types import Qubit
+
+
+@wraps(native.AtomLossChannel)
+def atom_loss_channel(
+    qargs: ilist.IList[Qubit, Any] | list, *, prob: float
+) -> None: ...
+
+
+@wraps(native.PauliChannel)
+def pauli_channel(
+    qargs: ilist.IList[Qubit, Any] | list, *, px: float, py: float, pz: float
+) -> None: ...
+
+
+@wraps(native.CZPauliChannel)
+def cz_pauli_channel(
+    ctrls: ilist.IList[Qubit, Any] | list,
+    qarg2: ilist.IList[Qubit, Any] | list,
+    *,
+    px_ctrl: float,
+    py_ctrl: float,
+    pz_ctrl: float,
+    px_qarg: float,
+    py_qarg: float,
+    pz_qarg: float,
+    paired: bool,
+) -> None: ...