bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/squin/op/stmts.py

@@ -0,0 +1,220 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from .types import OpType
+from .traits import Unitary, HasSites, FixedSites, MaybeUnitary
+from .complex import Complex
+from ._dialect import dialect
+
+
+@statement
+class Operator(ir.Statement):
+    pass
+
+
+@statement
+class PrimitiveOp(Operator):
+    pass
+
+
+@statement
+class CompositeOp(Operator):
+    pass
+
+
+@statement
+class BinaryOp(CompositeOp):
+    lhs: ir.SSAValue = info.argument(OpType)
+    rhs: ir.SSAValue = info.argument(OpType)
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement(dialect=dialect)
+class Kron(BinaryOp):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
+    is_unitary: bool = info.attribute(default=False)
+
+
+@statement(dialect=dialect)
+class Mult(BinaryOp):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
+    is_unitary: bool = info.attribute(default=False)
+
+
+@statement(dialect=dialect)
+class Adjoint(CompositeOp):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
+    is_unitary: bool = info.attribute(default=False)
+    op: ir.SSAValue = info.argument(OpType)
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement(dialect=dialect)
+class Scale(CompositeOp):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
+    is_unitary: bool = info.attribute(default=False)
+    op: ir.SSAValue = info.argument(OpType)
+    factor: ir.SSAValue = info.argument(Complex)
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement(dialect=dialect)
+class Control(CompositeOp):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
+    is_unitary: bool = info.attribute(default=False)
+    op: ir.SSAValue = info.argument(OpType)
+    n_controls: int = info.attribute()
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement(dialect=dialect)
+class Rot(CompositeOp):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary()})
+    axis: ir.SSAValue = info.argument(OpType)
+    angle: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement(dialect=dialect)
+class Identity(CompositeOp):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), HasSites()})
+    sites: int = info.attribute()
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement
+class ConstantOp(PrimitiveOp):
+    traits = frozenset(
+        {ir.Pure(), lowering.FromPythonCall(), ir.ConstantLike(), FixedSites(1)}
+    )
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement
+class ConstantUnitary(ConstantOp):
+    traits = frozenset(
+        {
+            ir.Pure(),
+            lowering.FromPythonCall(),
+            ir.ConstantLike(),
+            Unitary(),
+            FixedSites(1),
+        }
+    )
+
+
+@statement(dialect=dialect)
+class PhaseOp(PrimitiveOp):
+    """
+    A phase operator.
+
+    $$
+    PhaseOp(theta) = e^{i \theta} I
+    $$
+    """
+
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), FixedSites(1)})
+    theta: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement(dialect=dialect)
+class ShiftOp(PrimitiveOp):
+    """
+    A phase shift operator.
+
+    $$
+    Shift(theta) = \\begin{bmatrix} 1 & 0 \\\\ 0 & e^{i \\theta} \\end{bmatrix}
+    $$
+    """
+
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), FixedSites(1)})
+    theta: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(OpType)
+
+
+@statement
+class PauliOp(ConstantUnitary):
+    pass
+
+
+@statement(dialect=dialect)
+class X(PauliOp):
+    pass
+
+
+@statement(dialect=dialect)
+class Y(PauliOp):
+    pass
+
+
+@statement(dialect=dialect)
+class Z(PauliOp):
+    pass
+
+
+@statement(dialect=dialect)
+class H(ConstantUnitary):
+    pass
+
+
+@statement(dialect=dialect)
+class S(ConstantUnitary):
+    pass
+
+
+@statement(dialect=dialect)
+class T(ConstantUnitary):
+    pass
+
+
+@statement(dialect=dialect)
+class P0(ConstantOp):
+    """
+    The $P_0$ projection operator.
+
+    $$
+    P0 = \\begin{bmatrix} 1 & 0 \\\\ 0 & 0 \\end{bmatrix}
+    $$
+    """
+
+    pass
+
+
+@statement(dialect=dialect)
+class P1(ConstantOp):
+    """
+    The $P_1$ projection operator.
+
+    $$
+    P1 = \\begin{bmatrix} 0 & 0 \\\\ 0 & 1 \\end{bmatrix}
+    $$
+    """
+
+    pass
+
+
+@statement(dialect=dialect)
+class Sn(ConstantOp):
+    """
+    $S_{-}$ operator.
+
+    $$
+    Sn = \\frac{1}{2} (S_x - i S_y) = \\frac{1}{2} \\begin{bmatrix} 0 & 0 \\\\ 1 & 0 \\end{bmatrix}
+    $$
+    """
+
+    pass
+
+
+@statement(dialect=dialect)
+class Sp(ConstantOp):
+    """
+    $S_{+}$ operator.
+
+    $$
+    Sp = \\frac{1}{2} (S_x + i S_y) = \\frac{1}{2}\\begin{bmatrix} 0 & 1 \\\\ 0 & 0 \\end{bmatrix}
+    $$
+    """
+
+    pass

