bloqade-circuit 0.7.12__py3-none-any.whl → 0.8.0__py3-none-any.whl

bloqade/qubit/_prelude.py

@@ -0,0 +1,45 @@
+from typing import Annotated
+
+from kirin import ir
+from kirin.passes import Default
+from kirin.prelude import structural_no_opt
+from typing_extensions import Doc
+
+from . import _dialect as qubit
+
+
+@ir.dialect_group(structural_no_opt.union([qubit]))
+def kernel(self):
+    """Compile to a qubit kernel"""
+
+    def run_pass(
+        mt,
+        *,
+        verify: Annotated[
+            bool, Doc("run `verify` before running passes, default is `True`")
+        ] = True,
+        typeinfer: Annotated[
+            bool,
+            Doc(
+                "run type inference and apply the inferred type to IR, default `False`"
+            ),
+        ] = False,
+        fold: Annotated[bool, Doc("run folding passes, default is `True`")] = True,
+        aggressive: Annotated[
+            bool, Doc("run aggressive folding passes if `fold=True`")
+        ] = False,
+        no_raise: Annotated[
+            bool, Doc("do not raise exception during analysis, default is `True`")
+        ] = True,
+    ) -> None:
+        default_pass = Default(
+            self,
+            verify=verify,
+            fold=fold,
+            aggressive=aggressive,
+            typeinfer=typeinfer,
+            no_raise=no_raise,
+        )
+        default_pass.fixpoint(mt)
+
+    return run_pass

bloqade/qubit/analysis/__init__.py

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

bloqade/qubit/analysis/address_impl.py

@@ -0,0 +1,40 @@
+from kirin import interp
+from kirin.analysis import ForwardFrame
+
+from bloqade.analysis.address.lattice import (
+    Address,
+    AddressQubit,
+)
+from bloqade.analysis.address.analysis import AddressAnalysis
+
+from .. import stmts
+from .._dialect import dialect
+
+# 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
+
+
+@dialect.register(key="qubit.address")
+class SquinQubitMethodTable(interp.MethodTable):
+
+    @interp.impl(stmts.New)
+    def new_qubit(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: stmts.New,
+    ):
+        addr = AddressQubit(interp_.next_address)
+        interp_.next_address += 1
+        return (addr,)

bloqade/qubit/stdlib/__init__.py

@@ -0,0 +1,2 @@
+from . import simple as simple, broadcast as broadcast
+from ._new import new as new, qalloc as qalloc

bloqade/qubit/stdlib/_new.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from kirin.dialects import ilist
+
+from .. import _interface as qubit
+from .._prelude import kernel
+
+
+@kernel(typeinfer=True)
+def new() -> qubit.Qubit:
+    """Allocate a single new qubit
+
+    Returns:
+        (Qubit): The newly allocated qubit.
+    """
+    return qubit.new()
+
+
+# NOTE: this is a special case, that doesn't use the usual simple / broadcast semantics.
+@kernel(typeinfer=True)
+def qalloc(n_qubits: int) -> ilist.IList[qubit.Qubit, Any]:
+    """Allocate a new list of qubits.
+
+    Args:
+        n_qubits(int): The number of qubits to create.
+
+    Returns:
+        (ilist.IList[Qubit, n_qubits]) A list of qubits.
+    """
+
+    def _new(qid: int) -> qubit.Qubit:
+        return qubit.new()
+
+    return ilist.map(_new, ilist.range(n_qubits))

bloqade/qubit/stdlib/broadcast.py

@@ -0,0 +1,62 @@
+from typing import Any, TypeVar
+
+from kirin.dialects import ilist
+
+from bloqade.types import Qubit, MeasurementResult
+
+from .. import _interface as _qubit
+from .._prelude import kernel
+
+N = TypeVar("N", bound=int)
+
+
+@kernel
+def reset(qubits: ilist.IList[Qubit, Any]) -> None:
+    """
+    Reset a list of qubits to the zero state.
+
+    Args:
+        qubits (IList[Qubit, Any]): The list of qubits to reset.
+    """
+    _qubit.reset(qubits)
+
+
+@kernel
+def measure(qubits: ilist.IList[Qubit, N]) -> ilist.IList[MeasurementResult, N]:
+    """Measure a list of qubits.
+
+    Args:
+        qubits (IList[Qubit, N]): The list of qubits to measure.
+
+    Returns:
+        IList[MeasurementResult, N]: The list containing the results of the measurements.
+            A MeasurementResult can represent both 0 and 1 as well as atom loss.
+    """
+    return _qubit.measure(qubits)
+
+
+@kernel
+def get_qubit_id(qubits: ilist.IList[Qubit, N]) -> ilist.IList[int, N]:
+    """Get the global, unique ID of each qubit in the list.
+
+    Args:
+        qubits (IList[Qubit, N]): The list of qubits of which you want the ID.
+
+    Returns:
+        qubit_ids (IList[int, N]): The list of global, unique IDs of the qubits.
+    """
+    return _qubit.get_qubit_id(qubits)
+
+
+@kernel
+def get_measurement_id(
+    measurements: ilist.IList[MeasurementResult, N],
+) -> ilist.IList[int, N]:
+    """Get the global, unique ID of each of the measurement results in the list.
+
+    Args:
+        measurements (IList[MeasurementResult, N]): The previously taken measurement of which you want to know the ID.
+    Returns:
+        measurement_ids (IList[int, N]): The list of global, unique IDs of the measurements.
+    """
+    return _qubit.get_measurement_id(measurements)

