bloqade-circuit 0.6.8__py3-none-any.whl → 0.7.1__py3-none-any.whl

bloqade/squin/cirq/emit/op.py

@@ -2,6 +2,7 @@ import math
 
 import cirq
 import numpy as np
+from kirin.emit import EmitError
 from kirin.interp import MethodTable, impl
 
 from ... import op
@@ -9,11 +10,11 @@ from .runtime import (
     SnRuntime,
     SpRuntime,
     U3Runtime,
-    RotRuntime,
     KronRuntime,
     MultRuntime,
     ScaleRuntime,
     AdjointRuntime,
+    BasicOpRuntime,
     ControlRuntime,
     UnitaryRuntime,
     HermitianRuntime,
@@ -117,7 +118,7 @@ class EmitCirqOpMethods(MethodTable):
 
     @impl(op.stmts.Reset)
     def reset(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Reset):
-        return (HermitianRuntime(cirq.ResetChannel()),)
+        return (BasicOpRuntime(cirq.ResetChannel()),)
 
     @impl(op.stmts.PauliString)
     def pauli_string(
@@ -127,11 +128,42 @@ class EmitCirqOpMethods(MethodTable):
 
     @impl(op.stmts.Rot)
     def rot(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Rot):
-        axis_op: HermitianRuntime = frame.get(stmt.axis)
+        axis: OperatorRuntimeABC = frame.get(stmt.axis)
+
+        if not isinstance(axis, HermitianRuntime):
+            raise EmitError(
+                f"Circuit emission only supported for Pauli operators! Got axis {axis}"
+            )
+
         angle = frame.get(stmt.angle)
 
-        axis_name = str(axis_op.gate).lower()
-        return (RotRuntime(axis=axis_name, angle=angle),)
+        match axis.gate:
+            case cirq.X:
+                gate = cirq.Rx(rads=angle)
+            case cirq.Y:
+                gate = cirq.Ry(rads=angle)
+            case cirq.Z:
+                gate = cirq.Rz(rads=angle)
+            case _:
+                raise EmitError(
+                    f"Circuit emission only supported for Pauli operators! Got axis {axis.gate}"
+                )
+
+        return (HermitianRuntime(gate=gate),)
+
+    @impl(op.stmts.ResetToOne)
+    def reset_to_one(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.ResetToOne
+    ):
+        # NOTE: just apply a reset to 0 and flip in sequence (we re-use the multiplication runtime since it does exactly that)
+        gate1 = cirq.ResetChannel()
+        gate2 = cirq.X
+
+        rt1 = BasicOpRuntime(gate1)
+        rt2 = HermitianRuntime(gate2)
+
+        # NOTE: mind the order: rhs is applied first
+        return (MultRuntime(rt2, rt1),)
 
     @impl(op.stmts.SqrtX)
     def sqrt_x(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.SqrtX):

bloqade/squin/cirq/emit/qubit.py

@@ -25,7 +25,7 @@ class EmitCirqQubitMethods(MethodTable):
     @impl(qubit.Apply)
     def apply(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Apply):
         op: OperatorRuntimeABC = frame.get(stmt.operator)
-        qbits = frame.get(stmt.qubits)
+        qbits = [frame.get(qbit) for qbit in stmt.qubits]
         operations = op.apply(qbits)
         for operation in operations:
             frame.circuit.append(operation)
@@ -34,11 +34,11 @@ class EmitCirqQubitMethods(MethodTable):
     @impl(qubit.Broadcast)
     def broadcast(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Broadcast):
         op = frame.get(stmt.operator)
-        qbits = frame.get(stmt.qubits)
+        qbit_lists = [frame.get(qbit) for qbit in stmt.qubits]
 
         cirq_ops = []
-        for qbit in qbits:
-            cirq_ops.extend(op.apply([qbit]))
+        for qbits in zip(*qbit_lists):
+            cirq_ops.extend(op.apply(qbits))
 
         frame.circuit.append(cirq.Moment(cirq_ops))
         return ()

bloqade/squin/cirq/emit/runtime.py

@@ -240,18 +240,3 @@ class PauliStringRuntime(OperatorRuntimeABC):
             qbit: pauli_label for (qbit, pauli_label) in zip(qubits, self.string)
         }
         return [cirq.PauliString(pauli_mapping)]
-
-
-@dataclass
-class RotRuntime(OperatorRuntimeABC):
-    axis: str
-    angle: float
-
-    def num_qubits(self) -> int:
-        return 1
-
-    def unsafe_apply(
-        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
-    ) -> list[cirq.Operation]:
-        rot = getattr(cirq, "R" + self.axis.lower())(rads=self.angle)
-        return [rot(*qubits)]

bloqade/squin/cirq/lowering.py

@@ -44,14 +44,7 @@ class Squin(lowering.LoweringABC[CirqNode]):
         self, state: lowering.State[CirqNode], qids: list[cirq.Qid]
     ):
         qbits_getitem = [self.lower_qubit_getindex(state, qid) for qid in qids]