bloqade/squin/op/traits.py

@@ -0,0 +1,43 @@
+from typing import cast
+from dataclasses import dataclass
+
+from kirin import ir
+
+
+@dataclass(frozen=True)
+class FixedSites(ir.StmtTrait):
+    data: int
+
+
+@dataclass(frozen=True)
+class HasSites(ir.StmtTrait):
+    """An operator with a `sites` attribute."""
+
+    def get_sites(self, stmt: ir.Statement):
+        attr = stmt.get_attr_or_prop("sites")
+        if attr is None:
+            raise ValueError(f"Missing sites attribute in {stmt}")
+        return cast(ir.PyAttr[int], attr).data
+
+    def set_sites(self, stmt: ir.Statement, value: int):
+        stmt.attributes["sites"] = ir.PyAttr(value)
+        return
+
+
+@dataclass(frozen=True)
+class Unitary(ir.StmtTrait):
+    pass
+
+
+@dataclass(frozen=True)
+class MaybeUnitary(ir.StmtTrait):
+
+    def is_unitary(self, stmt: ir.Statement):
+        attr = stmt.get_attr_or_prop("is_unitary")
+        if attr is None:
+            return False
+        return cast(ir.PyAttr[bool], attr).data
+
+    def set_unitary(self, stmt: ir.Statement, value: bool):
+        stmt.attributes["is_unitary"] = ir.PyAttr(value)
+        return

bloqade/squin/op/types.py

@@ -0,0 +1,10 @@
+from kirin import types
+
+
+class Op:
+
+    def __matmul__(self, other: "Op") -> "Op":
+        raise NotImplementedError("@ can only be used within a squin kernel program")
+
+
+OpType = types.PyClass(Op)

bloqade/squin/qubit.py