bloqade/qubit/stdlib/simple.py

@@ -0,0 +1,59 @@
+from kirin.dialects import ilist
+
+from bloqade.types import Qubit, MeasurementResult
+
+from . import broadcast
+from .._prelude import kernel
+
+
+@kernel
+def reset(qubit: Qubit) -> None:
+    """
+    Reset a qubit to the zero state.
+
+    Args:
+        qubit (Qubit): The list qubit to reset.
+    """
+    return broadcast.reset(ilist.IList([qubit]))
+
+
+@kernel
+def measure(qubit: Qubit) -> MeasurementResult:
+    """Measure a qubit.
+
+    Args:
+        qubit (Qubit): The qubit to measure.
+
+    Returns:
+        MeasurementResult: The result of the measurement.
+            A MeasurementResult can represent both 0 and 1, but also atoms that are lost.
+    """
+    measurement_results = broadcast.measure(ilist.IList([qubit]))
+    return measurement_results[0]
+
+
+@kernel
+def get_qubit_id(qubit: Qubit) -> int:
+    """Get the global, unique ID of the qubit.
+
+    Args:
+        qubit (Qubit): The qubit of which you want the ID.
+
+    Returns:
+        qubit_id (int): The global, unique ID of the qubit.
+    """
+    ids = broadcast.get_qubit_id(ilist.IList([qubit]))
+    return ids[0]
+
+
+@kernel
+def get_measurement_id(measurement: MeasurementResult) -> int:
+    """Get the global, unique ID of the measurement result.
+
+    Args:
+        measurement (MeasurementResult): The previously taken measurement of which you want to know the ID.
+    Returns:
+        measurement_id (int): The global, unique ID of the measurement.
+    """
+    ids = broadcast.get_measurement_id(ilist.IList([measurement]))
+    return ids[0]

bloqade/qubit/stmts.py

@@ -0,0 +1,60 @@
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.dialects import ilist
+
+from bloqade.types import QubitType, MeasurementResultType
+
+from ._dialect import dialect
+
+
+@statement(dialect=dialect)
+class New(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    result: ir.ResultValue = info.result(QubitType)
+
+
+Len = types.TypeVar("Len", bound=types.Int)
+
+
+@statement(dialect=dialect)
+class Measure(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, Len])
+    result: ir.ResultValue = info.result(ilist.IListType[MeasurementResultType, Len])
+
+
+@statement(dialect=dialect)
+class QubitId(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, Len])
+    result: ir.ResultValue = info.result(ilist.IListType[types.Int, Len])
+
+
+@statement(dialect=dialect)
+class MeasurementId(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    measurements: ir.SSAValue = info.argument(
+        ilist.IListType[MeasurementResultType, Len]
+    )
+    result: ir.ResultValue = info.result(ilist.IListType[types.Int, Len])
+
+
+@statement(dialect=dialect)
+class Reset(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
+# TODO: investigate why this is needed to get type inference to be correct.
+@dialect.register(key="typeinfer")
+class __TypeInfer(interp.MethodTable):
+    @interp.impl(Measure)
+    def measure_list(self, _interp, frame: interp.AbstractFrame, stmt: Measure):
+        qubit_type = frame.get(stmt.qubits)
+
+        if isinstance(qubit_type, types.Generic):
+            len_type = qubit_type.vars[1]
+        else:
+            len_type = types.Any
+
+        return (ilist.IListType[MeasurementResultType, len_type],)

bloqade/rewrite/passes/aggressive_unroll.py

@@ -38,6 +38,7 @@ class Fold(Pass):
             InlineGetField(),
             InlineGetItem(),
             ilist.rewrite.InlineGetItem(),
+            ilist.rewrite.FlattenAdd(),
             ilist.rewrite.HintLen(),
         )
         result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
@@ -68,7 +69,7 @@ class AggressiveUnroll(Pass):
             .rewrite(mt.code)
             .join(result)
         )