-        qbits_stmt = ilist.New(values=qbits_getitem, elem_type=qubit.QubitType)
-        qbits_result = state.current_frame.get(qbits_stmt.name)
-
-        if qbits_result is not None:
-            return qbits_result
-
-        state.current_frame.push(qbits_stmt)
-        return qbits_stmt.result
+        return tuple(qbits_getitem)
 
     def run(
         self,
@@ -159,7 +152,8 @@ class Squin(lowering.LoweringABC[CirqNode]):
             stmt = state.current_frame.push(qubit.MeasureQubit(qbit))
         else:
             qbits = self.lower_qubit_getindices(state, node.qubits)
-            stmt = state.current_frame.push(qubit.MeasureQubitList(qbits))
+            qbits_list = state.current_frame.push(ilist.New(values=qbits))
+            stmt = state.current_frame.push(qubit.MeasureQubitList(qbits_list.result))
 
         key = node.gate.key
         if isinstance(key, cirq.MeasurementKey):

bloqade/squin/gate.py

@@ -137,6 +137,13 @@ def reset(qubit: Qubit) -> None:
     _qubit.apply(op, qubit)
 
 
+@kernel
+def reset_to_one(qubit: Qubit) -> None:
+    """Reset qubit to 1."""
+    op = _op.reset_to_one()
+    _qubit.apply(op, qubit)
+
+
 @kernel
 def cx(control: Qubit, target: Qubit) -> None:
     """Controlled x gate applied to control and target"""

bloqade/squin/lowering.py

@@ -52,3 +52,29 @@ class ApplyAnyCallLowering(lowering.FromPythonCall["qubit.ApplyAny"]):
             return op, qubits.elts
 
         return op, [qubits]
+
+
+@dataclass(frozen=True)
+class BroadcastCallLowering(lowering.FromPythonCall["qubit.Broadcast"]):
+    """
+    Custom lowering for broadcast vararg call.
+
+    NOTE: we can re-use this to lower Apply too once we remove the deprecated syntax
+    """
+
+    def lower(
+        self, stmt: type["qubit.Broadcast"], state: lowering.State, node: ast.Call
+    ):
+        if len(node.args) < 2:
+            raise lowering.BuildError(
+                "Broadcast requires at least one operator and one qubit list argument"
+            )
+
+        op, *qubit_lists = node.args
+
+        op_lowered = state.lower(op).expect_one()
+        qubits_lists_lowered = [
+            state.lower(qubit_list).expect_one() for qubit_list in qubit_lists
+        ]
+
+        return state.current_frame.push(stmt(op_lowered, tuple(qubits_lists_lowered)))

bloqade/squin/noise/__init__.py

@@ -1,7 +1,6 @@
 from . import stmts as stmts
 from ._dialect import dialect as dialect
 from ._wrapper import (
-    pp_error as pp_error,
     depolarize as depolarize,
     qubit_loss as qubit_loss,
     depolarize2 as depolarize2,

bloqade/squin/noise/_wrapper.py

@@ -3,17 +3,13 @@ from typing import Literal
 from kirin.dialects import ilist
 from kirin.lowering import wraps
 
-from bloqade.squin.op.types import Op
+from bloqade.squin.op.types import Op, MultiQubitPauliOp
 
 from . import stmts
 
 
 @wraps(stmts.PauliError)
-def pauli_error(basis: Op, p: float) -> Op: ...
-
-
-@wraps(stmts.PPError)
-def pp_error(op: Op, p: float) -> Op: ...
+def pauli_error(basis: MultiQubitPauliOp, p: float) -> Op: ...
 
 
 @wraps(stmts.Depolarize)

bloqade/squin/noise/rewrite.py

@@ -7,7 +7,6 @@ from kirin.dialects import py, ilist
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from .stmts import (
-    PPError,
     QubitLoss,
     Depolarize,
     PauliError,
@@ -86,16 +85,6 @@ class _RewriteNoiseStmts(RewriteRule):
         (operator_list := ilist.New(values=operators)).insert_before(node)
         return operator_list.result
 
-    def rewrite_p_p_error(self, node: PPError) -> RewriteResult:
-        (operators := ilist.New(values=(node.op,))).insert_before(node)
-        (ps := ilist.New(values=(node.p,))).insert_before(node)
-        stochastic_channel = StochasticUnitaryChannel(
-            operators=operators.result, probabilities=ps.result
-        )
-
-        node.replace_by(stochastic_channel)
-        return RewriteResult(has_done_something=True)
-
     def rewrite_depolarize(self, node: Depolarize) -> RewriteResult:
         paulis = (X(), Y(), Z())
         operators: list[ir.SSAValue] = []

bloqade/squin/noise/stmts.py

@@ -2,10 +2,8 @@ from kirin import ir, types, lowering
 from kirin.decl import info, statement
 from kirin.dialects import ilist
 
-from bloqade.squin.op.types import OpType
-
 from ._dialect import dialect
-from ..op.types import NumOperators
+from ..op.types import OpType, NumOperators, MultiQubitPauliOpType
 
 
 @statement
@@ -16,17 +14,7 @@ class NoiseChannel(ir.Statement):
 
 @statement(dialect=dialect)
 class PauliError(NoiseChannel):
-    basis: ir.SSAValue = info.argument(OpType)
-    p: ir.SSAValue = info.argument(types.Float)
-
-
-@statement(dialect=dialect)
-class PPError(NoiseChannel):
-    """
-    Pauli Product Error
-    """
-
-    op: ir.SSAValue = info.argument(OpType)
+    basis: ir.SSAValue = info.argument(MultiQubitPauliOpType)
     p: ir.SSAValue = info.argument(types.Float)
 
 

bloqade/squin/op/_wrapper.py

@@ -62,15 +62,15 @@ def phase(theta: float) -> types.Op: ...
 
 
 @wraps(stmts.X)
-def x() -> types.Op: ...
+def x() -> types.PauliOp: ...
 
 
 @wraps(stmts.Y)
-def y() -> types.Op: ...
+def y() -> types.PauliOp: ...
 
 
 @wraps(stmts.Z)
-def z() -> types.Op: ...
+def z() -> types.PauliOp: ...
 
 
 @wraps(stmts.SqrtX)
@@ -118,4 +118,4 @@ def u(theta: float, phi: float, lam: float) -> types.Op: ...
 
 
 @wraps(stmts.PauliString)
-def pauli_string(*, string: str) -> types.Op: ...
+def pauli_string(*, string: str) -> types.PauliStringOp: ...

bloqade/squin/op/stmts.py

@@ -1,7 +1,19 @@
 from kirin import ir, types, lowering
 from kirin.decl import info, statement
 
-from .types import OpType
+from .types import (
+    OpType,
+    ROpType,
+    XOpType,
+    YOpType,
+    ZOpType,
+    KronType,
+    MultType,
+    PauliOpType,
+    ControlOpType,
+    PauliStringType,
+    ControlledOpType,
+)
 from .number import NumberType
 from .traits import Unitary, HasSites, FixedSites, MaybeUnitary
 from ._dialect import dialect
@@ -22,22 +34,28 @@ class CompositeOp(Operator):
     pass
 
 
+LhsType = types.TypeVar("Lhs", bound=OpType)
+RhsType = types.TypeVar("Rhs", bound=OpType)
+
+
 @statement
 class BinaryOp(CompositeOp):
-    lhs: ir.SSAValue = info.argument(OpType)
-    rhs: ir.SSAValue = info.argument(OpType)
+    lhs: ir.SSAValue = info.argument(LhsType)
+    rhs: ir.SSAValue = info.argument(RhsType)
 
 
 @statement(dialect=dialect)
 class Kron(BinaryOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
     is_unitary: bool = info.attribute(default=False)
+    result: ir.ResultValue = info.result(KronType[LhsType, RhsType])
 
 
 @statement(dialect=dialect)
 class Mult(BinaryOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
     is_unitary: bool = info.attribute(default=False)
+    result: ir.ResultValue = info.result(MultType[LhsType, RhsType])
 
 
 @statement(dialect=dialect)
@@ -59,15 +77,20 @@ class Scale(CompositeOp):
 class Control(CompositeOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), MaybeUnitary()})
     is_unitary: bool = info.attribute(default=False)
-    op: ir.SSAValue = info.argument(OpType)
+    op: ir.SSAValue = info.argument(ControlledOpType)
     n_controls: int = info.attribute()
+    result: ir.ResultValue = info.result(ControlOpType[ControlledOpType])
+
+
+RotationAxisType = types.TypeVar("RotationAxis", bound=OpType)
 
 
 @statement(dialect=dialect)
 class Rot(CompositeOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary()})
-    axis: ir.SSAValue = info.argument(OpType)
+    axis: ir.SSAValue = info.argument(RotationAxisType)
     angle: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(ROpType[RotationAxisType])
 
 
 @statement(dialect=dialect)
@@ -166,13 +189,14 @@ class CliffordOp(ConstantUnitary):
 
 @statement
 class PauliOp(CliffordOp):
-    pass
+    result: ir.ResultValue = info.result(type=PauliOpType)
 
 
 @statement(dialect=dialect)
 class PauliString(ConstantUnitary):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), HasSites()})
     string: str = info.attribute()