@@ -0,0 +1,118 @@
+"""qubit dialect for squin language.
+
+This dialect defines the operations that can be performed on qubits.
+
+Depends on:
+- `bloqade.squin.op`: provides the `OpType` type and semantics for operators applied to qubits.
+- `kirin.dialects.ilist`: provides the `ilist.IListType` type for lists of qubits.
+"""
+
+from typing import Any
+
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.types import Qubit, QubitType
+from bloqade.squin.op.types import Op, OpType
+
+dialect = ir.Dialect("squin.qubit")
+
+
+@statement(dialect=dialect)
+class New(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    n_qubits: ir.SSAValue = info.argument(types.Int)
+    result: ir.ResultValue = info.result(ilist.IListType[QubitType, types.Any])
+
+
+@statement(dialect=dialect)
+class Apply(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    operator: ir.SSAValue = info.argument(OpType)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+
+
+@statement(dialect=dialect)
+class Measure(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+    result: ir.ResultValue = info.result(types.Int)
+
+
+@statement(dialect=dialect)
+class MeasureAndReset(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+    result: ir.ResultValue = info.result(types.Int)
+
+
+@statement(dialect=dialect)
+class Reset(ir.Statement):
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+
+
+# NOTE: no dependent types in Python, so we have to mark it Any...
+@wraps(New)
+def new(n_qubits: int) -> ilist.IList[Qubit, Any]:
+    """Create 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.
+    """
+    ...
+
+
+@wraps(Apply)
+def apply(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> None:
+    """Apply an operator to a list of qubits.
+
+    Args:
+        operator: The operator to apply.
+        qubits: The list of qubits to apply the operator to. The size of the list
+            must be inferable and match the number of qubits expected by the operator.
+
+    Returns:
+        None
+    """
+    ...
+
+
+@wraps(Measure)
+def measure(qubits: ilist.IList[Qubit, Any]) -> int:
+    """Measure the qubits in the list."
+
+    Args:
+        qubits: The list of qubits to measure.
+
+    Returns:
+        int: The result of the measurement.
+    """
+    ...
+
+
+@wraps(MeasureAndReset)
+def measure_and_reset(qubits: ilist.IList[Qubit, Any]) -> int:
+    """Measure the qubits in the list and reset them."
+
+    Args:
+        qubits: The list of qubits to measure and reset.
+
+    Returns:
+        int: The result of the measurement.
+    """
+    ...
+
+
+@wraps(Reset)
+def reset(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Reset the qubits in the list."
+
+    Args:
+        qubits: The list of qubits to reset.
+    """
+    ...

bloqade/squin/wire.py

@@ -0,0 +1,103 @@
+"""A NVIDIA QUAKE-like wire dialect.
+
+This dialect is expected to be used in combination with the operator dialect
+as an intermediate representation for analysis and optimization of quantum
+circuits. Thus we do not define wrapping functions for the statements in this
+dialect.
+"""
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+
+from bloqade.types import QubitType
+
+from .op.types import OpType
+
+# from kirin.lowering import wraps
+
+# from .op.types import Op, OpType
+
+dialect = ir.Dialect("squin.wire")
+
+
+class WireTerminator(ir.StmtTrait):
+    pass
+
+
+class Wire:
+    pass
+
+
+WireType = types.PyClass(Wire)
+
+
+# no return value for `wrap`
+@statement(dialect=dialect)
+class Wrap(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
+    wire: ir.SSAValue = info.argument(WireType)
+    qubit: ir.SSAValue = info.argument(QubitType)
+
+
+# "Unwrap the quantum references to expose wires" -> From Quake Dialect documentation
+# Unwrap(Qubit) -> Wire
+@statement(dialect=dialect)
+class Unwrap(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    qubit: ir.SSAValue = info.argument(QubitType)
+    result: ir.ResultValue = info.result(WireType)
+
+
+# In Quake, you put a wire in and get a wire out when you "apply" an operator
+# In this case though we just need to indicate that an operator is applied to list[wires]
+@statement(dialect=dialect)
+class Apply(ir.Statement):  # apply(op, w1, w2, ...)
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    operator: ir.SSAValue = info.argument(OpType)
+    inputs: tuple[ir.SSAValue, ...] = info.argument(WireType)
+
+    def __init__(self, operator: ir.SSAValue, *args: ir.SSAValue):
+        result_types = tuple(WireType for _ in args)
+        super().__init__(
+            args=(operator,) + args,
+            result_types=result_types,  # result types of the Apply statement, should all be WireTypes
+            args_slice={
+                "operator": 0,
+                "inputs": slice(1, None),
+            },  # pretty printing + syntax sugar
+        )  # custom lowering required for wrapper to work here
+
+
+# NOTE: measurement cannot be pure because they will collapse the state
+#       of the qubit. The state is a hidden state that is not visible to
+#      the user in the wire dialect.
+@statement(dialect=dialect)
+class Measure(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
+    wire: ir.SSAValue = info.argument(WireType)
+    result: ir.ResultValue = info.result(types.Int)
+
+
+@statement(dialect=dialect)
+class MeasureAndReset(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
+    wire: ir.SSAValue = info.argument(WireType)
+    result: ir.ResultValue = info.result(types.Int)
+    out_wire: ir.ResultValue = info.result(WireType)
+
+
+@statement(dialect=dialect)
+class Reset(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
+    wire: ir.SSAValue = info.argument(WireType)
+
+
+# Issue where constant propagation can't handle
+# multiple return values from Apply properly
+@dialect.register(key="constprop")
+class ConstPropWire(interp.MethodTable):
+
+    @interp.impl(Apply)
+    def apply(self, interp, frame, stmt: Apply):
+
+        return frame.get_values(stmt.inputs)

bloqade/stim/__init__.py

@@ -0,0 +1,6 @@
+from .groups import main as main
+from ._wrappers import *  # noqa: F403
+from .dialects.aux import *  # noqa F403
+from .dialects.gate import *  # noqa F403
+from .dialects.noise import *  # noqa F403
+from .dialects.collapse import *  # noqa F403

bloqade/stim/_wrappers.py

@@ -0,0 +1,186 @@
+from typing import Union
+
+from kirin.lowering import wraps
+
+from .dialects import aux, gate, noise, collapse
+
+
+# dialect:: gate
+## 1q
+@wraps(gate.X)
+def x(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+@wraps(gate.Y)
+def y(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+@wraps(gate.Z)
+def z(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+@wraps(gate.Identity)
+def identity(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+@wraps(gate.H)
+def h(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+@wraps(gate.S)
+def s(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+@wraps(gate.SqrtX)
+def sqrt_x(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+@wraps(gate.SqrtY)
+def sqrt_y(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+@wraps(gate.SqrtZ)
+def sqrt_z(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+## clif 2q
+@wraps(gate.Swap)
+def swap(targets: tuple[int, ...], dagger: bool = False) -> None: ...
+
+
+## ctrl 2q
+@wraps(gate.CX)
+def cx(
+    controls: tuple[int, ...], targets: tuple[int, ...], dagger: bool = False
+) -> None: ...
+
+
+@wraps(gate.CY)
+def cy(
+    controls: tuple[int, ...], targets: tuple[int, ...], dagger: bool = False
+) -> None: ...
+
+
+@wraps(gate.CZ)
+def cz(
+    controls: tuple[int, ...], targets: tuple[int, ...], dagger: bool = False
+) -> None: ...
+
+
+## pp
+@wraps(gate.SPP)
+def spp(targets: tuple[aux.PauliString, ...], dagger=False) -> None: ...
+
+
+# dialect:: aux
+@wraps(aux.GetRecord)
+def rec(id: int) -> aux.RecordResult: ...
+
+
+@wraps(aux.Detector)
+def detector(
+    coord: tuple[Union[int, float], ...], targets: tuple[aux.RecordResult, ...]
+) -> None: ...
+
+
+@wraps(aux.ObservableInclude)
+def observable_include(idx: int, targets: tuple[aux.RecordResult, ...]) -> None: ...
+
+
+@wraps(aux.Tick)
+def tick() -> None: ...
+
+
+@wraps(aux.NewPauliString)
+def pauli_string(
+    string: tuple[str, ...], flipped: tuple[bool, ...], targets: tuple[int, ...]
+) -> aux.PauliString: ...
+
+
+# dialect:: collapse
+@wraps(collapse.MZ)
+def mz(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(collapse.MY)
+def my(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(collapse.MX)
+def mx(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(collapse.MZZ)
+def mzz(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(collapse.MYY)
+def myy(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(collapse.MXX)
+def mxx(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(collapse.PPMeasurement)
+def mpp(p: float, targets: tuple[aux.PauliString, ...]) -> None: ...
+
+
+@wraps(collapse.RZ)
+def rz(targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(collapse.RY)
+def ry(targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(collapse.RX)
+def rx(targets: tuple[int, ...]) -> None: ...
+
+
+# dialect:: noise
+@wraps(noise.Depolarize1)
+def depolarize1(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.Depolarize2)
+def depolarize2(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.PauliChannel1)
+def pauli_channel1(
+    px: float, py: float, pz: float, targets: tuple[int, ...]
+) -> None: ...
+
+
+@wraps(noise.PauliChannel2)
+def pauli_channel2(
+    pix: float,
+    piy: float,
+    piz: float,
+    pxi: float,
+    pxx: float,
+    pxy: float,
+    pxz: float,
+    pyi: float,
+    pyx: float,
+    pyy: float,
+    pyz: float,
+    pzi: float,
+    pzx: float,
+    pzy: float,
+    pzz: float,
+    targets: tuple[int, ...],
+) -> None: ...
+
+
+@wraps(noise.XError)
+def x_error(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.YError)
+def y_error(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.ZError)
+def z_error(p: float, targets: tuple[int, ...]) -> None: ...

bloqade/stim/dialects/__init__.py

@@ -0,0 +1,5 @@
+from . import aux as aux, gate as gate, noise as noise, collapse as collapse
+from .aux.stmts import *  # noqa F403
+from .gate.stmts import *  # noqa F403
+from .noise.stmts import *  # noqa F403
+from .collapse.stmts import *  # noqa F403

bloqade/stim/dialects/aux/__init__.py

@@ -0,0 +1,11 @@
+from . import lowering as lowering
+from .emit import EmitStimAuxMethods as EmitStimAuxMethods
+from .stmts import *  # noqa F403
+from .types import (
+    RecordType as RecordType,
+    PauliString as PauliString,
+    RecordResult as RecordResult,
+    PauliStringType as PauliStringType,
+)
+from .interp import StimAuxMethods as StimAuxMethods
+from ._dialect import dialect as dialect

bloqade/stim/dialects/aux/_dialect.py

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