diff --git a/cuda_core/cuda/core/_memory/_buffer.pxd b/cuda_core/cuda/core/_memory/_buffer.pxd index b552e69554d..a9fa0d7e99c 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pxd +++ b/cuda_core/cuda/core/_memory/_buffer.pxd @@ -44,3 +44,9 @@ cdef Buffer Buffer_from_deviceptr_handle( object ipc_descriptor = *, type cls = *, ) + + +# Shared argument coercion for the batched free functions (copy_batch, +# prefetch_batch, discard_batch, discard_prefetch_batch). `single_hint` +# names the per-buffer API to use instead when a bare Buffer is passed. +cdef tuple Buffer_coerce_batch(object buffers, str what, str single_hint) diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 2506331d0fd..76837776383 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -29,6 +29,7 @@ from cuda.core._stream cimport Stream, Stream_accept, default_stream from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value import sys +from collections.abc import Sequence from typing import TYPE_CHECKING from cuda.core._utils.pycompat import BufferProtocol @@ -619,6 +620,32 @@ cdef Buffer Buffer_from_deviceptr_handle( return buf +cdef tuple Buffer_coerce_batch(object buffers, str what, str single_hint): + """Coerce ``buffers`` to a ``tuple[Buffer, ...]``; reject a bare Buffer. + + Shared by the batched free functions. Passing one Buffer is rejected + rather than treated as a one-element batch so that the per-buffer API + named by ``single_hint`` stays the single obvious way to do it. + """ + cdef list out + if isinstance(buffers, Buffer): + raise TypeError( + f"{what}: pass a sequence of Buffers; for a single buffer use {single_hint}" + ) + if not isinstance(buffers, Sequence): + raise TypeError( + f"{what}: buffers must be a sequence of Buffer, got {type(buffers).__name__}" + ) + if not buffers: + raise ValueError(f"{what}: empty buffers sequence") + out = [] + for item in buffers: + if not isinstance(item, Buffer): + raise TypeError(f"{what}: expected Buffer, got {type(item).__name__}") + out.append(item) + return tuple(out) + + cdef inline void Buffer_close(Buffer self, object stream): """Close a buffer, freeing its memory.""" cdef Stream s diff --git a/cuda_core/cuda/core/_memory/_copy_enums.py b/cuda_core/cuda/core/_memory/_copy_enums.py new file mode 100644 index 00000000000..20bfc469bf7 --- /dev/null +++ b/cuda_core/cuda/core/_memory/_copy_enums.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import dataclasses +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from cuda.core._utils.cuda_utils import driver +from cuda.core._utils.pycompat import StrEnum +from cuda.core._utils.version import binding_version + +if TYPE_CHECKING: + from cuda.core._device import Device + from cuda.core._host import Host + + +__all__ = ["CopyOptions", "MemcpyOverlapMode", "MemcpySrcAccessOrder"] + + +class MemcpySrcAccessOrder(StrEnum): + """Source access order hint for batched memcpy operations. + + Maps to ``CUmemcpySrcAccessOrder``. The ``INVALID`` and ``MAX`` + sentinel values from the driver enum are excluded from the public + Python surface. + """ + + STREAM = "stream" + DURING_API_CALL = "during_api_call" + ANY = "any" + + +class MemcpyOverlapMode(StrEnum): + """Overlap mode hint for batched memcpy operations. + + Maps to ``CUmemcpyFlags``. Renamed from "flags" to "overlap_mode" + for clarity; the only non-default flag is CE/compute overlap + (Tegra). + """ + + DEFAULT = "default" + PREFER_OVERLAP_WITH_COMPUTE = "prefer_overlap_with_compute" + + +@dataclasses.dataclass(frozen=True) +class CopyOptions: + """Attribute bundle for a single copy within a batched memcpy. + + Parameters + ---------- + src_access_order : :class:`MemcpySrcAccessOrder` or str + Hint describing how the source will be accessed. + Default is ``"stream"`` (stream-ordered access). + src_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None + Hint for the source memory location. ``None`` means no hint. + dst_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None + Hint for the destination memory location. ``None`` means no hint. + overlap_mode : :class:`MemcpyOverlapMode` or str + Hint for copy-engine / compute overlap. Only meaningful on + integrated (Tegra) GPUs; on discrete GPUs the driver silently + ignores it and a :class:`UserWarning` is emitted. + Default is ``"default"``. + """ + + src_access_order: MemcpySrcAccessOrder | str = "stream" + src_location_hint: Device | Host | None = None + dst_location_hint: Device | Host | None = None + overlap_mode: MemcpyOverlapMode | str = "default" + + def __post_init__(self): + # Frozen, unlike the other *Options dataclasses in cuda.core, because + # the batched-API contract agreed in NVIDIA/cuda-python#1775 specifies + # immutable per-call options: + # https://github.com/NVIDIA/cuda-python/pull/1775#issuecomment-4355502334 + # + # Normalizing str -> StrEnum therefore has to go through + # object.__setattr__; a plain assignment would raise + # FrozenInstanceError. Done here rather than at use so that a typo + # fails at construction and the field always holds the enum. + if isinstance(self.src_access_order, str): + try: + object.__setattr__( + self, + "src_access_order", + MemcpySrcAccessOrder(self.src_access_order), + ) + except ValueError as exc: + raise ValueError(f"invalid src_access_order: {self.src_access_order!r}") from exc + if isinstance(self.overlap_mode, str): + try: + object.__setattr__( + self, + "overlap_mode", + MemcpyOverlapMode(self.overlap_mode), + ) + except ValueError as exc: + raise ValueError(f"invalid overlap_mode: {self.overlap_mode!r}") from exc + + def _to_driver_enum(self) -> int: + """Return the driver CUmemcpySrcAccessOrder value.""" + if not _SRC_ACCESS_ORDER_TO_DRIVER: + raise NotImplementedError(_CUDA13_REQUIRED) + return _SRC_ACCESS_ORDER_TO_DRIVER[MemcpySrcAccessOrder(self.src_access_order)] + + def _to_driver_flags(self) -> int: + """Return the driver CUmemcpyFlags value.""" + if not _OVERLAP_MODE_TO_DRIVER: + raise NotImplementedError(_CUDA13_REQUIRED) + return _OVERLAP_MODE_TO_DRIVER[MemcpyOverlapMode(self.overlap_mode)] + + +_CUDA13_REQUIRED = "copy attributes require cuda.bindings 13.0 or newer" + +# CUmemcpySrcAccessOrder and CUmemcpyFlags are exposed by cuda.bindings 13.0+, +# so these maps are empty when it is older. Nothing reaches them there: +# copy_batch refuses non-default CopyOptions when the batched entry point is +# unavailable. +# +# Keyed by ``str``: under ``python_version = "3.10"`` mypy resolves StrEnum to +# the unstubbed backports shim and so infers the members as plain ``str``. +# StrEnum members are ``str`` instances, so this holds on every version. The +# values are wrapped in ``int()`` because the driver enums are untyped. +_SRC_ACCESS_ORDER_TO_DRIVER: dict[str, int] +_OVERLAP_MODE_TO_DRIVER: dict[str, int] + +if binding_version() >= (13, 0, 0): + _src_order = driver.CUmemcpySrcAccessOrder + _flags = driver.CUmemcpyFlags + _SRC_ACCESS_ORDER_TO_DRIVER = { + MemcpySrcAccessOrder.STREAM: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_STREAM), + MemcpySrcAccessOrder.DURING_API_CALL: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL), + MemcpySrcAccessOrder.ANY: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_ANY), + } + _OVERLAP_MODE_TO_DRIVER = { + MemcpyOverlapMode.DEFAULT: int(_flags.CU_MEMCPY_FLAG_DEFAULT), + MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE: int(_flags.CU_MEMCPY_FLAG_PREFER_OVERLAP_WITH_COMPUTE), + } + del _src_order, _flags +else: + _SRC_ACCESS_ORDER_TO_DRIVER = {} + _OVERLAP_MODE_TO_DRIVER = {} + + +def _attr_run_starts(attrs: Sequence[CopyOptions]) -> list[int]: + """Return the start index of each maximal run of equal attributes. + + This mirrors the ``attrsIdxs`` indirection that ``cuMemcpyBatchAsync`` + expects: ``attrs[k]`` applies to the copies in + ``[starts[k], starts[k + 1])``. Collapsing equal neighbours means a + broadcast attribute is passed to the driver once (``numAttrs == 1``) + rather than repeated per copy. + """ + starts: list[int] = [] + prev: CopyOptions | None = None + for i, attr in enumerate(attrs): + if i == 0 or attr != prev: + starts.append(i) + prev = attr + return starts diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyi b/cuda_core/cuda/core/_memory/_copy_ops.pyi new file mode 100644 index 00000000000..03def659ffc --- /dev/null +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyi @@ -0,0 +1,86 @@ +# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_memory/_copy_ops.pyx + +from __future__ import annotations + +from collections.abc import Sequence + +from cuda.core._memory._buffer import Buffer +from cuda.core._memory._copy_enums import CopyOptions +from cuda.core._stream import Stream + +_SINGLE_COPY_HINT = 'Buffer.copy_to / Buffer.copy_from' +_DEFAULT_COPY_OPTIONS = CopyOptions() + +def _batch_entry_point_in_use() -> bool: + """Internal: expose the dispatch predicate so tests can gate on it.""" + +def _normalize_copy_options(options: CopyOptions | Sequence[CopyOptions] | None, n: int) -> tuple[CopyOptions, ...]: + """Expand ``options`` to exactly one :class:`CopyOptions` per copy. + + ``None`` and a scalar broadcast; a sequence pairs by index and must + already have length ``n``. + + Internal, but deliberately importable: options are hints that change + how the driver stages a transfer and never the bytes it produces, so + this expansion (and the run encoding applied to it) is the only + observable evidence that a scalar reached every copy. + """ + +def copy_batch(stream: Stream, srcs: Sequence[Buffer], dsts: Sequence[Buffer], *, options: CopyOptions | Sequence[CopyOptions] | None=None) -> None: + """Copy a batch of buffers asynchronously. + + Sizes are taken from the source buffers and each destination must + match. For a single buffer, use :meth:`Buffer.copy_to` or + :meth:`Buffer.copy_from`. + + The driver provides no graph-node form of ``cuMemcpyBatchAsync``, so + this cannot be captured into a graph. Build graph copies with + :meth:`graph.GraphNode.memcpy` or per-buffer :meth:`Buffer.copy_to`. + + Parameters + ---------- + stream : :class:`~_stream.Stream` + Stream for the asynchronous copy. First positional and required + (mirrors :func:`launch`). Unlike most stream-taking APIs this does + not accept a :class:`~graph.GraphBuilder`; one is rejected with + ``TypeError`` because the copy cannot be captured. + srcs : Sequence[:class:`Buffer`] + Source buffers. Must be a sequence, not a single Buffer. + dsts : Sequence[:class:`Buffer`] + Destination buffers. Must match ``len(srcs)``. + options : :class:`CopyOptions` | Sequence[:class:`CopyOptions`] | None + Per-copy options. A single value applies to every copy; a + sequence pairs by index and must match ``len(srcs)``. ``None`` + uses stream-ordered defaults. + + Raises + ------ + ValueError + If lengths or sizes mismatch. + TypeError + If a single Buffer is passed instead of a sequence. + NotImplementedError + If non-default ``options`` are given where + ``cuMemcpyBatchAsync`` is unavailable (see Notes). + + Notes + ----- + Batching through ``cuMemcpyBatchAsync`` requires all three of: + ``cuda.core`` built against CUDA 13 headers, ``cuda.bindings`` 13.0 or + newer, and a driver reporting CUDA 13.0 or newer + (``cuDriverGetVersion() >= 13000``). ``cuda.bindings`` binds only the + CUDA 13.0 revision of the entry point, so a driver that predates it is + refused even where it implements the earlier CUDA 12.8 signature. + + Short of that, the copies fall back to a Python-level loop over + ``cuMemcpyAsync``, which is semantically equivalent but does not + amortize launch overhead. The fallback has no way to convey + :class:`CopyOptions` to the driver, so non-default options raise + :class:`NotImplementedError` there rather than being silently ignored. + + Warns + ----- + UserWarning + If ``overlap_mode='prefer_overlap_with_compute'`` is requested + on a non-integrated (discrete) GPU. + """ \ No newline at end of file diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyx b/cuda_core/cuda/core/_memory/_copy_ops.pyx new file mode 100644 index 00000000000..0b005f2f5b5 --- /dev/null +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyx @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Sequence + +import warnings + +IF CUDA_CORE_BUILD_MAJOR >= 13: + from libcpp.vector cimport vector + +from libc.string cimport memset + +from cuda.bindings cimport cydriver +from cuda.core._memory._buffer cimport Buffer, Buffer_coerce_batch +from cuda.core._memory._location cimport to_cumemlocation +from cuda.core._resource_handles cimport as_cu +from cuda.core._stream cimport Stream, Stream_accept +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN + +# cy_driver_version and _attr_run_starts are referenced only from CUDA 13 +# branches. cython-lint does not evaluate compile-time IF blocks, so they need +# a pragma to be seen as used. +from cuda.core._utils.version cimport cy_driver_version # no-cython-lint + +from cuda.core._device import Device +from cuda.core._memory._copy_enums import CopyOptions, _attr_run_starts # no-cython-lint +from cuda.core._memory._managed_location import _coerce_location + +_SINGLE_COPY_HINT = "Buffer.copy_to / Buffer.copy_from" + +# Attributes reach the driver only through cuMemcpyBatchAsync. The per-copy +# cuMemcpyAsync fallback has nowhere to put them, so anything other than the +# defaults is refused rather than dropped. +_DEFAULT_COPY_OPTIONS = CopyOptions() + + +cdef inline bint _batch_entry_point_available(): + """Whether cuMemcpyBatchAsync can actually be called here. + + Requires ``cuda.core`` built against CUDA 13 headers (compile time) and + a driver reporting CUDA 13.0 or newer, i.e. + ``cuDriverGetVersion() >= 13000`` (run time). + + The run-time bound is set by the binding layer, not by when the driver + gained the feature. CUDA 12.8 already exposed a ``cuMemcpyBatchAsync``, + but its signature carried a ``failIdx`` out-parameter that CUDA 13.0 + dropped. ``cuda.bindings`` resolves only the 13.0 revision, via + ``cuGetProcAddress_v2('cuMemcpyBatchAsync', ..., 13000, ...)``, so an + older driver yields a NULL pointer even though it may implement the + earlier entry point. + """ + IF CUDA_CORE_BUILD_MAJOR >= 13: + return cy_driver_version() >= (13, 0, 0) + ELSE: + return False + + +def _batch_entry_point_in_use() -> bool: + """Internal: expose the dispatch predicate so tests can gate on it.""" + return bool(_batch_entry_point_available()) + + +def _normalize_copy_options( + options: CopyOptions | Sequence[CopyOptions] | None, + Py_ssize_t n, +) -> tuple[CopyOptions, ...]: + """Expand ``options`` to exactly one :class:`CopyOptions` per copy. + + ``None`` and a scalar broadcast; a sequence pairs by index and must + already have length ``n``. + + Internal, but deliberately importable: options are hints that change + how the driver stages a transfer and never the bytes it produces, so + this expansion (and the run encoding applied to it) is the only + observable evidence that a scalar reached every copy. + """ + if options is None: + return (_DEFAULT_COPY_OPTIONS,) * n + if isinstance(options, CopyOptions): + return (options,) * n + if isinstance(options, Sequence): + if len(options) != n: + raise ValueError( + f"copy_batch: options length {len(options)} does not match " + f"buffers length {n}" + ) + for a in options: + if not isinstance(a, CopyOptions): + raise TypeError( + f"copy_batch: each options element must be CopyOptions, " + f"got {type(a).__name__}" + ) + return tuple(options) + raise TypeError( + f"copy_batch: options must be CopyOptions or a sequence of " + f"CopyOptions, got {type(options).__name__}" + ) + + +cdef cydriver.CUmemcpyAttributes _to_cu_memcpy_attributes(object attr): + """Convert a CopyOptions to a cydriver.CUmemcpyAttributes struct.""" + cdef cydriver.CUmemcpyAttributes cu_attr + memset(&cu_attr, 0, sizeof(cydriver.CUmemcpyAttributes)) + cu_attr.srcAccessOrder = (attr._to_driver_enum()) + cu_attr.flags = (attr._to_driver_flags()) + + cdef object src_loc = _coerce_location(attr.src_location_hint, allow_none=True) + cdef object dst_loc = _coerce_location(attr.dst_location_hint, allow_none=True) + + if src_loc is not None: + cu_attr.srcLocHint = to_cumemlocation(src_loc) + if dst_loc is not None: + cu_attr.dstLocHint = to_cumemlocation(dst_loc) + + return cu_attr + + +def copy_batch( + stream: Stream, + srcs: Sequence[Buffer], + dsts: Sequence[Buffer], + *, + options: CopyOptions | Sequence[CopyOptions] | None = None, +) -> None: + """Copy a batch of buffers asynchronously. + + Sizes are taken from the source buffers and each destination must + match. For a single buffer, use :meth:`Buffer.copy_to` or + :meth:`Buffer.copy_from`. + + The driver provides no graph-node form of ``cuMemcpyBatchAsync``, so + this cannot be captured into a graph. Build graph copies with + :meth:`graph.GraphNode.memcpy` or per-buffer :meth:`Buffer.copy_to`. + + Parameters + ---------- + stream : :class:`~_stream.Stream` + Stream for the asynchronous copy. First positional and required + (mirrors :func:`launch`). Unlike most stream-taking APIs this does + not accept a :class:`~graph.GraphBuilder`; one is rejected with + ``TypeError`` because the copy cannot be captured. + srcs : Sequence[:class:`Buffer`] + Source buffers. Must be a sequence, not a single Buffer. + dsts : Sequence[:class:`Buffer`] + Destination buffers. Must match ``len(srcs)``. + options : :class:`CopyOptions` | Sequence[:class:`CopyOptions`] | None + Per-copy options. A single value applies to every copy; a + sequence pairs by index and must match ``len(srcs)``. ``None`` + uses stream-ordered defaults. + + Raises + ------ + ValueError + If lengths or sizes mismatch. + TypeError + If a single Buffer is passed instead of a sequence. + NotImplementedError + If non-default ``options`` are given where + ``cuMemcpyBatchAsync`` is unavailable (see Notes). + + Notes + ----- + Batching through ``cuMemcpyBatchAsync`` requires all three of: + ``cuda.core`` built against CUDA 13 headers, ``cuda.bindings`` 13.0 or + newer, and a driver reporting CUDA 13.0 or newer + (``cuDriverGetVersion() >= 13000``). ``cuda.bindings`` binds only the + CUDA 13.0 revision of the entry point, so a driver that predates it is + refused even where it implements the earlier CUDA 12.8 signature. + + Short of that, the copies fall back to a Python-level loop over + ``cuMemcpyAsync``, which is semantically equivalent but does not + amortize launch overhead. The fallback has no way to convey + :class:`CopyOptions` to the driver, so non-default options raise + :class:`NotImplementedError` there rather than being silently ignored. + + Warns + ----- + UserWarning + If ``overlap_mode='prefer_overlap_with_compute'`` is requested + on a non-integrated (discrete) GPU. + """ + cdef tuple src_bufs = Buffer_coerce_batch(srcs, "copy_batch", _SINGLE_COPY_HINT) + cdef tuple dst_bufs = Buffer_coerce_batch(dsts, "copy_batch", _SINGLE_COPY_HINT) + cdef Py_ssize_t n = len(src_bufs) + + if len(dst_bufs) != n: + raise ValueError( + f"copy_batch: srcs length {n} does not match dsts length {len(dst_bufs)}" + ) + + cdef Stream s = Stream_accept(stream) + + cdef Buffer src_buf + cdef Buffer dst_buf + cdef Py_ssize_t i + + for i in range(n): + src_buf = src_bufs[i] + dst_buf = dst_bufs[i] + if src_buf.size != dst_buf.size: + raise ValueError( + f"copy_batch: buffer size mismatch at index {i} " + f"(src={src_buf.size}, dst={dst_buf.size})" + ) + + cdef tuple attr_tuple = _normalize_copy_options(options, n) + + # Without the batched entry point there is nowhere to put attributes, + # so refuse them here rather than dropping them in the fallback. Doing + # this before the overlap warning keeps the error the only diagnostic. + if not _batch_entry_point_available(): + for i in range(n): + if attr_tuple[i] != _DEFAULT_COPY_OPTIONS: + raise NotImplementedError( + "copy_batch: non-default CopyOptions requires cuMemcpyBatchAsync, " + "which needs cuda.core built against CUDA 13, cuda.bindings 13.0+, " + "and a driver reporting CUDA 13.0 or newer; omit options to use the " + "per-copy fallback" + ) + + # Check for overlap_mode warning on non-integrated GPUs + cdef bint any_overlap = False + cdef object ca_attr + for i in range(n): + ca_attr = attr_tuple[i] + if ca_attr.overlap_mode != "default": + any_overlap = True + break + + if any_overlap: + device = Device() + if not device.properties.integrated: + warnings.warn( + "overlap_mode='prefer_overlap_with_compute' has no effect on " + "non-integrated (non-Tegra) GPUs; the transfer will use " + "default copy behavior.", + UserWarning, + stacklevel=2, + ) + + _do_copy_batch(src_bufs, dst_bufs, s, attr_tuple) + + +cdef void _do_copy_batch(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tuple): + IF CUDA_CORE_BUILD_MAJOR >= 13: + # Building against CUDA 13 headers says nothing about the installed + # driver, so the run-time version still has to be checked before + # calling a 13.0-only entry point (see PRs #2054 / #2064). + if _batch_entry_point_available(): + _do_copy_batch_native(src_bufs, dst_bufs, s, attr_tuple) + else: + _do_copy_batch_loop(src_bufs, dst_bufs, s) + ELSE: + _do_copy_batch_loop(src_bufs, dst_bufs, s) + + +cdef void _do_copy_batch_loop(tuple src_bufs, tuple dst_bufs, Stream s): + """Per-copy cuMemcpyAsync fallback where the batch entry point is absent. + + Equivalent in effect to the batched call, minus the launch-overhead + amortization. Callers guarantee the options are defaults; copy_batch + rejects anything else before reaching here. + """ + cdef Py_ssize_t n = len(src_bufs) + cdef Py_ssize_t i + cdef Buffer src_buf + cdef Buffer dst_buf + cdef size_t nbytes + cdef cydriver.CUstream hstream = as_cu(s._h_stream) + + for i in range(n): + src_buf = src_bufs[i] + dst_buf = dst_bufs[i] + nbytes = src_buf._size + with nogil: + HANDLE_RETURN(cydriver.cuMemcpyAsync( + as_cu(dst_buf._h_ptr), as_cu(src_buf._h_ptr), nbytes, hstream)) + + +IF CUDA_CORE_BUILD_MAJOR >= 13: + cdef void _do_copy_batch_native(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tuple): + cdef Py_ssize_t n = len(src_bufs) + cdef cydriver.CUstream hstream = as_cu(s._h_stream) + cdef vector[cydriver.CUdeviceptr] dst_ptrs + cdef vector[cydriver.CUdeviceptr] src_ptrs + cdef vector[size_t] sizes + cdef vector[size_t] attrs_idxs + dst_ptrs.resize(n) + src_ptrs.resize(n) + sizes.resize(n) + + cdef Buffer src_buf + cdef Buffer dst_buf + cdef Py_ssize_t i + + # Collapse equal neighbouring attributes into runs so a broadcast + # attribute reaches the driver once (numAttrs == 1) instead of being + # repeated per copy. attrs[k] applies to [attrsIdxs[k], attrsIdxs[k+1]). + cdef list run_starts = _attr_run_starts(attr_tuple) + cdef vector[cydriver.CUmemcpyAttributes] cu_attrs + cdef size_t num_attrs = len(run_starts) + cu_attrs.reserve(num_attrs) + attrs_idxs.reserve(num_attrs) + for i in run_starts: + cu_attrs.push_back(_to_cu_memcpy_attributes(attr_tuple[i])) + attrs_idxs.push_back(i) + + for i in range(n): + src_buf = src_bufs[i] + dst_buf = dst_bufs[i] + src_ptrs[i] = as_cu(src_buf._h_ptr) + dst_ptrs[i] = as_cu(dst_buf._h_ptr) + sizes[i] = src_buf.size + + with nogil: + HANDLE_RETURN(cydriver.cuMemcpyBatchAsync( + dst_ptrs.data(), + src_ptrs.data(), + sizes.data(), + n, + cu_attrs.data(), + attrs_idxs.data(), + num_attrs, + hstream, + )) diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyi b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyi index ca29265f103..a72bc52827f 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyi +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyi @@ -11,6 +11,7 @@ from cuda.core._memory._buffer import Buffer from cuda.core._stream import Stream from cuda.core._utils.cuda_utils import driver +_SINGLE_MANAGED_HINT = 'the ManagedBuffer instance method' def discard_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer]) -> None: """Discard a batch of managed-memory ranges. diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx index b2ecde29f39..dd9b21bf150 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx @@ -11,7 +11,11 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: from libcpp.vector cimport vector from cuda.bindings cimport cydriver -from cuda.core._memory._buffer cimport Buffer +from cuda.core._memory._buffer cimport Buffer, Buffer_coerce_batch + +# to_cumemlocation is referenced only from CUDA 13 branches. cython-lint does +# not evaluate compile-time IF blocks, so it needs a pragma to be seen as used. +from cuda.core._memory._location cimport to_cumemlocation # no-cython-lint from cuda.core._resource_handles cimport as_cu from cuda.core._stream cimport Stream, Stream_accept from cuda.core._utils.cuda_utils cimport HANDLE_RETURN @@ -52,31 +56,16 @@ cdef void _require_managed_buffer(Buffer self, str what): raise ValueError(f"{what} requires a managed-memory allocation") -cdef tuple _coerce_batch_buffers(object buffers, str what): +_SINGLE_MANAGED_HINT = "the ManagedBuffer instance method" + + +cdef inline tuple _coerce_batch_buffers(object buffers, str what): """Coerce ``buffers`` to a tuple[Buffer, ...]; rejects a single Buffer. For single-buffer operations, use the corresponding ManagedBuffer instance method instead. """ - cdef Buffer buf - cdef list out - if isinstance(buffers, Buffer): - raise TypeError( - f"{what}: pass a sequence of Buffers; for a single buffer use " - f"the ManagedBuffer instance method" - ) - if isinstance(buffers, Sequence): - if not buffers: - raise ValueError(f"{what}: empty buffers sequence") - out = [] - for t in buffers: - buf = t - out.append(buf) - return tuple(out) - raise TypeError( - f"{what}: buffers must be a sequence of Buffer, " - f"got {type(buffers).__name__}" - ) + return Buffer_coerce_batch(buffers, what, _SINGLE_MANAGED_HINT) cdef tuple _broadcast_locations(object location, Py_ssize_t n, bint allow_none, str what): @@ -91,27 +80,7 @@ cdef tuple _broadcast_locations(object location, Py_ssize_t n, bint allow_none, return tuple([coerced] * n) -IF CUDA_CORE_BUILD_MAJOR >= 13: - # Convert a _LocSpec dataclass to a cydriver.CUmemLocation struct. - cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc): - cdef str kind = loc.kind - if kind == "device": - return cydriver.CUmemLocation( - type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, - id=loc.id) - elif kind == "host": - return cydriver.CUmemLocation( - type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, - id=0) - elif kind == "host_numa": - return cydriver.CUmemLocation( - type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, - id=loc.id) - else: # host_numa_current - return cydriver.CUmemLocation( - type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, - id=0) -ELSE: +IF CUDA_CORE_BUILD_MAJOR < 13: # CUDA 12 cuMemPrefetchAsync takes a device ordinal (-1 = host). cdef inline int _to_legacy_device(object loc) except? -2: cdef str kind = loc.kind @@ -228,7 +197,7 @@ cdef void _do_single_advise(Buffer buf, object advice_value, object loc, bint al type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, id=0) else: - cu_loc = _to_cumemlocation(loc) + cu_loc = to_cumemlocation(loc) with nogil: HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, cu_loc)) ELSE: @@ -292,7 +261,7 @@ cdef void _do_single_prefetch(Buffer buf, object loc, Stream s): cdef size_t nbytes = buf._size cdef cydriver.CUstream hstream = as_cu(s._h_stream) IF CUDA_CORE_BUILD_MAJOR >= 13: - cdef cydriver.CUmemLocation cu_loc = _to_cumemlocation(loc) + cdef cydriver.CUmemLocation cu_loc = to_cumemlocation(loc) with nogil: HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, cu_loc, 0, hstream)) ELSE: @@ -365,7 +334,7 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: buf = bufs[i] ptrs[i] = as_cu(buf._h_ptr) sizes[i] = buf._size - loc_arr[i] = _to_cumemlocation(locs[i]) + loc_arr[i] = to_cumemlocation(locs[i]) loc_indices[i] = i with nogil: HANDLE_RETURN(fn( diff --git a/cuda_core/cuda/core/utils/__init__.py b/cuda_core/cuda/core/utils/__init__.py index 93a4c14c083..bc0a38f2b40 100644 --- a/cuda_core/cuda/core/utils/__init__.py +++ b/cuda_core/cuda/core/utils/__init__.py @@ -2,6 +2,12 @@ # # SPDX-License-Identifier: Apache-2.0 +from cuda.core._memory._copy_enums import ( + CopyOptions, + MemcpyOverlapMode, + MemcpySrcAccessOrder, +) +from cuda.core._memory._copy_ops import copy_batch from cuda.core._memory._managed_memory_ops import ( discard_batch, discard_prefetch_batch, @@ -19,11 +25,15 @@ ) __all__ = [ + "CopyOptions", "FileStreamProgramCache", "InMemoryProgramCache", + "MemcpyOverlapMode", + "MemcpySrcAccessOrder", "ProgramCacheResource", "StridedMemoryView", "args_viewable_as_strided_memory", + "copy_batch", "discard_batch", "discard_prefetch_batch", "make_program_cache_key", diff --git a/cuda_core/docs/source/api.rst b/cuda_core/docs/source/api.rst index e903a46a7ee..38ff695ace5 100644 --- a/cuda_core/docs/source/api.rst +++ b/cuda_core/docs/source/api.rst @@ -377,6 +377,7 @@ Utility functions :toctree: generated/ utils.args_viewable_as_strided_memory + utils.copy_batch utils.prefetch_batch utils.discard_batch utils.discard_prefetch_batch @@ -384,3 +385,20 @@ Utility functions :template: autosummary/cyclass.rst utils.StridedMemoryView + +Data transfer options +````````````````````` + +.. currentmodule:: cuda.core + +.. autosummary:: + :toctree: generated/ + + :template: dataclass.rst + + utils.CopyOptions + + :template: class.rst + + utils.MemcpySrcAccessOrder + utils.MemcpyOverlapMode diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index 120d2c2a253..b41d507699b 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -6,6 +6,29 @@ ``cuda.core`` 1.2.0 Release Notes ================================== +New features +------------ + +- Added :func:`utils.copy_batch`, which submits many buffer-to-buffer copies + in a single ``cuMemcpyBatchAsync`` call, amortizing the per-copy launch + overhead of issuing :meth:`Buffer.copy_to` in a loop. Sizes are taken from + the source buffers. Per-copy behavior is expressed with the new + :class:`utils.CopyOptions` dataclass, carrying a + :class:`utils.MemcpySrcAccessOrder` source access-order hint, optional + ``Device`` / ``Host`` location hints, and a + :class:`utils.MemcpyOverlapMode` copy-engine overlap hint. A single + ``CopyOptions`` applies to every copy; a sequence pairs by index. + Batching requires ``cuda.core`` built against CUDA 13, ``cuda.bindings`` + 13.0 or newer, and a driver reporting CUDA 13.0 or newer + (``cuDriverGetVersion() >= 13000``); ``cuda.bindings`` binds only the CUDA + 13.0 revision of the entry point, so a driver predating it is refused even + where it implements the earlier CUDA 12.8 signature. Short of that, the + copies fall back to a per-copy ``cuMemcpyAsync`` loop, and non-default + options raise ``NotImplementedError`` rather than being silently dropped. + Batched memcpy has no graph-node form and cannot be captured into a graph; + use :meth:`graph.GraphNode.memcpy` or :meth:`Buffer.copy_to` there. + (`#1333 `__) + Fixes and enhancements ---------------------- diff --git a/cuda_core/examples/batched_memcpy.py b/cuda_core/examples/batched_memcpy.py new file mode 100644 index 00000000000..ee0f2c7eb65 --- /dev/null +++ b/cuda_core/examples/batched_memcpy.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# ################################################################################ +# +# This example demonstrates the batched memory copy API (copy_batch) for +# performing multiple async memory transfers in a single driver call. It +# covers homogeneous batches (all copies share one CopyOptions), +# heterogeneous batches (per-copy attributes), and verifies equivalence +# with sequential Buffer.copy_to calls. +# +# Requires CUDA 13+ (cuMemcpyBatchAsync is not available on CUDA 12). +# +# ################################################################################ + +# /// script +# dependencies = ["cuda_bindings", "cuda_core"] +# /// + +import sys + +from cuda.core import Device, Host, LegacyPinnedMemoryResource, ManagedMemoryResource +from cuda.core.utils import CopyOptions, MemcpySrcAccessOrder, copy_batch + + +def readback(any_buf, pinned_mr, *, stream): + """Copy a buffer to a new pinned buffer and return the bytes.""" + host_buf = pinned_mr.allocate(any_buf.size) + any_buf.copy_to(host_buf, stream=stream) + stream.sync() + import ctypes + + ptr = ctypes.cast(int(host_buf.handle), ctypes.POINTER(ctypes.c_byte)) + data = ctypes.string_at(ptr, host_buf.size) + host_buf.close() + return data + + +def main(dev: Device): + dev.set_current() + stream = dev.create_stream() + pinned_mr = LegacyPinnedMemoryResource() + device_mr = dev.memory_resource + + num_copies = 4 + buf_size = 4096 + + # ---- Allocate source (pinned) and destination (device) buffers ---------- + + srcs = [] + dsts = [] + for i in range(num_copies): + src = pinned_mr.allocate(buf_size) + dst = device_mr.allocate(buf_size, stream=stream) + + # Fill each source with a distinct byte pattern so we can verify + fill_byte = (i + 1) % 256 + src.fill(fill_byte, stream=stream) + + srcs.append(src) + dsts.append(dst) + + # ---- 1. Homogeneous batch: all copies share a single CopyOptions ----- + + print("1. Homogeneous batched H2D copy...", file=sys.stderr) + + options = CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM) + copy_batch(stream, srcs, dsts, options=options) + + for i, dst in enumerate(dsts): + expected_byte = (i + 1) % 256 + data = readback(dst, pinned_mr, stream=stream) + assert all(b == expected_byte for b in data), f"Copy {i}: expected byte {expected_byte}, got {data[:8]!r}..." + + print(" All copies verified.", file=sys.stderr) + + # ---- 2. Equivalence with sequential Buffer.copy_to ---------------------- + + print("2. Verifying batched == sequential copy_to...", file=sys.stderr) + + # Re-fill sources with new patterns + for i, src in enumerate(srcs): + src.fill((i + 100) % 256, stream=stream) + + # Sequential path: individual copy_to calls + seq_dsts = [device_mr.allocate(buf_size, stream=stream) for _ in range(num_copies)] + for src, dst in zip(srcs, seq_dsts): + src.copy_to(dst, stream=stream) + + # Batched path: single copy_batch call + batch_dsts = [device_mr.allocate(buf_size, stream=stream) for _ in range(num_copies)] + copy_batch(stream, srcs, batch_dsts, options=CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM)) + + # Compare results + for i in range(num_copies): + seq_data = readback(seq_dsts[i], pinned_mr, stream=stream) + batch_data = readback(batch_dsts[i], pinned_mr, stream=stream) + assert seq_data == batch_data, f"Copy {i}: sequential and batched results differ" + + print(" Batched and sequential results match.", file=sys.stderr) + + # ---- 3. Heterogeneous batch: per-copy attributes ------------------------ + # + # src_access_order controls how the driver accesses source memory: + # STREAM - source read respects stream ordering (pinned/device memory) + # DURING_API_CALL - source read during the API call itself (ephemeral host ptrs) + # ANY - driver picks best strategy (pageable or HMM-backed memory) + + print("3. Heterogeneous batch with per-copy attributes...", file=sys.stderr) + + per_copy_options = [ + CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM), + CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY), + CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM), + CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY), + ] + hetero_dsts = [device_mr.allocate(buf_size, stream=stream) for _ in range(num_copies)] + copy_batch(stream, srcs, hetero_dsts, options=per_copy_options) + + for i in range(num_copies): + expected_byte = (i + 100) % 256 + data = readback(hetero_dsts[i], pinned_mr, stream=stream) + assert all(b == expected_byte for b in data), f"Heterogeneous copy {i}: expected byte {expected_byte}" + + print(" Heterogeneous batch verified.", file=sys.stderr) + + # ---- 4. Location hints with managed memory ------------------------------ + # + # When copying managed-memory buffers, src_location_hint and + # dst_location_hint tell the driver where the data currently lives and + # where it is going, enabling optimized transfer paths. + + print("4. Batched copy with location hints (managed memory)...", file=sys.stderr) + + managed_mr = ManagedMemoryResource() + managed_srcs = [managed_mr.allocate(buf_size, stream=stream) for _ in range(2)] + managed_dsts = [managed_mr.allocate(buf_size, stream=stream) for _ in range(2)] + + for i, src in enumerate(managed_srcs): + src.fill((i + 200) % 256, stream=stream) + + hint_options = CopyOptions( + src_access_order=MemcpySrcAccessOrder.STREAM, + src_location_hint=dev, + dst_location_hint=Host(), + ) + copy_batch(stream, managed_srcs, managed_dsts, options=hint_options) + + for i in range(2): + expected_byte = (i + 200) % 256 + data = readback(managed_dsts[i], pinned_mr, stream=stream) + assert all(b == expected_byte for b in data), f"Managed copy {i}: expected byte {expected_byte}" + + print(" Location-hinted batch verified.", file=sys.stderr) + + # ---- Cleanup ------------------------------------------------------------ + + all_bufs = srcs + dsts + seq_dsts + batch_dsts + hetero_dsts + managed_srcs + managed_dsts + for buf in all_bufs: + buf.close(stream) + stream.close() + + print("Batched memcpy example completed!") + + +if __name__ == "__main__": + main(Device(0)) diff --git a/cuda_core/tests/conftest.py b/cuda_core/tests/conftest.py index dfe97b265eb..b212633ebcf 100644 --- a/cuda_core/tests/conftest.py +++ b/cuda_core/tests/conftest.py @@ -179,6 +179,8 @@ def create_managed_memory_resource_or_skip(*args, xfail_device=None, **kwargs): except RuntimeError as e: if "requires CUDA 13.0" in str(e): pytest.skip("ManagedMemoryResource requires CUDA 13.0 or later") + if "concurrent managed access is not available" in str(e).lower(): + pytest.skip("Device does not support concurrent managed memory access") raise diff --git a/cuda_core/tests/example_tests/test_basic_examples.py b/cuda_core/tests/example_tests/test_basic_examples.py index bf423758366..4f4fd79f371 100644 --- a/cuda_core/tests/example_tests/test_basic_examples.py +++ b/cuda_core/tests/example_tests/test_basic_examples.py @@ -74,8 +74,25 @@ def has_recent_memory_pool_support() -> bool: # Specific system requirements for each of the examples. +def has_copy_batch_support() -> bool: + """Check if cuMemcpyBatchAsync is available (CUDA 13+).""" + from cuda.core._utils.version import binding_version + + if binding_version() < (13, 0, 0): + return False + try: + from cuda.bindings import driver + + if not hasattr(driver, "cuMemcpyBatchAsync"): + return False + except AttributeError: + return False + return True + + SYSTEM_REQUIREMENTS = { "memory_pool_resources.py": has_recent_memory_pool_support, + "batched_memcpy.py": lambda: has_copy_batch_support() and has_recent_memory_pool_support(), "gl_interop_plasma.py": has_display, "gl_interop_fluid.py": has_display, "gl_interop_mipmap_lod.py": has_display, diff --git a/cuda_core/tests/helpers/copy_batch.py b/cuda_core/tests/helpers/copy_batch.py new file mode 100644 index 00000000000..0eafff2094e --- /dev/null +++ b/cuda_core/tests/helpers/copy_batch.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Shared constants and helpers for the ``copy_batch`` tests. + +Fixtures live in ``tests/memory/conftest.py``; this module holds the +pieces that tests import by name. +""" + +from cuda.core import LegacyPinnedMemoryResource +from helpers.buffers import compare_equal_buffers, make_scratch_buffer + +COPY_BATCH_SIZE = 4096 +COPY_BATCH_COUNT = 4 + +# Matches the UserWarning raised when prefer_overlap_with_compute is +# requested on a discrete GPU. Tests that are not about the warning +# silence it so they stay green under -W error. +OVERLAP_WARNING_FILTER = "ignore:overlap_mode:UserWarning" + + +def assert_managed_holds(dev, buf, value, *, stream): + """Assert a managed buffer holds ``value``. + + Reads via an explicit device-to-host copy rather than dereferencing + the managed pointer from the host. Managed pages carry residency and + ``cuMemAdvise`` state that earlier tests in the suite can leave + behind, which makes direct host reads order-dependent. Also avoids + ``compare_buffer_to_constant``, which resolves a ``Device`` from + ``memory_resource.device_id`` -- that is -1 for + ``ManagedMemoryResource``. + """ + host = LegacyPinnedMemoryResource().allocate(buf.size) + expected = make_scratch_buffer(dev, value, buf.size) + try: + buf.copy_to(host, stream=stream) + stream.sync() + assert compare_equal_buffers(expected, host) + finally: + expected.close() + host.close(stream) + stream.sync() diff --git a/cuda_core/tests/memory/__init__.py b/cuda_core/tests/memory/__init__.py new file mode 100644 index 00000000000..61904faa870 --- /dev/null +++ b/cuda_core/tests/memory/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Marks `tests/memory` as a package. Under pytest's default "prepend" import +# mode, a test file's sys.path entry is the first parent directory *without* an +# __init__.py. Adding this file moves that entry up from `tests/memory` to +# `tests/`, which is what we want for two reasons: +# +# * `tests/memory` is no longer on sys.path, so the `from conftest import ...` +# in test_managed_ops.py resolves to the root tests/conftest.py. Without it +# the local memory/conftest.py wins and the import fails with ImportError. +# * Modules are named `memory.test_x` rather than `test_x`, so a file basename +# reused under another directory cannot collide. +# +# Both follow from module identity tracking the directory layout instead of +# whichever directory happens to land on sys.path. `tests/memory_ipc/` carries +# an __init__.py for the same reasons. diff --git a/cuda_core/tests/memory/conftest.py b/cuda_core/tests/memory/conftest.py new file mode 100644 index 00000000000..c0683f9aff3 --- /dev/null +++ b/cuda_core/tests/memory/conftest.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-directory conftest for the ``copy_batch`` test modules. + +Provides the device, stream and buffer fixtures shared by +``test_copy_batch.py`` (data movement) and ``test_copy_batch_options.py`` +(options and validation). +""" + +import pytest +from helpers.copy_batch import COPY_BATCH_COUNT, COPY_BATCH_SIZE + +from cuda.core import Device, LegacyPinnedMemoryResource +from cuda.core._memory._copy_ops import _batch_entry_point_in_use + + +@pytest.fixture +def copy_batch_device(init_cuda): + """``copy_batch`` works on every supported toolkit, so this never skips. + + Only non-default ``CopyOptions`` need CUDA 13; those tests take + ``requires_copy_options`` as well. + """ + device = Device() + device.set_current() + return device + + +@pytest.fixture +def requires_copy_options(): + """Skip when ``CopyOptions`` cannot reach the driver. + + Options travel only through ``cuMemcpyBatchAsync``. Where that is + unavailable, ``copy_batch`` falls back to per-copy ``cuMemcpyAsync``, + which cannot convey attributes, so non-default options are rejected. + """ + if not _batch_entry_point_in_use(): + pytest.skip( + "non-default CopyOptions requires cuda.core built against CUDA 13, " + "cuda.bindings 13.0+, and a driver reporting CUDA 13.0 or newer" + ) + + +@pytest.fixture +def copy_stream(copy_batch_device): + """The single stream used for both allocation and copies in a test. + + Stream-ordered pool allocations are only guaranteed usable on the + stream that allocated them, so tests allocate and copy on this one + stream rather than mixing it with ``device.default_stream``. + """ + s = copy_batch_device.create_stream() + yield s + s.close() + + +@pytest.fixture +def h2d_bufs(copy_batch_device, copy_stream): + """Pinned-host source / device destination pairs.""" + pinned_mr = LegacyPinnedMemoryResource() + device_mr = copy_batch_device.memory_resource + + srcs = [pinned_mr.allocate(COPY_BATCH_SIZE) for _ in range(COPY_BATCH_COUNT)] + dsts = [device_mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in range(COPY_BATCH_COUNT)] + + yield srcs, dsts + + for buf in srcs + dsts: + buf.close(copy_stream) + copy_stream.sync() + + +@pytest.fixture +def device_bufs(copy_batch_device, copy_stream): + """Device source / device destination pairs.""" + device_mr = copy_batch_device.memory_resource + + srcs = [device_mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in range(COPY_BATCH_COUNT)] + dsts = [device_mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in range(COPY_BATCH_COUNT)] + + yield srcs, dsts + + for buf in srcs + dsts: + buf.close(copy_stream) + copy_stream.sync() diff --git a/cuda_core/tests/memory/test_copy_batch.py b/cuda_core/tests/memory/test_copy_batch.py new file mode 100644 index 00000000000..04c7cc41af9 --- /dev/null +++ b/cuda_core/tests/memory/test_copy_batch.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Data movement behaviour of ``copy_batch``. + +Covers that the right bytes reach the right destination, that batched +results agree with the per-buffer ``Buffer.copy_to`` path, and that the +batch is correctly ordered on its stream. +""" + +import pytest +from helpers.buffers import ( + compare_buffer_to_constant, + compare_equal_buffers, + make_scratch_buffer, + set_buffer, +) +from helpers.copy_batch import COPY_BATCH_SIZE + +from cuda.core import LegacyPinnedMemoryResource +from cuda.core.utils import copy_batch + + +class TestCopyBatchCore: + """Each transfer direction moves the expected bytes.""" + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_h2d_batch(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + for i, src in enumerate(srcs): + set_buffer(src, i + 1) + + copy_batch(copy_stream, srcs, dsts) + copy_stream.sync() + + for i, dst in enumerate(dsts): + assert compare_buffer_to_constant(dst, i + 1) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_d2h_batch(self, copy_batch_device, h2d_bufs, copy_stream): + dev = copy_batch_device + _, device_dsts = h2d_bufs + pinned_mr = LegacyPinnedMemoryResource() + + for i, buf in enumerate(device_dsts): + buf.fill(i + 10, stream=copy_stream) + + host_bufs = [pinned_mr.allocate(COPY_BATCH_SIZE) for _ in device_dsts] + copy_batch(copy_stream, device_dsts, host_bufs) + copy_stream.sync() + + for i, host_buf in enumerate(host_bufs): + expected = make_scratch_buffer(dev, i + 10, COPY_BATCH_SIZE) + assert compare_equal_buffers(expected, host_buf) + expected.close() + host_buf.close(copy_stream) + copy_stream.sync() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_d2d_batch(self, device_bufs, copy_stream): + srcs, dsts = device_bufs + for i, src in enumerate(srcs): + src.fill(i + 20, stream=copy_stream) + + copy_batch(copy_stream, srcs, dsts) + copy_stream.sync() + + for i, dst in enumerate(dsts): + assert compare_buffer_to_constant(dst, i + 20) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_various_sizes(self, copy_batch_device, copy_stream): + pinned_mr = LegacyPinnedMemoryResource() + device_mr = copy_batch_device.memory_resource + sizes = [1024, 2048, 512, 4096] + + srcs = [pinned_mr.allocate(size) for size in sizes] + dsts = [device_mr.allocate(size, stream=copy_stream) for size in sizes] + for i, src in enumerate(srcs): + set_buffer(src, i + 1) + + copy_batch(copy_stream, srcs, dsts) + copy_stream.sync() + + for i, dst in enumerate(dsts): + assert compare_buffer_to_constant(dst, i + 1) + + for buf in srcs + dsts: + buf.close(copy_stream) + copy_stream.sync() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_single_element_batch(self, copy_batch_device, copy_stream): + """A one-element batch is legal; only a bare Buffer is rejected.""" + pinned_mr = LegacyPinnedMemoryResource() + src = pinned_mr.allocate(COPY_BATCH_SIZE) + dst = copy_batch_device.memory_resource.allocate(COPY_BATCH_SIZE, stream=copy_stream) + set_buffer(src, 7) + + copy_batch(copy_stream, [src], [dst]) + copy_stream.sync() + + assert compare_buffer_to_constant(dst, 7) + src.close(copy_stream) + dst.close(copy_stream) + copy_stream.sync() + + +class TestCopyBatchEquivalence: + """Batched results must agree with the already-tested per-buffer path. + + ``Buffer.copy_to`` and ``Buffer.copy_from`` have their own coverage in + ``tests/test_memory.py``, so agreement between the two paths is the + property under test here. + """ + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_batch_matches_sequential_copy_to(self, copy_batch_device, h2d_bufs, copy_stream): + srcs, _ = h2d_bufs + device_mr = copy_batch_device.memory_resource + pinned_mr = LegacyPinnedMemoryResource() + + for i, src in enumerate(srcs): + set_buffer(src, i + 50) + + seq_dsts = [device_mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in srcs] + for src, dst in zip(srcs, seq_dsts): + src.copy_to(dst, stream=copy_stream) + + batch_dsts = [device_mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in srcs] + copy_batch(copy_stream, srcs, batch_dsts) + copy_stream.sync() + + for seq_dst, batch_dst in zip(seq_dsts, batch_dsts): + seq_host = pinned_mr.allocate(COPY_BATCH_SIZE) + batch_host = pinned_mr.allocate(COPY_BATCH_SIZE) + seq_dst.copy_to(seq_host, stream=copy_stream) + batch_dst.copy_to(batch_host, stream=copy_stream) + copy_stream.sync() + assert compare_equal_buffers(seq_host, batch_host) + seq_host.close(copy_stream) + batch_host.close(copy_stream) + + for buf in seq_dsts + batch_dsts: + buf.close(copy_stream) + copy_stream.sync() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_batch_matches_sequential_varied_sizes(self, copy_batch_device, copy_stream): + device_mr = copy_batch_device.memory_resource + pinned_mr = LegacyPinnedMemoryResource() + sizes = [1024, 2048, 512] + + srcs = [pinned_mr.allocate(size) for size in sizes] + for i, src in enumerate(srcs): + set_buffer(src, i + 60) + + seq_dsts = [device_mr.allocate(size, stream=copy_stream) for size in sizes] + for src, dst in zip(srcs, seq_dsts): + src.copy_to(dst, stream=copy_stream) + + batch_dsts = [device_mr.allocate(size, stream=copy_stream) for size in sizes] + copy_batch(copy_stream, srcs, batch_dsts) + copy_stream.sync() + + for size, seq_dst, batch_dst in zip(sizes, seq_dsts, batch_dsts): + seq_host = pinned_mr.allocate(size) + batch_host = pinned_mr.allocate(size) + seq_dst.copy_to(seq_host, stream=copy_stream) + batch_dst.copy_to(batch_host, stream=copy_stream) + copy_stream.sync() + assert compare_equal_buffers(seq_host, batch_host) + seq_host.close(copy_stream) + batch_host.close(copy_stream) + + for buf in srcs + seq_dsts + batch_dsts: + buf.close(copy_stream) + copy_stream.sync() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_batch_matches_sequential_d2d(self, copy_batch_device, device_bufs, copy_stream): + srcs, seq_dsts = device_bufs + device_mr = copy_batch_device.memory_resource + pinned_mr = LegacyPinnedMemoryResource() + + for i, src in enumerate(srcs): + src.fill(i + 70, stream=copy_stream) + + for src, dst in zip(srcs, seq_dsts): + src.copy_to(dst, stream=copy_stream) + + batch_dsts = [device_mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in srcs] + copy_batch(copy_stream, srcs, batch_dsts) + copy_stream.sync() + + for seq_dst, batch_dst in zip(seq_dsts, batch_dsts): + seq_host = pinned_mr.allocate(COPY_BATCH_SIZE) + batch_host = pinned_mr.allocate(COPY_BATCH_SIZE) + seq_dst.copy_to(seq_host, stream=copy_stream) + batch_dst.copy_to(batch_host, stream=copy_stream) + copy_stream.sync() + assert compare_equal_buffers(seq_host, batch_host) + seq_host.close(copy_stream) + batch_host.close(copy_stream) + + for buf in batch_dsts: + buf.close(copy_stream) + copy_stream.sync() + + +class TestCopyBatchStreamSemantics: + """Where the batch sits in stream order, and what it cannot be part of.""" + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_ordered_between_prior_and_later_stream_work(self, device_bufs, copy_stream): + """The batch must observe prior stream work and precede later work. + + Each source is filled with ``before``, copied, then refilled with + ``after`` -- all enqueued on one stream with no intervening sync. + Destinations holding ``before`` prove the copy ran after the first + fill and before the second, rather than racing either. + """ + srcs, dsts = device_bufs + before, after = 11, 22 + + for src in srcs: + src.fill(before, stream=copy_stream) + copy_batch(copy_stream, srcs, dsts) + for src in srcs: + src.fill(after, stream=copy_stream) + + copy_stream.sync() + + for dst in dsts: + assert compare_buffer_to_constant(dst, before) + for src in srcs: + assert compare_buffer_to_constant(src, after) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_graph_builder_is_rejected(self, copy_batch_device, device_bufs, copy_stream): + """Batched memcpy cannot be captured into a graph. + + ``cuMemcpyBatchAsync`` has no graph-node form and the driver + rejects it mid-capture, so ``copy_batch`` is typed to take only a + ``Stream`` and refuses a ``GraphBuilder`` at the boundary rather + than failing later with ``CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED``. + Use ``GraphNode.memcpy`` or per-buffer ``Buffer.copy_to`` to build + copies into a graph. + """ + srcs, dsts = device_bufs + gb = copy_batch_device.create_graph_builder().begin_building() + try: + with pytest.raises(TypeError, match="Argument 'stream' has incorrect type"): + copy_batch(gb, srcs, dsts) + finally: + # Nothing was captured, so the builder still ends cleanly. + gb.end_building() + gb.close() diff --git a/cuda_core/tests/memory/test_copy_batch_options.py b/cuda_core/tests/memory/test_copy_batch_options.py new file mode 100644 index 00000000000..7b6065cccaf --- /dev/null +++ b/cuda_core/tests/memory/test_copy_batch_options.py @@ -0,0 +1,406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``CopyOptions`` handling and argument validation for ``copy_batch``. + +Covers how options are encoded into the driver's attribute runs, how each +option field behaves, and every rejection path. +""" + +import warnings + +import pytest + +# Shared with test_managed_ops.py: handles the CUDA 13 requirement, mempool +# OOM, and CUDA_ERROR_NOT_SUPPORTED (managed pools are unavailable on +# Windows), so the location-hint tests skip rather than error there. +from conftest import create_managed_memory_resource_or_skip +from helpers.buffers import compare_buffer_to_constant, set_buffer +from helpers.copy_batch import ( + COPY_BATCH_SIZE, + OVERLAP_WARNING_FILTER, + assert_managed_holds, +) + +from cuda.core import Device, Host, LegacyPinnedMemoryResource +from cuda.core._memory._copy_enums import _attr_run_starts +from cuda.core._memory._copy_ops import _batch_entry_point_in_use, _normalize_copy_options +from cuda.core.utils import ( + CopyOptions, + MemcpyOverlapMode, + MemcpySrcAccessOrder, + copy_batch, +) + + +class TestOptionsEncoding: + """How ``options`` becomes the driver's ``attrs`` / ``attrsIdxs`` pair. + + Pure logic, no CUDA. This is the only place the effect of ``options`` + is observable: they are hints that change how the driver stages a + transfer, never the bytes it produces, so no data comparison can + distinguish an option that was applied from one that was dropped. + """ + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_scalar_broadcasts_to_every_copy(self): + """A scalar must reach all N copies, not just the first.""" + n = 4 + scalar = CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY) + + # copy_batch expands the scalar to one entry per copy... + assert _normalize_copy_options(scalar, n) == (scalar,) * n + # ...and the encoder collapses those to a single driver attribute. + assert _attr_run_starts(_normalize_copy_options(scalar, n)) == [0] + + # An explicit list of the same option is indistinguishable. + assert _normalize_copy_options([scalar] * n, n) == _normalize_copy_options(scalar, n) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_none_broadcasts_defaults(self): + assert _normalize_copy_options(None, 3) == (CopyOptions(),) * 3 + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_sequence_is_never_broadcast(self): + """A sequence pairs by index, so a short one is an error.""" + scalar = CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY) + with pytest.raises(ValueError, match="options length"): + _normalize_copy_options([scalar], 4) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_equal_but_distinct_instances_collapse(self): + # Structural equality, not identity, drives the collapse. + attrs = [CopyOptions(src_access_order="stream") for _ in range(3)] + assert len({id(a) for a in attrs}) == 3 + assert _attr_run_starts(attrs) == [0] + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_all_distinct_yields_one_run_each(self): + attrs = [ + CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM), + CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY), + CopyOptions(src_access_order=MemcpySrcAccessOrder.DURING_API_CALL), + ] + assert _attr_run_starts(attrs) == [0, 1, 2] + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_adjacent_runs_are_grouped(self): + stream_attr = CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM) + any_attr = CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY) + attrs = [stream_attr, stream_attr, any_attr, any_attr, stream_attr] + # Runs start at 0 (stream), 2 (any) and 4 (stream again). + assert _attr_run_starts(attrs) == [0, 2, 4] + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_single_element(self): + assert _attr_run_starts([CopyOptions()]) == [0] + + +class TestCopyBatchOptions: + """Each ``CopyOptions`` field is accepted and does not corrupt the copy.""" + + @pytest.mark.agent_authored(model="Claude Opus 5") + @pytest.mark.parametrize( + ("order", "marker"), + [ + (MemcpySrcAccessOrder.STREAM, 31), + (MemcpySrcAccessOrder.DURING_API_CALL, 32), + (MemcpySrcAccessOrder.ANY, 33), + ], + ) + def test_src_access_order(self, requires_copy_options, h2d_bufs, copy_stream, order, marker): + srcs, dsts = h2d_bufs + for i, src in enumerate(srcs): + set_buffer(src, i + marker) + + copy_batch(copy_stream, srcs, dsts, options=CopyOptions(src_access_order=order)) + copy_stream.sync() + + for i, dst in enumerate(dsts): + assert compare_buffer_to_constant(dst, i + marker) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_per_copy_options(self, requires_copy_options, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + for i, src in enumerate(srcs): + set_buffer(src, i + 40) + + per_copy_options = [ + CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM), + CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY), + CopyOptions(src_access_order=MemcpySrcAccessOrder.DURING_API_CALL), + CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM), + ] + # The encoding itself is covered by TestOptionsEncoding; here the + # point is that distinct per-copy options do not corrupt the data. + copy_batch(copy_stream, srcs, dsts, options=per_copy_options) + copy_stream.sync() + + for i, dst in enumerate(dsts): + assert compare_buffer_to_constant(dst, i + 40) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_location_hints_do_not_corrupt_copy(self, requires_copy_options, copy_batch_device, copy_stream): + """Device and host hints are accepted and leave the bytes intact. + + Hints only steer how the driver stages a transfer, so no data + comparison can show one was *applied*; what this catches is a hint + that errors or corrupts. It is also the only test that drives the + ``device`` and ``host`` branches of ``to_cumemlocation`` and the + ``src_location_hint`` path through ``copy_batch``. + """ + dev = copy_batch_device + mr = create_managed_memory_resource_or_skip() + srcs = [mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in range(2)] + dsts = [mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in range(2)] + + for i, src in enumerate(srcs): + src.fill(i + 80, stream=copy_stream) + + options = CopyOptions( + src_access_order=MemcpySrcAccessOrder.STREAM, + src_location_hint=dev, + dst_location_hint=Host(), + ) + copy_batch(copy_stream, srcs, dsts, options=options) + copy_stream.sync() + + for i, dst in enumerate(dsts): + assert_managed_holds(dev, dst, i + 80, stream=copy_stream) + + for buf in srcs + dsts: + buf.close(copy_stream) + copy_stream.sync() + mr.close() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_host_numa_location_hint(self, requires_copy_options, copy_batch_device, copy_stream): + """A NUMA-specific host hint is accepted and does not corrupt the copy. + + ``requires_copy_options`` already implies CUDA 13, where + ``Host(numa_id=...)`` is representable; the CUDA 12 rejection is + covered by ``_coerce_location``'s own tests. + """ + dev = copy_batch_device + mr = create_managed_memory_resource_or_skip() + srcs = [mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in range(2)] + dsts = [mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in range(2)] + for i, src in enumerate(srcs): + src.fill(i + 85, stream=copy_stream) + + copy_batch(copy_stream, srcs, dsts, options=CopyOptions(dst_location_hint=Host(numa_id=0))) + copy_stream.sync() + for i, dst in enumerate(dsts): + assert_managed_holds(dev, dst, i + 85, stream=copy_stream) + + for buf in srcs + dsts: + buf.close(copy_stream) + copy_stream.sync() + mr.close() + + @pytest.mark.agent_authored(model="Claude Opus 5") + @pytest.mark.filterwarnings(OVERLAP_WARNING_FILTER) + def test_overlap_mode_copies_correctly(self, requires_copy_options, h2d_bufs, copy_stream): + """The overlap hint is advisory and must not change the bytes copied.""" + srcs, dsts = h2d_bufs + for i, src in enumerate(srcs): + set_buffer(src, i + 90) + + copy_batch( + copy_stream, + srcs, + dsts, + options=CopyOptions(overlap_mode=MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE), + ) + copy_stream.sync() + + for i, dst in enumerate(dsts): + assert compare_buffer_to_constant(dst, i + 90) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_overlap_mode_warns_only_on_discrete_gpu( + self, requires_copy_options, copy_batch_device, h2d_bufs, copy_stream + ): + srcs, dsts = h2d_bufs + options = CopyOptions(overlap_mode=MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE) + + if copy_batch_device.properties.integrated: + # Tegra honours the hint, so no warning should be emitted. + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + copy_batch(copy_stream, srcs, dsts, options=options) + else: + with pytest.warns(UserWarning, match="non-integrated"): + copy_batch(copy_stream, srcs, dsts, options=options) + copy_stream.sync() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_default_overlap_mode_does_not_warn(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + copy_batch(copy_stream, srcs, dsts, options=CopyOptions()) + copy_stream.sync() + + +class TestCopyOptionsValidation: + """``CopyOptions`` rejects invalid enum values at construction.""" + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_invalid_access_order(self): + with pytest.raises(ValueError, match="invalid src_access_order"): + CopyOptions(src_access_order="invalid_order") + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_invalid_overlap_mode(self): + with pytest.raises(ValueError, match="invalid overlap_mode"): + CopyOptions(overlap_mode="invalid_mode") + + +class TestCopyBatchValidation: + """``copy_batch`` rejects malformed buffer and option arguments.""" + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_rejects_single_buffer(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + + with pytest.raises(TypeError, match="sequence of Buffers"): + copy_batch(copy_stream, srcs[0], dsts) + + with pytest.raises(TypeError, match="sequence of Buffers"): + copy_batch(copy_stream, srcs, dsts[0]) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_rejects_empty_sequence(self, h2d_bufs, copy_stream): + srcs, _ = h2d_bufs + + with pytest.raises(ValueError, match="empty buffers sequence"): + copy_batch(copy_stream, [], []) + + with pytest.raises(ValueError, match="empty buffers sequence"): + copy_batch(copy_stream, srcs, []) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_rejects_non_buffer_element(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + + with pytest.raises(TypeError, match="expected Buffer, got int"): + copy_batch(copy_stream, [srcs[0], 42], dsts[:2]) + + with pytest.raises(TypeError, match="expected Buffer, got NoneType"): + copy_batch(copy_stream, srcs[:2], [dsts[0], None]) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_rejects_non_sequence(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + + with pytest.raises(TypeError, match="must be a sequence of Buffer"): + copy_batch(copy_stream, 42, dsts) + + with pytest.raises(TypeError, match="must be a sequence of Buffer"): + copy_batch(copy_stream, srcs, None) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_length_mismatch(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + + with pytest.raises(ValueError, match="does not match dsts length"): + copy_batch(copy_stream, srcs[:2], dsts[:3]) + + @pytest.mark.agent_authored(model="Claude Opus 5") + @pytest.mark.parametrize(("src_size", "dst_size"), [(1024, 2048), (2048, 1024)]) + def test_size_mismatch(self, copy_batch_device, copy_stream, src_size, dst_size): + """Sizes come from the buffers, so any inequality is an error.""" + pinned_mr = LegacyPinnedMemoryResource() + src = pinned_mr.allocate(src_size) + dst = copy_batch_device.memory_resource.allocate(dst_size, stream=copy_stream) + + with pytest.raises(ValueError, match="size mismatch at index 0"): + copy_batch(copy_stream, [src], [dst]) + + src.close(copy_stream) + dst.close(copy_stream) + copy_stream.sync() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_options_length_mismatch(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + + with pytest.raises(ValueError, match="options length"): + copy_batch(copy_stream, srcs, dsts, options=[CopyOptions()] * 3) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_rejects_bad_options_type(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + + with pytest.raises(TypeError, match="options must be CopyOptions"): + copy_batch(copy_stream, srcs, dsts, options=42) + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_rejects_bad_options_element(self, h2d_bufs, copy_stream): + srcs, dsts = h2d_bufs + bad = [CopyOptions()] * (len(srcs) - 1) + ["nope"] + + with pytest.raises(TypeError, match="each options element must be CopyOptions"): + copy_batch(copy_stream, srcs, dsts, options=bad) + + +class TestPerCopyFallback: + """Behaviour where ``cuMemcpyBatchAsync`` is unavailable. + + Reached whenever any of the three requirements for + ``cuMemcpyBatchAsync`` is unmet: ``cuda.core`` built against CUDA 13, + ``cuda.bindings`` 13.0+, and a driver reporting CUDA 13.0 or newer. + Skipped when the batched entry point is actually in use. + """ + + @pytest.fixture(autouse=True) + def _skip_if_batched(self): + if _batch_entry_point_in_use(): + pytest.skip("cuMemcpyBatchAsync is in use; fallback path not exercised") + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_default_options_still_copy(self, init_cuda): + """The copies still happen, just one driver call at a time.""" + device = Device() + device.set_current() + stream = device.create_stream() + pinned_mr = LegacyPinnedMemoryResource() + srcs = [pinned_mr.allocate(COPY_BATCH_SIZE) for _ in range(3)] + dsts = [device.memory_resource.allocate(COPY_BATCH_SIZE, stream=stream) for _ in srcs] + for i, src in enumerate(srcs): + set_buffer(src, i + 5) + + copy_batch(stream, srcs, dsts) + stream.sync() + + for i, dst in enumerate(dsts): + assert compare_buffer_to_constant(dst, i + 5) + + for buf in srcs + dsts: + buf.close(stream) + stream.sync() + stream.close() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_non_default_options_are_rejected(self, init_cuda): + """Options have no per-copy equivalent, so they must not be dropped.""" + device = Device() + device.set_current() + stream = device.create_stream() + pinned_mr = LegacyPinnedMemoryResource() + src = pinned_mr.allocate(1024) + dst = device.memory_resource.allocate(1024, stream=stream) + + with pytest.raises(NotImplementedError, match="non-default CopyOptions"): + copy_batch( + stream, + [src], + [dst], + options=CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY), + ) + + src.close(stream) + dst.close(stream) + stream.sync() + stream.close()