+    result: ir.ResultValue = info.result(type=PauliStringType)
 
     def verify(self) -> None:
         if not set("XYZ").issuperset(self.string):
@@ -183,17 +207,17 @@ class PauliString(ConstantUnitary):
 
 @statement(dialect=dialect)
 class X(PauliOp):
-    pass
+    result: ir.ResultValue = info.result(XOpType)
 
 
 @statement(dialect=dialect)
 class Y(PauliOp):
-    pass
+    result: ir.ResultValue = info.result(YOpType)
 
 
 @statement(dialect=dialect)
 class Z(PauliOp):
-    pass
+    result: ir.ResultValue = info.result(ZOpType)
 
 
 @statement(dialect=dialect)

bloqade/squin/op/types.py

@@ -1,4 +1,4 @@
-from typing import overload
+from typing import Generic, TypeVar, overload
 
 from kirin import types
 
@@ -23,4 +23,106 @@ class Op:
 
 OpType = types.PyClass(Op)
 
-NumOperators = types.TypeVar("NumOperators")
+
+class CompositeOp(Op):
+    pass
+
+
+CompositeOpType = types.PyClass(CompositeOp)
+
+LhsType = TypeVar("LhsType", bound=Op)
+RhsType = TypeVar("RhsType", bound=Op)
+
+
+class BinaryOp(Op, Generic[LhsType, RhsType]):
+    lhs: LhsType
+    rhs: RhsType
+
+
+BinaryOpType = types.Generic(BinaryOp, OpType, OpType)
+
+
+class Mult(BinaryOp[LhsType, RhsType]):
+    pass
+
+
+MultType = types.Generic(Mult, OpType, OpType)
+
+
+class Kron(BinaryOp[LhsType, RhsType]):
+    pass
+
+
+KronType = types.Generic(Kron, OpType, OpType)
+
+
+class MultiQubitPauliOp(Op):
+    pass
+
+
+MultiQubitPauliOpType = types.PyClass(MultiQubitPauliOp)
+
+
+class PauliStringOp(MultiQubitPauliOp):
+    pass
+
+
+PauliStringType = types.PyClass(PauliStringOp)
+
+
+class PauliOp(MultiQubitPauliOp):
+    pass
+
+
+PauliOpType = types.PyClass(PauliOp)
+
+
+class XOp(PauliOp):
+    pass
+
+
+XOpType = types.PyClass(XOp)
+
+
+class YOp(PauliOp):
+    pass
+
+
+YOpType = types.PyClass(YOp)
+
+
+class ZOp(PauliOp):
+    pass
+
+
+ZOpType = types.PyClass(ZOp)
+
+
+ControlledOp = TypeVar("ControlledOp", bound=Op)
+
+
+class ControlOp(CompositeOp, Generic[ControlledOp]):
+    op: ControlledOp
+
+
+ControlledOpType = types.TypeVar("ControlledOp", bound=OpType)
+ControlOpType = types.Generic(ControlOp, ControlledOpType)
+CXOpType = ControlOpType[XOpType]
+CYOpType = ControlOpType[YOpType]
+CZOpType = ControlOpType[ZOpType]
+
+RotationAxis = TypeVar("RotationAxis", bound=Op)
+
+
+class ROp(CompositeOp, Generic[RotationAxis]):
+    axis: RotationAxis
+    angle: float
+
+
+ROpType = types.Generic(ROp, OpType)
+RxOpType = ROpType[XOpType]
+RyOpType = ROpType[YOpType]
+RzOpType = ROpType[ZOpType]
+
+
+NumOperators = types.TypeVar("NumOperators", bound=types.Int)

