"""GPU execution engine backed by CUDA-Q (cuStateVec / cuTensorNet).
This engine mirrors :class:`QarpEngine` so algorithms
(VQE, QPE, …) swap to GPU with a one-line change::
engine = CudaqEngine(backend="statevector") # instead of QarpEngine()
The difference from QarpEngine is purely the target: circuits are transpiled to
``qx.cudaq_gateset()`` and run through ``qx.CudaqSimulator`` on the GPU. There
is no device-routing machinery here — a state-vector / tensor-network simulator
is all-to-all, so the router (used by QarpEngine for
hardware connectivity) is irrelevant. Add it back only if you want to
benchmark *routed* circuits on the GPU.
Importing this module is always safe (including on CPU-only wheels); the CUDA
requirement is checked lazily at construction by :func:`_require_cudaq`.
"""
from __future__ import annotations
from typing import Optional, Union
import qarpx as qx
from qarp.operators import QubitOperator
from .._types import Consumes, PrimitiveResult, Shots, Target
from ..errors import CapabilityError
from ._engine import (
Engine,
_exact_result,
_operator_n_qubits,
_qubit_operator_to_observable,
_substituted,
)
from ._runnable import Runnable
def _is_gpu_expectation(prim) -> bool:
"""True for an EXPECTATION_VALUE ``StateVector`` over a ``QubitOperator``.
These are routed to the GPU ``batch_expectation`` path (sampling-free
⟨ψ|H|ψ⟩ on-device) instead of pulling a full statevector back to host —
the dominant VQE/QAOA objective and the main GPU win. Anything else
(OVERLAP, TRANSITION_AMPLITUDE, non-QubitOperator) falls back to
``run_from_amplitudes`` with the GPU simulator injected.
"""
return (
getattr(prim, "consumes", None) is Consumes.AMPLITUDES
and getattr(prim, "target", None) is Target.EXPECTATION_VALUE
and isinstance(getattr(prim, "operator", None), QubitOperator)
)
# String → C++ enum member name. Kept as plain strings so this module imports
# cleanly even when the GPU enums are absent (CPU-only build); the actual enum
# is resolved after the availability guard in __init__.
_BACKEND_NAMES = {
"statevector": "StateVector",
"statevector-mgpu": "StateVectorMGPU",
"tensornet": "TensorNet",
"tensornet-mps": "TensorNetMPS",
}
_PRECISION_NAMES = {"fp32": "FP32", "fp64": "FP64"}
# Above this qubit count, pulling a full statevector back to host is refused
# (2^33 * 16 B ≈ 137 GB — already an H200's worth). Sampling/expectation are
# unaffected; only the host-statevector (amplitude / EXACT) path is capped.
_STATEVECTOR_HOST_QUBIT_CAP = 30
def _check_statevector_host_width(n: int, consumer: str) -> None:
"""Refuse a device→host full-statevector transfer above the qubit cap.
``consumer`` names the primitive or algorithm in the message. Above the
cap a ``2**n``-amplitude copy is a multi-hundred-GB pull.
"""
if n > _STATEVECTOR_HOST_QUBIT_CAP:
raise CapabilityError(
f"CudaqEngine: {consumer} needs a host statevector of {n} "
f"qubits (2**{n} amplitudes), above the {_STATEVECTOR_HOST_QUBIT_CAP}-qubit "
"cap. Re-express the objective as an EXPECTATION_VALUE (computed "
"on-device via batch_expectation), or raise _STATEVECTOR_HOST_QUBIT_CAP "
"if you have the host memory."
)
def _check_statevector_host(prim) -> None:
"""Primitive-level form of the cap: the amplitude path (OVERLAP /
TRANSITION_AMPLITUDE / EXACT readout) copies ``2**n`` amplitudes to host.
EXPECTATION_VALUE is exempt — ``batch_expectation`` keeps everything on-device.
"""
width = max((c.n_qubits for c in prim.compiled), default=0)
_check_statevector_host_width(width, type(prim).__name__)
def _require_cudaq():
"""Return the ``qx.CudaqSimulator`` class or raise ``CapabilityError``.
A CPU-only qarpx registers ``CudaqSimulator`` with ``available()`` False;
the ``getattr`` also tolerates a build that omits the class entirely.
"""
sim_cls = getattr(qx, "CudaqSimulator", None)
if sim_cls is None or not sim_cls.available():
raise CapabilityError(
"CudaqEngine requires a qarpx built with CUDA-Q support, and the "
"installed qarpx is CPU-only. There is no GPU wheel — published "
"wheels are QARP_WITH_CUDAQ=OFF — so this needs a source build of "
"qarpx against the CUDA-Q runtime: install `cuda-quantum-cu12` "
"(the `cudaq-runtime` extra), then rebuild with "
"QARP_WITH_CUDAQ=ON, QARP_CUDAQ_WHEEL_DIR and "
"QARP_CUDAQ_NVQIR_BACKEND=nvqir-custatevec-fp64. The full command "
"is under 'GPU execution (CUDA-Q)' in docs/source/installation.rst. "
"For CPU simulation use QarpEngine."
)
return sim_cls
[docs]
class CudaqEngine(Engine):
"""Pure-Python engine wrapping the C++ Transpiler + ``qx.CudaqSimulator``.
``provides_amplitudes = True`` (base default) — noiseless simulator; the
device→host statevector transfer is separately capped by
``_check_statevector_host``.
Gradients: ``"default"`` resolves to the batched parameter shift, which
stays on-device for StateVector expectation values over a QubitOperator
(``batch_expectation``); other targets pull the host statevector per
point. ``"adjoint"`` is refused — the adjoint is a CPU kernel and this
engine never constructs a CPU simulator behind the caller's back (a GPU
adjoint is post-release work); use QarpEngine for it.
Args:
backend: One of ``statevector`` (single GPU, default),
``statevector-mgpu`` (multi-GPU + MPI), ``tensornet``,
``tensornet-mps``.
precision: ``fp64`` (default) or ``fp32``.
max_bond_dim: Bond-dimension cap for ``tensornet-mps`` (ignored
otherwise).
target: Optional NVQIR backend name passed verbatim to the simulator
(``CUDAQ_DEFAULT_SIMULATOR``), overriding ``backend``/``precision``.
Use ``"qpp"`` to run on CPU (e.g. for tests/CI without a GPU) or to
reach a backend not covered by the ``backend`` enum.
n_shots: Default shot count used when a primitive has ``n_shots=None``;
``qarp.EXACT`` makes exact readout the engine-wide default.
seed: Optional RNG seed (passed to the simulator). EXACT readouts
involve no RNG and are seed-independent.
"""
gradient_methods = frozenset({"default", "parameter-shift", "finite-diff", "spsa"})
def __init__(
self,
*,
backend: str = "statevector",
precision: str = "fp64",
max_bond_dim: int = 0,
target: Optional[str] = None,
n_shots: Union[int, Shots] = 10_000,
seed: Optional[int] = None,
):
if backend not in _BACKEND_NAMES:
raise ValueError(
f"CudaqEngine: unknown backend {backend!r}; choose one of {sorted(_BACKEND_NAMES)}."
)
if precision not in _PRECISION_NAMES:
raise ValueError(
f"CudaqEngine: unknown precision {precision!r}; "
f"choose one of {sorted(_PRECISION_NAMES)}."
)
sim_cls = _require_cudaq()
cfg = qx.CudaqConfig()
cfg.backend = getattr(qx.CudaqBackend, _BACKEND_NAMES[backend])
cfg.precision = getattr(qx.CudaqPrecision, _PRECISION_NAMES[precision])
cfg.max_bond_dim = max_bond_dim
if target is not None:
cfg.target = target
self._sim = sim_cls(cfg)
# Rebase to the CUDA-Q native gate set (no fusion — see _compile_one).
self._transpiler = qx.Transpiler(qx.cudaq_gateset())
self._n_shots = n_shots
self._seed = seed
self._primitives: list[Runnable] = []
# ── Template hooks (build()/run() live on the base Engine) ─────────────
def _compile_one(self, flat, block_n_qubits: int):
"""Transpile (rebase) one flat command list to cudaq_gateset.
Uses ``transpile`` rather than ``transpile_and_optimize``: the
single-qubit fusion pass emits dense ``Custom`` gates, which the CUDA-Q
lowerer (named-gate kernel builder) cannot represent. Plain transpile
keeps the stream to named gates in ``cudaq_gateset``. (A future
optimisation could ZYZ-decompose fused ``Custom`` gates back to
Rz/Ry/Rz and re-enable fusion.)
No routing: an all-to-all simulator needs none, so the layout is
always None and there is no reindex step in run()/batch_run(). A gate the
CUDA-Q gate set cannot reach raises ``CapabilityError`` (§14).
"""
return self._transpiler.transpile(flat), block_n_qubits, None
def _host_statevector_simulator(self, n_qubits: int):
# Algorithm-level fast paths (ADAPT pool scans, QSE, MonteCarlo's
# exact reference) pull 2**n amplitudes to host: same cap as the
# primitive path, so the gate cannot bypass it.
_check_statevector_host_width(n_qubits, "a statevector fast path")
return self._sim
def _post_compile_check(self, prim: Runnable, compiled, layout) -> None:
# lower() (cudaq_simulator.cpp) handles unitary gates plus
# end-of-circuit measurement only — reject the rest with a clear
# message rather than emitting a circuit that silently drops them.
self._reject_true_mcm(
compiled,
primitive_name=type(prim).__name__,
reason=("The CUDA-Q lowerer handles unitary gates plus end-of-circuit sampling only."),
)
def _dispatch_one(self, prim: Runnable, compiled, ordinal: int, params):
substituted = _substituted(compiled, params)
widths = [c.n_qubits for c in compiled]
if prim.consumes is Consumes.AMPLITUDES:
if _is_gpu_expectation(prim):
# Pad to the operator width (idle qubits stay |0>).
n_q = max(widths[0], _operator_n_qubits(prim.operator))
obs = _qubit_operator_to_observable(prim.operator, n_q)
# One already-substituted circuit; empty param set = no-op.
val = self._sim.batch_expectation(substituted[0], n_q, obs, [{}])[0]
return complex(val), ordinal
_check_statevector_host(prim)
return prim.run_from_amplitudes(substituted, simulator=self._sim), ordinal
n_shots = self._resolve_shots(prim)
if n_shots is Shots.EXACT:
# Exact Born distribution needs the statevector on host.
_check_statevector_host(prim)
sampling_results = [
_exact_result(self._sim, cmds, n_q)
for cmds, n_q in zip(substituted, widths, strict=True)
]
else:
sampling_results = [
self._sim.run(cmds, n_q, n_shots, self._circuit_seed(ordinal + k))
for k, (cmds, n_q) in enumerate(zip(substituted, widths, strict=True))
]
ordinal += len(sampling_results)
return prim.run(sampling_results), ordinal
def _sweep(
self,
primitives,
compiled_per_prim,
param_sets,
shots_override,
) -> list[list[PrimitiveResult]]:
"""GPU inner loop: the parameter sweep stays on-device in
``CudaqSimulator.batch_run`` / ``batch_expectation`` — no Python
round-trip per parameter set. This is the VQE inner loop and the main
reason to reach for the GPU engine. Shot resolution per primitive:
``shots_override`` > ``prim.n_shots`` > engine default; EXACT
primitives evaluate per parameter set (host statevector, capped)."""
# Pre-compute the batched payload per primitive, keeping every sweep on
# the GPU. Three kinds:
# "expect" — EXPECTATION_VALUE: one batch_expectation call sweeps all
# param sets on-device → list[n_sets] of scalars.
# "amps" — other amplitude targets (OVERLAP, …): per-set GPU
# statevector via run_from_amplitudes (no up-front sampling).
# "sample" — sampling primitives: batch_run over all param sets.
# Second element is always a list ([] for the "amps" case) so indexing
# stays well-typed: "expect" → scalars per set, "sample" → per-circuit
# batch_run results.
prim_data: list[tuple[str, list]] = []
ordinal = 0
for prim_idx, prim in enumerate(primitives):
compiled = compiled_per_prim[prim_idx]
if _is_gpu_expectation(prim):
# Pad to the operator width (idle qubits stay |0>).
n_q = max(compiled[0].n_qubits, _operator_n_qubits(prim.operator))
obs = _qubit_operator_to_observable(prim.operator, n_q)
vals = self._sim.batch_expectation(compiled[0].commands, n_q, obs, list(param_sets))
prim_data.append(("expect", list(vals)))
elif prim.consumes is Consumes.AMPLITUDES:
_check_statevector_host(prim)
prim_data.append(("amps", []))
elif self._resolve_shots(prim, shots_override) is Shots.EXACT:
_check_statevector_host(prim)
prim_data.append(("exact", []))
else:
shots = self._resolve_shots(prim, shots_override)
circ_batch = [
self._sim.batch_run(
c.commands, c.n_qubits, shots, param_sets, self._circuit_seed(ordinal + k)
)
for k, c in enumerate(compiled)
]
ordinal += len(circ_batch)
prim_data.append(("sample", circ_batch))
results_by_set: list[list[PrimitiveResult]] = []
for set_idx, ps in enumerate(param_sets):
set_results: list[PrimitiveResult] = []
for prim_idx, prim in enumerate(primitives):
kind, data = prim_data[prim_idx]
compiled = compiled_per_prim[prim_idx]
if kind == "expect":
set_results.append(complex(data[set_idx]))
elif kind == "amps":
substituted = _substituted(compiled, ps)
set_results.append(prim.run_from_amplitudes(substituted, simulator=self._sim))
elif kind == "exact":
sampling_results = [
_exact_result(self._sim, cmds, c.n_qubits)
for cmds, c in zip(_substituted(compiled, ps), compiled, strict=True)
]
set_results.append(prim.run(sampling_results))
else: # "sample"
sampling_results = [
data[circ_idx][set_idx] for circ_idx in range(len(compiled))
]
set_results.append(prim.run(sampling_results))
results_by_set.append(set_results)
return results_by_set