-        result = self.typeinfer.unsafe_run(mt).join(result)
+        self.typeinfer.unsafe_run(mt)
         result = self.fold.unsafe_run(mt).join(result)
         result = Walk(Inline(self.inline_heuristic)).rewrite(mt.code).join(result)
         result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)

bloqade/squin/__init__.py

@@ -1,24 +1,51 @@
 from . import (
-    op as op,
-    wire as wire,
+    gate as gate,
     noise as noise,
-    qubit as qubit,
     analysis as analysis,
-    lowering as lowering,
-    _typeinfer as _typeinfer,
 )
-from .groups import wired as wired, kernel as kernel
+from .. import qubit as qubit
+from ..qubit import (
+    reset as reset,
+    qalloc as qalloc,
+    measure as measure,
+    get_qubit_id as get_qubit_id,
+    get_measurement_id as get_measurement_id,
+)
+from .groups import kernel as kernel
+from .stdlib.simple import (
+    h as h,
+    s as s,
+    t as t,
+    x as x,
+    y as y,
+    z as z,
+    cx as cx,
+    cy as cy,
+    cz as cz,
+    rx as rx,
+    ry as ry,
+    rz as rz,
+    u3 as u3,
+    s_adj as s_adj,
+    shift as shift,
+    t_adj as t_adj,
+    sqrt_x as sqrt_x,
+    sqrt_y as sqrt_y,
+    sqrt_z as sqrt_z,
+    bit_flip as bit_flip,
+    depolarize as depolarize,
+    qubit_loss as qubit_loss,
+    sqrt_x_adj as sqrt_x_adj,
+    sqrt_y_adj as sqrt_y_adj,
+    sqrt_z_adj as sqrt_z_adj,
+    depolarize2 as depolarize2,
+    correlated_qubit_loss as correlated_qubit_loss,
+    two_qubit_pauli_channel as two_qubit_pauli_channel,
+    single_qubit_pauli_channel as single_qubit_pauli_channel,
+)
 
 # NOTE: it's important to keep these imports here since they import squin.kernel
 # we skip isort here
-from . import parallel as parallel  # isort: skip
-from .stdlib import gate as gate, channel as channel  # isort: skip
-
-try:
-    # NOTE: make sure optional cirq dependency is installed
-    import cirq as cirq_package  # noqa: F401
-except ImportError:
-    pass
-else:
-    from . import cirq as cirq
-    from .cirq import load_circuit as load_circuit
+from .stdlib import (  # isort: skip
+    broadcast as broadcast,
+)

bloqade/squin/analysis/__init__.py

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

bloqade/squin/analysis/schedule.py

@@ -210,8 +210,8 @@ class DagScheduleAnalysis(Forward[GateSchedule]):
                 if old_stmt is not None:
                     self.stmt_dag.add_edge(old_stmt, stmt)
                 self.use_def[idx] = stmt
-        elif isinstance(addr, address.AddressTuple):
-            for sub_addr in addr.data:
+        elif isinstance(addr, address.AddressReg):
+            for sub_addr in addr.qubits:
                 self._update_dag(stmt, sub_addr)
 
     def update_dag(self, stmt: ir.Statement, args: Sequence[ir.SSAValue]):

bloqade/squin/gate/__init__.py

@@ -0,0 +1,2 @@
+from . import stmts as stmts
+from ._dialect import dialect as dialect

bloqade/squin/gate/_dialect.py

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

bloqade/squin/gate/_interface.py