bloqade/squin/qubit.py

@@ -7,7 +7,7 @@ Depends on:
 - `kirin.dialects.ilist`: provides the `ilist.IListType` type for lists of qubits.
 """
 
-from typing import Any, overload
+from typing import Any, TypeVar, overload
 
 from kirin import ir, types, lowering
 from kirin.decl import info, statement
@@ -18,7 +18,7 @@ from bloqade.types import Qubit, QubitType
 from bloqade.squin.op.types import Op, OpType
 
 from .types import MeasurementResult, MeasurementResultType
-from .lowering import ApplyAnyCallLowering
+from .lowering import ApplyAnyCallLowering, BroadcastCallLowering
 
 dialect = ir.Dialect("squin.qubit")
 
@@ -34,7 +34,7 @@ class New(ir.Statement):
 class Apply(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     operator: ir.SSAValue = info.argument(OpType)
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+    qubits: tuple[ir.SSAValue, ...] = info.argument(QubitType)
 
 
 @statement(dialect=dialect)
@@ -47,9 +47,9 @@ class ApplyAny(ir.Statement):
 
 @statement(dialect=dialect)
 class Broadcast(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
+    traits = frozenset({BroadcastCallLowering()})
     operator: ir.SSAValue = info.argument(OpType)
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+    qubits: tuple[ir.SSAValue, ...] = info.argument(ilist.IListType[QubitType])
 
 
 @statement(dialect=dialect)
@@ -93,26 +93,10 @@ def new(n_qubits: int) -> ilist.IList[Qubit, Any]:
     ...
 
 
-@overload
-def apply(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> None:
-    """Apply an operator to a list of qubits.
-
-    Note, that when considering atom loss, lost qubits will be skipped.
-
-    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
-    """
-    ...
-
-
-@overload
+@wraps(ApplyAny)
 def apply(operator: Op, *qubits: Qubit) -> None:
-    """Apply and operator to any number of qubits.
+    """Apply an operator to qubits. The number of qubit arguments must match the
+    size of the operator.
 
     Note, that when considering atom loss, lost qubits will be skipped.
 
@@ -127,10 +111,6 @@ def apply(operator: Op, *qubits: Qubit) -> None:
     ...
 
 
-@wraps(ApplyAny)
-def apply(operator: Op, *qubits) -> None: ...
-
-
 @overload
 def measure(input: Qubit) -> MeasurementResult: ...
 @overload
@@ -154,23 +134,30 @@ def measure(input: Any) -> Any:
     ...
 
 
+OpSize = TypeVar("OpSize")
+
+
 @wraps(Broadcast)
-def broadcast(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> None:
-    """Broadcast and apply an operator to a list of qubits. For example, an operator
-    that expects 2 qubits can be applied to a list of 2n qubits, where n is an integer > 0.
+def broadcast(operator: Op, *qubits: ilist.IList[Qubit, OpSize] | list[Qubit]) -> None:
+    """Broadcast and apply an operator to lists of qubits. The number of qubit lists must
+    match the size of the operator and the lists must be of same length. The operator is
+    then applied to the list elements similar to what python's map function does.
 
-    For controlled operators, the list of qubits is interpreted as sets of (controls, targets).
-    For example
+    ## Usage examples
 
-    ```
-    apply(CX, [q0, q1, q2, q3])
-    ```
+    ```python
+    from bloqade import squin
 
-    is equivalent to
+    @squin.kernel
+    def ghz():
+        controls = squin.qubit.new(4)
+        targets = squin.qubit.new(4)
 
-    ```
-    apply(CX, [q0, q1])
-    apply(CX, [q2, q3])
+        h = squin.op.h()
+        squin.qubit.broadcast(h, controls)
+
+        cx = squin.op.cx()
+        squin.qubit.broadcast(cx, controls, targets)
     ```
 
     Args: