bloqade-circuit 0.6.7__py3-none-any.whl → 0.7.0__py3-none-any.whl

bloqade/pyqrack/squin/qubit.py

@@ -25,7 +25,7 @@ class PyQrackMethods(interp.MethodTable):
 
     @interp.impl(qubit.Apply)
     def apply(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Apply):
-        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        qubits: list[PyQrackQubit] = [frame.get(qbit) for qbit in stmt.qubits]
         operator: OperatorRuntimeABC = frame.get(stmt.operator)
         operator.apply(*qubits)
 
@@ -34,7 +34,9 @@ class PyQrackMethods(interp.MethodTable):
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Broadcast
     ):
         operator: OperatorRuntimeABC = frame.get(stmt.operator)
-        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        qubits: list[ilist.IList[PyQrackQubit, Any]] = [
+            frame.get(qbit) for qbit in stmt.qubits
+        ]
         operator.broadcast_apply(qubits)
 
     def _measure_qubit(self, qbit: PyQrackQubit, interp: PyQrackInterpreter):

bloqade/pyqrack/squin/runtime.py

@@ -28,17 +28,25 @@ class OperatorRuntimeABC:
     ) -> None:
         raise RuntimeError(f"Can't apply controlled version of {self}")
 
-    def broadcast_apply(self, qubits: ilist.IList[PyQrackQubit, Any], **kwargs) -> None:
+    def broadcast_apply(
+        self, qubit_lists: list[ilist.IList[PyQrackQubit, Any]], **kwargs
+    ) -> None:
         n = self.n_sites
 
-        if len(qubits) % n != 0:
+        if n != len(qubit_lists):
             raise RuntimeError(
-                f"Cannot broadcast operator {self} that applies to {n} over {len(qubits)} qubits."
+                f"Cannot apply operator of size {n} to {len(qubit_lists)} qubits!"
             )
 
-        for qubit_index in range(0, len(qubits), n):
-            targets = qubits[qubit_index : qubit_index + n]
-            self.apply(*targets, **kwargs)
+        m = len(qubit_lists[0])
+        for qubit_list in qubit_lists:
+            if m != len(qubit_list):
+                raise RuntimeError(
+                    "Cannot broadcast operator on qubit lists of varying length!"
+                )
+
+        for qubits in zip(*qubit_lists):
+            self.apply(*qubits, **kwargs)
 
 
 @dataclass(frozen=True)

bloqade/pyqrack/task.py

@@ -2,6 +2,7 @@ from typing import TypeVar, ParamSpec, cast
 from dataclasses import dataclass
 
 from bloqade.task import AbstractSimulatorTask
+from bloqade.pyqrack.reg import QubitState, PyQrackQubit
 from bloqade.pyqrack.base import (
     MemoryABC,
     PyQrackInterpreter,
@@ -36,3 +37,17 @@ class PyQrackSimulatorTask(AbstractSimulatorTask[Param, RetType, MemoryType]):
         """Returns the state vector of the simulator."""
         self.run()
         return self.state.sim_reg.out_ket()
+
+    def qubits(self) -> list[PyQrackQubit]:
+        """Returns the qubits in the simulator."""
+        try:
+            N = self.state.sim_reg.num_qubits()
+            return [
+                PyQrackQubit(
+                    addr=i, sim_reg=self.state.sim_reg, state=QubitState.Active
+                )
+                for i in range(N)
+            ]
+        except AttributeError:
+            Warning("Task has not been run, there are no qubits!")
+            return []

bloqade/rewrite/rules/split_ifs.py

@@ -46,9 +46,13 @@ class SplitIfStmts(RewriteRule):
         if not isinstance(node, scf.IfElse):
             return RewriteResult()
 
+        # NOTE: only empty else bodies are allowed in valid QASM2
+        if not self._has_empty_else(node):
+            return RewriteResult()
+
         *stmts, yield_or_return = node.then_body.stmts()
 
-        if len(stmts) == 1:
+        if len(stmts) <= 1:
             return RewriteResult()
 
         is_yield = isinstance(yield_or_return, scf.Yield)
@@ -71,3 +75,16 @@ class SplitIfStmts(RewriteRule):
         node.delete()
 
         return RewriteResult(has_done_something=True)
+
+    def _has_empty_else(self, node: scf.IfElse) -> bool:
+        else_stmts = list(node.else_body.stmts())
+        if len(else_stmts) > 1:
+            return False
+
+        if len(else_stmts) == 0:
+            return True
+
+        if not isinstance(else_stmts[0], scf.Yield):
+            return False
+
+        return len(else_stmts[0].values) == 0

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
@@ -13,6 +14,7 @@ from .runtime import (
     MultRuntime,
     ScaleRuntime,
     AdjointRuntime,
+    BasicOpRuntime,
     ControlRuntime,
     UnitaryRuntime,
     HermitianRuntime,
@@ -116,10 +118,59 @@ 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(
         self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.PauliString
     ):
         return (PauliStringRuntime(stmt.string),)
+
+    @impl(op.stmts.Rot)
+    def rot(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Rot):
+        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)
+
+        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):
+        cirq_op = cirq.XPowGate(exponent=0.5)
+        return (UnitaryRuntime(cirq_op),)
+
+    @impl(op.stmts.SqrtY)
+    def sqrt_y(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.SqrtY):
+        cirq_op = cirq.YPowGate(exponent=0.5)
+        return (UnitaryRuntime(cirq_op),)

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/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)