@@ -0,0 +1,98 @@
+from typing import Any, TypeVar
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.types import Qubit
+
+from .stmts import (
+    CX,
+    CY,
+    CZ,
+    U3,
+    H,
+    S,
+    T,
+    X,
+    Y,
+    Z,
+    Rx,
+    Ry,
+    Rz,
+    SqrtX,
+    SqrtY,
+)
+
+
+@wraps(X)
+def x(qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(Y)
+def y(qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(Z)
+def z(qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(H)
+def h(qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(T)
+def t(qubits: ilist.IList[Qubit, Any], *, adjoint: bool) -> None: ...
+
+
+@wraps(S)
+def s(qubits: ilist.IList[Qubit, Any], *, adjoint: bool) -> None: ...
+
+
+@wraps(SqrtX)
+def sqrt_x(qubits: ilist.IList[Qubit, Any], *, adjoint: bool) -> None: ...
+
+
+@wraps(SqrtY)
+def sqrt_y(qubits: ilist.IList[Qubit, Any], *, adjoint: bool) -> None: ...
+
+
+@wraps(Rx)
+def rx(angle: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(Ry)
+def ry(angle: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(Rz)
+def rz(angle: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+Len = TypeVar("Len", bound=int)
+
+
+@wraps(CX)
+def cx(
+    controls: ilist.IList[Qubit, Len],
+    targets: ilist.IList[Qubit, Len],
+) -> None: ...
+
+
+@wraps(CY)
+def cy(
+    controls: ilist.IList[Qubit, Len],
+    targets: ilist.IList[Qubit, Len],
+) -> None: ...
+
+
+@wraps(CZ)
+def cz(
+    controls: ilist.IList[Qubit, Len],
+    targets: ilist.IList[Qubit, Len],
+) -> None: ...
+
+
+@wraps(U3)
+def u3(
+    theta: float, phi: float, lam: float, qubits: ilist.IList[Qubit, Any]
+) -> None: ...

bloqade/squin/gate/stmts.py

@@ -0,0 +1,119 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+from kirin.dialects import ilist
+
+from bloqade.types import QubitType
+
+from ._dialect import dialect
+
+
+@statement
+class SingleQubitGate(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
+@statement(dialect=dialect)
+class X(SingleQubitGate):
+    pass
+
+
+@statement(dialect=dialect)
+class Y(SingleQubitGate):
+    pass
+
+
+@statement(dialect=dialect)
+class Z(SingleQubitGate):
+    pass
+
+
+@statement(dialect=dialect)
+class H(SingleQubitGate):
+    pass
+
+
+@statement
+class SingleQubitNonHermitianGate(SingleQubitGate):
+    adjoint: bool = info.attribute(default=False)
+
+
+@statement(dialect=dialect)
+class T(SingleQubitNonHermitianGate):
+    pass
+
+
+@statement(dialect=dialect)
+class S(SingleQubitNonHermitianGate):
+    pass
+
+
+@statement(dialect=dialect)
+class SqrtX(SingleQubitNonHermitianGate):
+    pass
+
+
+@statement(dialect=dialect)
+class SqrtY(SingleQubitNonHermitianGate):
+    pass
+
+
+@statement
+class RotationGate(ir.Statement):
+    # NOTE: don't inherit from SingleQubitGate here so the wrapper doesn't have qubits as first arg
+    traits = frozenset({lowering.FromPythonCall()})
+    angle: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
+@statement(dialect=dialect)
+class Rx(RotationGate):
+    pass
+
+
+@statement(dialect=dialect)
+class Ry(RotationGate):
+    pass
+
+
+@statement(dialect=dialect)
+class Rz(RotationGate):
+    pass
+
+
+N = types.TypeVar("N", bound=types.Int)
+
+
+@statement
+class ControlledGate(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    controls: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+    targets: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+
+
+@statement(dialect=dialect)
+class CX(ControlledGate):
+    name = "cx"
+    pass
+
+
+@statement(dialect=dialect)
+class CY(ControlledGate):
+    name = "cy"
+    pass
+
+
+@statement(dialect=dialect)
+class CZ(ControlledGate):
+    name = "cz"
+    pass
+
+
+@statement(dialect=dialect)
+class U3(ir.Statement):
+    # NOTE: don't inherit from SingleQubitGate here so the wrapper doesn't have qubits as first arg
+    traits = frozenset({lowering.FromPythonCall()})
+    theta: ir.SSAValue = info.argument(types.Float)
+    phi: ir.SSAValue = info.argument(types.Float)
+    lam: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])

bloqade/squin/groups.py

@@ -1,31 +1,24 @@
 from kirin import ir, passes
 from kirin.prelude import structural_no_opt
-from kirin.rewrite import Walk, Chain
 from kirin.dialects import debug, ilist
 
-from . import op, wire, noise, qubit
-from .op.rewrite import PyMultToSquinMult
-from .rewrite.desugar import ApplyDesugarRule, MeasureDesugarRule
+from . import gate, noise
+from .. import qubit
 
 
-@ir.dialect_group(structural_no_opt.union([op, qubit, noise, debug]))
+@ir.dialect_group(structural_no_opt.union([qubit, noise, gate, debug]))
 def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
     ilist_desugar_pass = ilist.IListDesugar(self)
-    desugar_pass = Walk(Chain(MeasureDesugarRule(), ApplyDesugarRule()))
-    py_mult_to_mult_pass = PyMultToSquinMult(self)
 
     def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
         method.verify()
         if fold:
             fold_pass.fixpoint(method)
 
-        py_mult_to_mult_pass(method)
-
         if typeinfer:
-            typeinfer_pass(method)
-            desugar_pass.rewrite(method.code)
+            typeinfer_pass(method)  # infer types before desugaring
 
         ilist_desugar_pass(method)
 
@@ -34,13 +27,3 @@ def kernel(self):
             method.verify_type()
 
     return run_pass
-
-
-@ir.dialect_group(structural_no_opt.union([op, wire, noise]))
-def wired(self):
-    py_mult_to_mult_pass = PyMultToSquinMult(self)
-
-    def run_pass(method):
-        py_mult_to_mult_pass(method)
-
-    return run_pass