From bf41f9351f591df3a127ff5e71ffffa35e0b4c66 Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Mon, 10 Aug 2026 10:39:18 -0700 Subject: [PATCH 1/6] cuda.core: Add copy_batch to cuda.core.utils --- cuda_core/cuda/core/_memory/_copy_enums.py | 147 +++++++ cuda_core/cuda/core/_memory/_copy_ops.pyi | 46 +++ cuda_core/cuda/core/_memory/_copy_ops.pyx | 266 ++++++++++++ cuda_core/cuda/core/utils/__init__.py | 10 + cuda_core/docs/source/api.rst | 18 + cuda_core/examples/batched_memcpy.py | 168 ++++++++ .../example_tests/test_basic_examples.py | 17 + cuda_core/tests/helpers/copy_batch.py | 60 +++ cuda_core/tests/memory/__init__.py | 3 + cuda_core/tests/memory/conftest.py | 71 ++++ cuda_core/tests/memory/test_copy_batch.py | 288 +++++++++++++ .../tests/memory/test_copy_batch_options.py | 382 ++++++++++++++++++ 12 files changed, 1476 insertions(+) create mode 100644 cuda_core/cuda/core/_memory/_copy_enums.py create mode 100644 cuda_core/cuda/core/_memory/_copy_ops.pyi create mode 100644 cuda_core/cuda/core/_memory/_copy_ops.pyx create mode 100644 cuda_core/examples/batched_memcpy.py create mode 100644 cuda_core/tests/helpers/copy_batch.py create mode 100644 cuda_core/tests/memory/__init__.py create mode 100644 cuda_core/tests/memory/conftest.py create mode 100644 cuda_core/tests/memory/test_copy_batch.py create mode 100644 cuda_core/tests/memory/test_copy_batch_options.py 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..823c0e3c965 --- /dev/null +++ b/cuda_core/cuda/core/_memory/_copy_enums.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import dataclasses +import functools +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 + +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): + # Validate enum fields while still in __init__ (frozen dataclass). + # Use __setattr__ because fields are frozen. + 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.""" + return _src_access_order_to_cu()[MemcpySrcAccessOrder(self.src_access_order)] + + def _to_driver_flags(self) -> int: + """Return the driver CUmemcpyFlags value.""" + return _overlap_mode_to_cu()[MemcpyOverlapMode(self.overlap_mode)] + + +# Bridges between the public StrEnums and the driver integer values. Built on +# first use rather than at import: the CUmemcpy* enums only exist on toolkits +# that ship the batched memcpy entry points, and importing cuda.core must not +# depend on them. +# +# Keyed by ``str`` rather than by the enum: under ``python_version = "3.10"`` +# mypy resolves ``StrEnum`` to the unstubbed ``backports.strenum`` shim and so +# infers the members as plain ``str``. StrEnum members are ``str`` instances, +# so this annotation is accurate on every supported version. +@functools.cache +def _src_access_order_to_cu() -> dict[str, int]: + cu = driver.CUmemcpySrcAccessOrder + return { + MemcpySrcAccessOrder.STREAM: int(cu.CU_MEMCPY_SRC_ACCESS_ORDER_STREAM), + MemcpySrcAccessOrder.DURING_API_CALL: int(cu.CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL), + MemcpySrcAccessOrder.ANY: int(cu.CU_MEMCPY_SRC_ACCESS_ORDER_ANY), + } + + +@functools.cache +def _overlap_mode_to_cu() -> dict[str, int]: + cu = driver.CUmemcpyFlags + return { + MemcpyOverlapMode.DEFAULT: int(cu.CU_MEMCPY_FLAG_DEFAULT), + MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE: int(cu.CU_MEMCPY_FLAG_PREFER_OVERLAP_WITH_COMPUTE), + } + + +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..a0ce96ce4dd --- /dev/null +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyi @@ -0,0 +1,46 @@ +# 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 + + +def copy_batch(stream: object, srcs: Sequence[Buffer], dsts: Sequence[Buffer], *, options: object=None) -> None: + """Copy a batch of buffers asynchronously. + + Requires CUDA 13+. 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. Passing a + :class:`~graph.GraphBuilder` raises ``CUDAError`` + (``CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED``). + 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 + ------ + NotImplementedError + On a CUDA 12 build of ``cuda.core``. + ValueError + If lengths or sizes mismatch. + TypeError + If a single Buffer is passed instead of a sequence. + 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..fc2d7477506 --- /dev/null +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyx @@ -0,0 +1,266 @@ +# 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 + +# as_cu, HANDLE_RETURN and _attr_run_starts are referenced only from the +# CUDA 13 branch of _do_copy_batch. cython-lint does not evaluate +# compile-time IF blocks, so it needs a pragma to see them as used. +from cuda.bindings cimport cydriver +from cuda.core._memory._buffer cimport Buffer +from cuda.core._resource_handles cimport as_cu # no-cython-lint +from cuda.core._stream cimport Stream, Stream_accept +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN # 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 + +cdef tuple _coerce_batch_buffers(object buffers, str what): + """Coerce ``buffers`` to a tuple[Buffer, ...]; rejects a single Buffer.""" + cdef list out + if isinstance(buffers, Buffer): + raise TypeError( + f"{what}: pass a sequence of Buffers; for a single buffer use " + f"the Buffer.copy_to / Buffer.copy_from instance method" + ) + if isinstance(buffers, Sequence): + if not buffers: + raise ValueError(f"{what}: empty buffers sequence") + out = [] + for t in buffers: + if not isinstance(t, Buffer): + raise TypeError( + f"{what}: expected Buffer, got {type(t).__name__}" + ) + out.append(t) + return tuple(out) + raise TypeError( + f"{what}: buffers must be a sequence of Buffer, " + f"got {type(buffers).__name__}" + ) + + +IF CUDA_CORE_BUILD_MAJOR >= 13: + cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc): + """Convert a _LocSpec dataclass to a cydriver.CUmemLocation struct.""" + 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: + cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc): + raise NotImplementedError( + "_to_cumemlocation requires a CUDA 13 build of cuda.core" + ) + + +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: object, + srcs: Sequence[Buffer], + dsts: Sequence[Buffer], + *, + options: object = None, +) -> None: + """Copy a batch of buffers asynchronously. + + Requires CUDA 13+. 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. Passing a + :class:`~graph.GraphBuilder` raises ``CUDAError`` + (``CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED``). + 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 + ------ + NotImplementedError + On a CUDA 12 build of ``cuda.core``. + ValueError + If lengths or sizes mismatch. + TypeError + If a single Buffer is passed instead of a sequence. + UserWarning + If ``overlap_mode='prefer_overlap_with_compute'`` is requested + on a non-integrated (discrete) GPU. + """ + cdef tuple src_bufs = _coerce_batch_buffers(srcs, "copy_batch") + cdef tuple dst_bufs = _coerce_batch_buffers(dsts, "copy_batch") + 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})" + ) + + # Expand `options` to one CopyOptions per copy; the encoder below + # collapses equal neighbours back into driver attribute runs. + cdef tuple attr_tuple + if options is None: + attr_tuple = (CopyOptions(),) * n + elif isinstance(options, CopyOptions): + attr_tuple = (options,) * n + elif isinstance(options, Sequence): + if len(options) != n: + raise ValueError( + f"copy_batch: options length {len(options)} does not match " + f"buffers length {n}" + ) + attr_list = [] + 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__}" + ) + attr_list.append(a) + attr_tuple = tuple(attr_list) + else: + raise TypeError( + f"copy_batch: options must be CopyOptions or a sequence of " + f"CopyOptions, got {type(options).__name__}" + ) + + # 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: + 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, + )) + ELSE: + raise NotImplementedError( + "copy_batch requires a CUDA 13 build of cuda.core" + ) 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/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/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..4510c8c83c6 --- /dev/null +++ b/cuda_core/tests/helpers/copy_batch.py @@ -0,0 +1,60 @@ +# 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. +""" + +import pytest + +from cuda.bindings import driver +from cuda.core import LegacyPinnedMemoryResource, ManagedMemoryResource +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 skip_if_copy_batch_unsupported(): + if not hasattr(driver, "cuMemcpyBatchAsync"): + pytest.skip("cuMemcpyBatchAsync unavailable (CUDA 13+ required)") + + +def managed_mr_or_skip(): + try: + return ManagedMemoryResource() + except RuntimeError as exc: + if "requires CUDA 13.0" in str(exc) or "managed allocations" in str(exc): + pytest.skip("ManagedMemoryResource not available") + raise + + +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..27422b3cb7e --- /dev/null +++ b/cuda_core/tests/memory/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/cuda_core/tests/memory/conftest.py b/cuda_core/tests/memory/conftest.py new file mode 100644 index 00000000000..b23bc99d894 --- /dev/null +++ b/cuda_core/tests/memory/conftest.py @@ -0,0 +1,71 @@ +# 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). Constants and helper functions that tests +import by name live in ``helpers.copy_batch``. +""" + +import pytest +from helpers.copy_batch import ( + COPY_BATCH_COUNT, + COPY_BATCH_SIZE, + skip_if_copy_batch_unsupported, +) + +from cuda.core import Device, LegacyPinnedMemoryResource + + +@pytest.fixture +def copy_batch_device(init_cuda): + skip_if_copy_batch_unsupported() + device = Device() + device.set_current() + return device + + +@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..b1b777792bf --- /dev/null +++ b/cuda_core/tests/memory/test_copy_batch.py @@ -0,0 +1,288 @@ +# 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. Options handling and argument +validation live in ``test_copy_batch_options.py``. +""" + +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.cuda_utils import CUDAError +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_visible_on_other_stream_after_explicit_wait(self, copy_batch_device, device_bufs, copy_stream): + """A batch on one stream is not ordered against an unrelated stream. + + The second stream must be made to wait explicitly; once it does, + it observes the copied bytes. + """ + srcs, dsts = device_bufs + other = copy_batch_device.create_stream() + try: + for src in srcs: + src.fill(33, stream=copy_stream) + copy_batch(copy_stream, srcs, dsts) + + # Explicit cross-stream dependency, then observe from `other`. + other.wait(copy_stream) + probes = [copy_batch_device.memory_resource.allocate(COPY_BATCH_SIZE, stream=other) for _ in dsts] + for dst, probe in zip(dsts, probes): + dst.copy_to(probe, stream=other) + other.sync() + + for probe in probes: + assert compare_buffer_to_constant(probe, 33) + for probe in probes: + probe.close(other) + other.sync() + finally: + other.close() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_graph_builder_capture_is_unsupported(self, copy_batch_device, device_bufs, copy_stream): + """Batched memcpy cannot be captured into a graph. + + ``Stream_accept`` takes a ``GraphBuilder``, so the batch reaches + the driver, but ``cuMemcpyBatchAsync`` has no graph-node form and + the driver rejects it mid-capture. Pinned here so the limitation + is asserted rather than rediscovered; 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(CUDAError, match="STREAM_CAPTURE_UNSUPPORTED"): + copy_batch(gb, srcs, dsts) + finally: + # The rejection leaves the capture intact, so it 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..c3f0c8b5df6 --- /dev/null +++ b/cuda_core/tests/memory/test_copy_batch_options.py @@ -0,0 +1,382 @@ +# 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. The data-movement +behaviour itself lives in ``test_copy_batch.py``. +""" + +import warnings + +import pytest +from helpers.buffers import compare_buffer_to_constant, compare_equal_buffers, set_buffer +from helpers.copy_batch import ( + COPY_BATCH_SIZE, + OVERLAP_WARNING_FILTER, + assert_managed_holds, + managed_mr_or_skip, +) + +from cuda.core import Device, Host, LegacyPinnedMemoryResource +from cuda.core._memory._copy_enums import _attr_run_starts +from cuda.core._utils.version import binding_version +from cuda.core.utils import ( + CopyOptions, + MemcpyOverlapMode, + MemcpySrcAccessOrder, + copy_batch, +) + + +class TestAttrRunStarts: + """Unit tests for the attrsIdxs run-length encoding. + + Pure logic, no CUDA: ``attrs[k]`` applies to the copies in + ``[starts[k], starts[k + 1])``. + """ + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_broadcast_collapses_to_one_run(self): + attrs = [CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY)] * 4 + assert _attr_run_starts(attrs) == [0] + + @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, 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, 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), + ] + assert _attr_run_starts(per_copy_options) == [0, 1, 2, 3] + + 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_scalar_options_broadcast(self, copy_batch_device, h2d_bufs, copy_stream): + """A scalar option must apply to every copy. + + Verified three ways: the scalar collapses to a single driver + attribute, a scalar and an equivalent explicit per-copy list give + identical bytes, and a short list is *not* silently broadcast. + """ + srcs, scalar_dsts = h2d_bufs + device_mr = copy_batch_device.memory_resource + pinned_mr = LegacyPinnedMemoryResource() + n = len(srcs) + scalar_option = CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY) + + for i, src in enumerate(srcs): + set_buffer(src, i + 95) + + # A scalar is expanded internally to n copies of one option, which + # the encoder then collapses back to a single driver entry. + assert _attr_run_starts([scalar_option] * n) == [0] + + copy_batch(copy_stream, srcs, scalar_dsts, options=scalar_option) + + listed_dsts = [device_mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in srcs] + copy_batch(copy_stream, srcs, listed_dsts, options=[scalar_option] * n) + copy_stream.sync() + + # Every copy received the option, and both spellings agree. + for i, (scalar_dst, listed_dst) in enumerate(zip(scalar_dsts, listed_dsts)): + assert compare_buffer_to_constant(scalar_dst, i + 95) + scalar_host = pinned_mr.allocate(COPY_BATCH_SIZE) + listed_host = pinned_mr.allocate(COPY_BATCH_SIZE) + scalar_dst.copy_to(scalar_host, stream=copy_stream) + listed_dst.copy_to(listed_host, stream=copy_stream) + copy_stream.sync() + assert compare_equal_buffers(scalar_host, listed_host) + scalar_host.close(copy_stream) + listed_host.close(copy_stream) + + # A sequence is paired by index and never broadcast, so a + # one-element list is a length error rather than a scalar. + with pytest.raises(ValueError, match="options length"): + copy_batch(copy_stream, srcs, listed_dsts, options=[scalar_option]) + + for buf in listed_dsts: + buf.close(copy_stream) + copy_stream.sync() + + @pytest.mark.agent_authored(model="Claude Opus 5") + def test_location_hints(self, copy_batch_device, copy_stream): + dev = copy_batch_device + mr = managed_mr_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, copy_batch_device, copy_stream): + """NUMA host hints round-trip on CUDA 13 and are rejected on CUDA 12.""" + dev = copy_batch_device + mr = managed_mr_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) + + options = CopyOptions(dst_location_hint=Host(numa_id=0)) + + if binding_version() < (13, 0, 0): + with pytest.raises(TypeError, match="CUDA 13"): + copy_batch(copy_stream, srcs, dsts, options=options) + else: + copy_batch(copy_stream, srcs, dsts, options=options) + 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, 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, 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") + @pytest.mark.filterwarnings(OVERLAP_WARNING_FILTER) + 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) + + +@pytest.mark.agent_authored(model="Claude Opus 5") +def test_cuda12_raises_not_implemented(init_cuda): + """cuMemcpyBatchAsync is CUDA 13+; single copies use Buffer.copy_to.""" + if binding_version() >= (13, 0, 0): + pytest.skip("Only relevant on CUDA 12 builds") + + 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="CUDA 13"): + copy_batch(stream, [src], [dst]) + + src.close(stream) + dst.close(stream) + stream.sync() + stream.close() From fd58e8fd8ca9cfd5cbd312840dc204d40fc75eb5 Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Mon, 10 Aug 2026 13:14:36 -0700 Subject: [PATCH 2/6] fallback for CUDA 12 and type annotations --- cuda_core/cuda/core/_memory/_buffer.pxd | 6 + cuda_core/cuda/core/_memory/_buffer.pyx | 27 ++++ cuda_core/cuda/core/_memory/_copy_enums.py | 59 ++++--- cuda_core/cuda/core/_memory/_copy_ops.pyi | 43 +++-- cuda_core/cuda/core/_memory/_copy_ops.pyx | 153 ++++++++++++------ .../cuda/core/_memory/_managed_memory_ops.pyi | 1 + .../cuda/core/_memory/_managed_memory_ops.pyx | 27 +--- cuda_core/docs/source/release/1.2.0-notes.rst | 19 +++ cuda_core/tests/helpers/copy_batch.py | 25 ++- cuda_core/tests/memory/conftest.py | 20 ++- cuda_core/tests/memory/test_copy_batch.py | 18 +-- .../tests/memory/test_copy_batch_options.py | 99 ++++++++---- 12 files changed, 341 insertions(+), 156 deletions(-) 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 index 823c0e3c965..b7503febe68 100644 --- a/cuda_core/cuda/core/_memory/_copy_enums.py +++ b/cuda_core/cuda/core/_memory/_copy_enums.py @@ -5,12 +5,12 @@ from __future__ import annotations import dataclasses -import functools 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 @@ -94,39 +94,46 @@ def __post_init__(self): def _to_driver_enum(self) -> int: """Return the driver CUmemcpySrcAccessOrder value.""" - return _src_access_order_to_cu()[MemcpySrcAccessOrder(self.src_access_order)] + 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.""" - return _overlap_mode_to_cu()[MemcpyOverlapMode(self.overlap_mode)] + if not _OVERLAP_MODE_TO_DRIVER: + raise NotImplementedError(_CUDA13_REQUIRED) + return _OVERLAP_MODE_TO_DRIVER[MemcpyOverlapMode(self.overlap_mode)] -# Bridges between the public StrEnums and the driver integer values. Built on -# first use rather than at import: the CUmemcpy* enums only exist on toolkits -# that ship the batched memcpy entry points, and importing cuda.core must not -# depend on them. +_CUDA13_REQUIRED = "copy attributes require a CUDA 13 build of cuda-bindings" + +# CUmemcpySrcAccessOrder and CUmemcpyFlags are CUDA 13 additions, so these +# maps are empty on a CUDA 12 build. Nothing reaches them there: copy_batch +# refuses non-default CopyOptions when the batched entry point is absent. # -# Keyed by ``str`` rather than by the enum: under ``python_version = "3.10"`` -# mypy resolves ``StrEnum`` to the unstubbed ``backports.strenum`` shim and so -# infers the members as plain ``str``. StrEnum members are ``str`` instances, -# so this annotation is accurate on every supported version. -@functools.cache -def _src_access_order_to_cu() -> dict[str, int]: - cu = driver.CUmemcpySrcAccessOrder - return { - MemcpySrcAccessOrder.STREAM: int(cu.CU_MEMCPY_SRC_ACCESS_ORDER_STREAM), - MemcpySrcAccessOrder.DURING_API_CALL: int(cu.CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL), - MemcpySrcAccessOrder.ANY: int(cu.CU_MEMCPY_SRC_ACCESS_ORDER_ANY), +# 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), } - - -@functools.cache -def _overlap_mode_to_cu() -> dict[str, int]: - cu = driver.CUmemcpyFlags - return { - MemcpyOverlapMode.DEFAULT: int(cu.CU_MEMCPY_FLAG_DEFAULT), - MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE: int(cu.CU_MEMCPY_FLAG_PREFER_OVERLAP_WITH_COMPUTE), + _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]: diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyi b/cuda_core/cuda/core/_memory/_copy_ops.pyi index a0ce96ce4dd..218ba0b7e30 100644 --- a/cuda_core/cuda/core/_memory/_copy_ops.pyi +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyi @@ -5,13 +5,21 @@ 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 copy_batch(stream: object, srcs: Sequence[Buffer], dsts: Sequence[Buffer], *, options: object=None) -> None: +def _batch_entry_point_in_use() -> bool: + """Internal: expose the dispatch predicate so tests can gate on it.""" + +def copy_batch(stream: Stream, srcs: Sequence[Buffer], dsts: Sequence[Buffer], *, options: CopyOptions | Sequence[CopyOptions] | None=None) -> None: """Copy a batch of buffers asynchronously. - Requires CUDA 13+. For a single buffer, use - :meth:`Buffer.copy_to` or :meth:`Buffer.copy_from`. + 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 @@ -20,13 +28,14 @@ def copy_batch(stream: object, srcs: Sequence[Buffer], dsts: Sequence[Buffer], * Parameters ---------- stream : :class:`~_stream.Stream` - Stream for the asynchronous copy. Passing a - :class:`~graph.GraphBuilder` raises ``CUDAError`` - (``CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED``). + 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. + Source buffers. Must be a sequence, not a single Buffer. dsts : Sequence[:class:`Buffer`] - Destination buffers. Must match ``len(srcs)``. + 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`` @@ -34,12 +43,26 @@ def copy_batch(stream: object, srcs: Sequence[Buffer], dsts: Sequence[Buffer], * Raises ------ - NotImplementedError - On a CUDA 12 build of ``cuda.core``. 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 + ----- + ``cuMemcpyBatchAsync`` needs both a CUDA 13 build of ``cuda.core`` + and a CUDA 13 driver. Otherwise the copies fall back to a + Python-level loop over ``cuMemcpyAsync``, which is semantically + equivalent but does not amortize launch overhead. That 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. diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyx b/cuda_core/cuda/core/_memory/_copy_ops.pyx index fc2d7477506..e95733e3ca3 100644 --- a/cuda_core/cuda/core/_memory/_copy_ops.pyx +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyx @@ -13,42 +13,44 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: from libc.string cimport memset -# as_cu, HANDLE_RETURN and _attr_run_starts are referenced only from the -# CUDA 13 branch of _do_copy_batch. cython-lint does not evaluate -# compile-time IF blocks, so it needs a pragma to see them as used. from cuda.bindings cimport cydriver -from cuda.core._memory._buffer cimport Buffer -from cuda.core._resource_handles cimport as_cu # no-cython-lint +from cuda.core._memory._buffer cimport Buffer, Buffer_coerce_batch +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 # no-cython-lint +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 -cdef tuple _coerce_batch_buffers(object buffers, str what): - """Coerce ``buffers`` to a tuple[Buffer, ...]; rejects a single Buffer.""" - cdef list out - if isinstance(buffers, Buffer): - raise TypeError( - f"{what}: pass a sequence of Buffers; for a single buffer use " - f"the Buffer.copy_to / Buffer.copy_from instance method" - ) - if isinstance(buffers, Sequence): - if not buffers: - raise ValueError(f"{what}: empty buffers sequence") - out = [] - for t in buffers: - if not isinstance(t, Buffer): - raise TypeError( - f"{what}: expected Buffer, got {type(t).__name__}" - ) - out.append(t) - return tuple(out) - raise TypeError( - f"{what}: buffers must be a sequence of Buffer, " - f"got {type(buffers).__name__}" - ) +_SINGLE_COPY_HINT = "Buffer.copy_to / Buffer.copy_from" + +# Attributes reach the driver only through cuMemcpyBatchAsync, which is +# CUDA 13+. 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. + + Needs a CUDA 13 build (compile time) and a CUDA 13 driver (run time); + a CUDA 13 build against a CUDA 12 driver has no such symbol. + """ + 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()) IF CUDA_CORE_BUILD_MAJOR >= 13: @@ -97,16 +99,17 @@ cdef cydriver.CUmemcpyAttributes _to_cu_memcpy_attributes(object attr): def copy_batch( - stream: object, + stream: Stream, srcs: Sequence[Buffer], dsts: Sequence[Buffer], *, - options: object = None, + options: CopyOptions | Sequence[CopyOptions] | None = None, ) -> None: """Copy a batch of buffers asynchronously. - Requires CUDA 13+. For a single buffer, use - :meth:`Buffer.copy_to` or :meth:`Buffer.copy_from`. + 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 @@ -115,13 +118,14 @@ def copy_batch( Parameters ---------- stream : :class:`~_stream.Stream` - Stream for the asynchronous copy. Passing a - :class:`~graph.GraphBuilder` raises ``CUDAError`` - (``CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED``). + 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. + Source buffers. Must be a sequence, not a single Buffer. dsts : Sequence[:class:`Buffer`] - Destination buffers. Must match ``len(srcs)``. + 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`` @@ -129,18 +133,32 @@ def copy_batch( Raises ------ - NotImplementedError - On a CUDA 12 build of ``cuda.core``. 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 + ----- + ``cuMemcpyBatchAsync`` needs both a CUDA 13 build of ``cuda.core`` + and a CUDA 13 driver. Otherwise the copies fall back to a + Python-level loop over ``cuMemcpyAsync``, which is semantically + equivalent but does not amortize launch overhead. That 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 = _coerce_batch_buffers(srcs, "copy_batch") - cdef tuple dst_bufs = _coerce_batch_buffers(dsts, "copy_batch") + 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: @@ -191,6 +209,18 @@ def copy_batch( f"CopyOptions, got {type(options).__name__}" ) + # 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 both a CUDA 13 build of cuda.core and a CUDA 13 " + "driver; 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 @@ -216,6 +246,41 @@ def copy_batch( cdef void _do_copy_batch(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tuple): IF CUDA_CORE_BUILD_MAJOR >= 13: + # A CUDA 13 build can still be running against a CUDA 12 driver, + # which has no cuMemcpyBatchAsync (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 @@ -260,7 +325,3 @@ cdef void _do_copy_batch(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tu num_attrs, hstream, )) - ELSE: - raise NotImplementedError( - "copy_batch requires a CUDA 13 build of cuda.core" - ) 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..73175af87b1 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx @@ -11,7 +11,7 @@ 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 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 +52,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): 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..dc5a0d99b20 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,25 @@ ``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 both a CUDA 13 build of ``cuda.core`` and a CUDA 13 + driver; otherwise the copies fall back to a per-copy 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/tests/helpers/copy_batch.py b/cuda_core/tests/helpers/copy_batch.py index 4510c8c83c6..337a73438ad 100644 --- a/cuda_core/tests/helpers/copy_batch.py +++ b/cuda_core/tests/helpers/copy_batch.py @@ -8,10 +8,8 @@ pieces that tests import by name. """ -import pytest - -from cuda.bindings import driver -from cuda.core import LegacyPinnedMemoryResource, ManagedMemoryResource +from cuda.core import LegacyPinnedMemoryResource +from cuda.core._memory._copy_ops import _batch_entry_point_in_use from helpers.buffers import compare_equal_buffers, make_scratch_buffer COPY_BATCH_SIZE = 4096 @@ -23,18 +21,15 @@ OVERLAP_WARNING_FILTER = "ignore:overlap_mode:UserWarning" -def skip_if_copy_batch_unsupported(): - if not hasattr(driver, "cuMemcpyBatchAsync"): - pytest.skip("cuMemcpyBatchAsync unavailable (CUDA 13+ required)") - +def uses_batch_entry_point() -> bool: + """Whether copy_batch reaches ``cuMemcpyBatchAsync`` on this system. -def managed_mr_or_skip(): - try: - return ManagedMemoryResource() - except RuntimeError as exc: - if "requires CUDA 13.0" in str(exc) or "managed allocations" in str(exc): - pytest.skip("ManagedMemoryResource not available") - raise + Delegates to the implementation's own dispatch predicate rather than + re-deriving it, so the tests cannot drift from it. ``copy_batch`` + itself works either way -- only non-default ``CopyOptions`` need the + batched entry point. + """ + return _batch_entry_point_in_use() def assert_managed_holds(dev, buf, value, *, stream): diff --git a/cuda_core/tests/memory/conftest.py b/cuda_core/tests/memory/conftest.py index b23bc99d894..6550cab753d 100644 --- a/cuda_core/tests/memory/conftest.py +++ b/cuda_core/tests/memory/conftest.py @@ -13,7 +13,7 @@ from helpers.copy_batch import ( COPY_BATCH_COUNT, COPY_BATCH_SIZE, - skip_if_copy_batch_unsupported, + uses_batch_entry_point, ) from cuda.core import Device, LegacyPinnedMemoryResource @@ -21,12 +21,28 @@ @pytest.fixture def copy_batch_device(init_cuda): - skip_if_copy_batch_unsupported() + """``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. + + The per-copy ``cuMemcpyAsync`` fallback used on a CUDA 12 build, or on + a CUDA 13 build against a CUDA 12 driver, has no way to convey copy + attributes, so ``copy_batch`` rejects non-default options there. + """ + if not uses_batch_entry_point(): + pytest.skip("non-default CopyOptions requires a CUDA 13 build and a CUDA 13 driver") + + @pytest.fixture def copy_stream(copy_batch_device): """The single stream used for both allocation and copies in a test. diff --git a/cuda_core/tests/memory/test_copy_batch.py b/cuda_core/tests/memory/test_copy_batch.py index b1b777792bf..eb6b89dd045 100644 --- a/cuda_core/tests/memory/test_copy_batch.py +++ b/cuda_core/tests/memory/test_copy_batch.py @@ -19,7 +19,6 @@ from helpers.copy_batch import COPY_BATCH_SIZE from cuda.core import LegacyPinnedMemoryResource -from cuda.core._utils.cuda_utils import CUDAError from cuda.core.utils import copy_batch @@ -268,21 +267,22 @@ def test_visible_on_other_stream_after_explicit_wait(self, copy_batch_device, de other.close() @pytest.mark.agent_authored(model="Claude Opus 5") - def test_graph_builder_capture_is_unsupported(self, copy_batch_device, device_bufs, copy_stream): + def test_graph_builder_is_rejected(self, copy_batch_device, device_bufs, copy_stream): """Batched memcpy cannot be captured into a graph. - ``Stream_accept`` takes a ``GraphBuilder``, so the batch reaches - the driver, but ``cuMemcpyBatchAsync`` has no graph-node form and - the driver rejects it mid-capture. Pinned here so the limitation - is asserted rather than rediscovered; use ``GraphNode.memcpy`` or - per-buffer ``Buffer.copy_to`` to build copies 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(CUDAError, match="STREAM_CAPTURE_UNSUPPORTED"): + with pytest.raises(TypeError, match="Argument 'stream' has incorrect type"): copy_batch(gb, srcs, dsts) finally: - # The rejection leaves the capture intact, so it still ends cleanly. + # 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 index c3f0c8b5df6..b4d11b4040b 100644 --- a/cuda_core/tests/memory/test_copy_batch_options.py +++ b/cuda_core/tests/memory/test_copy_batch_options.py @@ -11,12 +11,17 @@ 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, compare_equal_buffers, set_buffer from helpers.copy_batch import ( COPY_BATCH_SIZE, OVERLAP_WARNING_FILTER, assert_managed_holds, - managed_mr_or_skip, + uses_batch_entry_point, ) from cuda.core import Device, Host, LegacyPinnedMemoryResource @@ -83,7 +88,7 @@ class TestCopyBatchOptions: (MemcpySrcAccessOrder.ANY, 33), ], ) - def test_src_access_order(self, h2d_bufs, copy_stream, order, marker): + 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) @@ -95,7 +100,7 @@ def test_src_access_order(self, h2d_bufs, copy_stream, order, marker): assert compare_buffer_to_constant(dst, i + marker) @pytest.mark.agent_authored(model="Claude Opus 5") - def test_per_copy_options(self, h2d_bufs, copy_stream): + 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) @@ -115,7 +120,7 @@ def test_per_copy_options(self, h2d_bufs, copy_stream): assert compare_buffer_to_constant(dst, i + 40) @pytest.mark.agent_authored(model="Claude Opus 5") - def test_scalar_options_broadcast(self, copy_batch_device, h2d_bufs, copy_stream): + def test_scalar_options_broadcast(self, requires_copy_options, copy_batch_device, h2d_bufs, copy_stream): """A scalar option must apply to every copy. Verified three ways: the scalar collapses to a single driver @@ -163,9 +168,9 @@ def test_scalar_options_broadcast(self, copy_batch_device, h2d_bufs, copy_stream copy_stream.sync() @pytest.mark.agent_authored(model="Claude Opus 5") - def test_location_hints(self, copy_batch_device, copy_stream): + def test_location_hints(self, requires_copy_options, copy_batch_device, copy_stream): dev = copy_batch_device - mr = managed_mr_or_skip() + 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)] @@ -189,10 +194,10 @@ def test_location_hints(self, copy_batch_device, copy_stream): mr.close() @pytest.mark.agent_authored(model="Claude Opus 5") - def test_host_numa_location_hint(self, copy_batch_device, copy_stream): + def test_host_numa_location_hint(self, requires_copy_options, copy_batch_device, copy_stream): """NUMA host hints round-trip on CUDA 13 and are rejected on CUDA 12.""" dev = copy_batch_device - mr = managed_mr_or_skip() + 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): @@ -216,7 +221,7 @@ def test_host_numa_location_hint(self, copy_batch_device, copy_stream): @pytest.mark.agent_authored(model="Claude Opus 5") @pytest.mark.filterwarnings(OVERLAP_WARNING_FILTER) - def test_overlap_mode_copies_correctly(self, h2d_bufs, copy_stream): + 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): @@ -234,7 +239,9 @@ def test_overlap_mode_copies_correctly(self, h2d_bufs, copy_stream): 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, copy_batch_device, h2d_bufs, copy_stream): + 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) @@ -360,23 +367,61 @@ def test_rejects_bad_options_element(self, h2d_bufs, copy_stream): copy_batch(copy_stream, srcs, dsts, options=bad) -@pytest.mark.agent_authored(model="Claude Opus 5") -def test_cuda12_raises_not_implemented(init_cuda): - """cuMemcpyBatchAsync is CUDA 13+; single copies use Buffer.copy_to.""" - if binding_version() >= (13, 0, 0): - pytest.skip("Only relevant on CUDA 12 builds") +class TestPerCopyFallback: + """Behaviour where ``cuMemcpyBatchAsync`` is unavailable. + + Reached on a CUDA 12 build of ``cuda.core`` and on a CUDA 13 build + running against a CUDA 12 driver. Skipped when the batched entry + point is actually in use. + """ + + @pytest.fixture(autouse=True) + def _skip_if_batched(self): + if uses_batch_entry_point(): + pytest.skip("cuMemcpyBatchAsync is in use; fallback path not exercised") - 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) + @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) - with pytest.raises(NotImplementedError, match="CUDA 13"): - copy_batch(stream, [src], [dst]) + copy_batch(stream, srcs, dsts) + stream.sync() - src.close(stream) - dst.close(stream) - stream.sync() - stream.close() + 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() From 24525d2064da6e53ac4ac27a873a65fa4664f981 Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Mon, 10 Aug 2026 14:40:47 -0700 Subject: [PATCH 3/6] be more precise about CUDA requirements --- cuda_core/cuda/core/_memory/_copy_enums.py | 20 +++++--- cuda_core/cuda/core/_memory/_copy_ops.pyi | 19 ++++--- cuda_core/cuda/core/_memory/_copy_ops.pyx | 50 ++++++++++++------- cuda_core/docs/source/release/1.2.0-notes.rst | 8 ++- cuda_core/tests/helpers/copy_batch.py | 9 ++-- cuda_core/tests/memory/__init__.py | 15 ++++++ cuda_core/tests/memory/conftest.py | 11 ++-- .../tests/memory/test_copy_batch_options.py | 7 +-- 8 files changed, 96 insertions(+), 43 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_copy_enums.py b/cuda_core/cuda/core/_memory/_copy_enums.py index b7503febe68..20bfc469bf7 100644 --- a/cuda_core/cuda/core/_memory/_copy_enums.py +++ b/cuda_core/cuda/core/_memory/_copy_enums.py @@ -71,8 +71,15 @@ class CopyOptions: overlap_mode: MemcpyOverlapMode | str = "default" def __post_init__(self): - # Validate enum fields while still in __init__ (frozen dataclass). - # Use __setattr__ because fields are frozen. + # 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__( @@ -105,11 +112,12 @@ def _to_driver_flags(self) -> int: return _OVERLAP_MODE_TO_DRIVER[MemcpyOverlapMode(self.overlap_mode)] -_CUDA13_REQUIRED = "copy attributes require a CUDA 13 build of cuda-bindings" +_CUDA13_REQUIRED = "copy attributes require cuda.bindings 13.0 or newer" -# CUmemcpySrcAccessOrder and CUmemcpyFlags are CUDA 13 additions, so these -# maps are empty on a CUDA 12 build. Nothing reaches them there: copy_batch -# refuses non-default CopyOptions when the batched entry point is absent. +# 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``. diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyi b/cuda_core/cuda/core/_memory/_copy_ops.pyi index 218ba0b7e30..6c33af3b19d 100644 --- a/cuda_core/cuda/core/_memory/_copy_ops.pyi +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyi @@ -53,13 +53,18 @@ def copy_batch(stream: Stream, srcs: Sequence[Buffer], dsts: Sequence[Buffer], * Notes ----- - ``cuMemcpyBatchAsync`` needs both a CUDA 13 build of ``cuda.core`` - and a CUDA 13 driver. Otherwise the copies fall back to a - Python-level loop over ``cuMemcpyAsync``, which is semantically - equivalent but does not amortize launch overhead. That fallback has - no way to convey :class:`CopyOptions` to the driver, so non-default - options raise :class:`NotImplementedError` there rather than being - silently ignored. + 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 ----- diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyx b/cuda_core/cuda/core/_memory/_copy_ops.pyx index e95733e3ca3..41bebefb936 100644 --- a/cuda_core/cuda/core/_memory/_copy_ops.pyx +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyx @@ -30,17 +30,26 @@ 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, which is -# CUDA 13+. The per-copy cuMemcpyAsync fallback has nowhere to put them, -# so anything other than the defaults is refused rather than dropped. +# 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. - Needs a CUDA 13 build (compile time) and a CUDA 13 driver (run time); - a CUDA 13 build against a CUDA 12 driver has no such symbol. + 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) @@ -76,7 +85,7 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: ELSE: cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc): raise NotImplementedError( - "_to_cumemlocation requires a CUDA 13 build of cuda.core" + "_to_cumemlocation requires cuda.core built against CUDA 13 headers" ) @@ -143,13 +152,18 @@ def copy_batch( Notes ----- - ``cuMemcpyBatchAsync`` needs both a CUDA 13 build of ``cuda.core`` - and a CUDA 13 driver. Otherwise the copies fall back to a - Python-level loop over ``cuMemcpyAsync``, which is semantically - equivalent but does not amortize launch overhead. That fallback has - no way to convey :class:`CopyOptions` to the driver, so non-default - options raise :class:`NotImplementedError` there rather than being - silently ignored. + 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 ----- @@ -217,8 +231,9 @@ def copy_batch( if attr_tuple[i] != _DEFAULT_COPY_OPTIONS: raise NotImplementedError( "copy_batch: non-default CopyOptions requires cuMemcpyBatchAsync, " - "which needs both a CUDA 13 build of cuda.core and a CUDA 13 " - "driver; omit options to use the per-copy fallback" + "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 @@ -246,8 +261,9 @@ def copy_batch( cdef void _do_copy_batch(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tuple): IF CUDA_CORE_BUILD_MAJOR >= 13: - # A CUDA 13 build can still be running against a CUDA 12 driver, - # which has no cuMemcpyBatchAsync (see PRs #2054 / #2064). + # 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: 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 dc5a0d99b20..b41d507699b 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -18,8 +18,12 @@ New features ``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 both a CUDA 13 build of ``cuda.core`` and a CUDA 13 - driver; otherwise the copies fall back to a per-copy loop, and non-default + 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. diff --git a/cuda_core/tests/helpers/copy_batch.py b/cuda_core/tests/helpers/copy_batch.py index 337a73438ad..ceed4d1dc34 100644 --- a/cuda_core/tests/helpers/copy_batch.py +++ b/cuda_core/tests/helpers/copy_batch.py @@ -24,10 +24,11 @@ def uses_batch_entry_point() -> bool: """Whether copy_batch reaches ``cuMemcpyBatchAsync`` on this system. - Delegates to the implementation's own dispatch predicate rather than - re-deriving it, so the tests cannot drift from it. ``copy_batch`` - itself works either way -- only non-default ``CopyOptions`` need the - batched entry point. + True only with cuda.core built against CUDA 13, cuda.bindings 13.0+, + and a driver reporting CUDA 13.0 or newer. Delegates to the + implementation's own dispatch predicate rather than re-deriving it, so + the tests cannot drift from it. ``copy_batch`` itself works either way + -- only non-default ``CopyOptions`` need the batched entry point. """ return _batch_entry_point_in_use() diff --git a/cuda_core/tests/memory/__init__.py b/cuda_core/tests/memory/__init__.py index 27422b3cb7e..61904faa870 100644 --- a/cuda_core/tests/memory/__init__.py +++ b/cuda_core/tests/memory/__init__.py @@ -1,3 +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 index 6550cab753d..aa83720fbc1 100644 --- a/cuda_core/tests/memory/conftest.py +++ b/cuda_core/tests/memory/conftest.py @@ -35,12 +35,15 @@ def copy_batch_device(init_cuda): def requires_copy_options(): """Skip when ``CopyOptions`` cannot reach the driver. - The per-copy ``cuMemcpyAsync`` fallback used on a CUDA 12 build, or on - a CUDA 13 build against a CUDA 12 driver, has no way to convey copy - attributes, so ``copy_batch`` rejects non-default options there. + 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 uses_batch_entry_point(): - pytest.skip("non-default CopyOptions requires a CUDA 13 build and a CUDA 13 driver") + 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 diff --git a/cuda_core/tests/memory/test_copy_batch_options.py b/cuda_core/tests/memory/test_copy_batch_options.py index b4d11b4040b..41d718612ed 100644 --- a/cuda_core/tests/memory/test_copy_batch_options.py +++ b/cuda_core/tests/memory/test_copy_batch_options.py @@ -370,9 +370,10 @@ def test_rejects_bad_options_element(self, h2d_bufs, copy_stream): class TestPerCopyFallback: """Behaviour where ``cuMemcpyBatchAsync`` is unavailable. - Reached on a CUDA 12 build of ``cuda.core`` and on a CUDA 13 build - running against a CUDA 12 driver. Skipped when the batched entry - point is actually in use. + 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) From 8e75a9e8dd4c9f669fcc712902a5bad7aabc633a Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Mon, 10 Aug 2026 14:48:28 -0700 Subject: [PATCH 4/6] skip tests on Windows that require managed memory --- cuda_core/tests/conftest.py | 2 ++ 1 file changed, 2 insertions(+) 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 From b9ef7ac2b8bcb45be98cc6cd509b46991f44c0d4 Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Mon, 10 Aug 2026 16:48:06 -0700 Subject: [PATCH 5/6] rework some tests --- cuda_core/cuda/core/_memory/_copy_ops.pyi | 12 ++ cuda_core/cuda/core/_memory/_copy_ops.pyx | 65 ++++++---- cuda_core/tests/helpers/copy_batch.py | 13 -- cuda_core/tests/memory/conftest.py | 12 +- cuda_core/tests/memory/test_copy_batch.py | 32 +---- .../tests/memory/test_copy_batch_options.py | 120 +++++++----------- 6 files changed, 100 insertions(+), 154 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyi b/cuda_core/cuda/core/_memory/_copy_ops.pyi index 6c33af3b19d..03def659ffc 100644 --- a/cuda_core/cuda/core/_memory/_copy_ops.pyi +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyi @@ -14,6 +14,18 @@ _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. diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyx b/cuda_core/cuda/core/_memory/_copy_ops.pyx index 41bebefb936..11964b99f8e 100644 --- a/cuda_core/cuda/core/_memory/_copy_ops.pyx +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyx @@ -62,6 +62,43 @@ def _batch_entry_point_in_use() -> bool: 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__}" + ) + + IF CUDA_CORE_BUILD_MAJOR >= 13: cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc): """Convert a _LocSpec dataclass to a cydriver.CUmemLocation struct.""" @@ -195,33 +232,7 @@ def copy_batch( f"(src={src_buf.size}, dst={dst_buf.size})" ) - # Expand `options` to one CopyOptions per copy; the encoder below - # collapses equal neighbours back into driver attribute runs. - cdef tuple attr_tuple - if options is None: - attr_tuple = (CopyOptions(),) * n - elif isinstance(options, CopyOptions): - attr_tuple = (options,) * n - elif isinstance(options, Sequence): - if len(options) != n: - raise ValueError( - f"copy_batch: options length {len(options)} does not match " - f"buffers length {n}" - ) - attr_list = [] - 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__}" - ) - attr_list.append(a) - attr_tuple = tuple(attr_list) - else: - raise TypeError( - f"copy_batch: options must be CopyOptions or a sequence of " - f"CopyOptions, got {type(options).__name__}" - ) + 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 diff --git a/cuda_core/tests/helpers/copy_batch.py b/cuda_core/tests/helpers/copy_batch.py index ceed4d1dc34..0eafff2094e 100644 --- a/cuda_core/tests/helpers/copy_batch.py +++ b/cuda_core/tests/helpers/copy_batch.py @@ -9,7 +9,6 @@ """ from cuda.core import LegacyPinnedMemoryResource -from cuda.core._memory._copy_ops import _batch_entry_point_in_use from helpers.buffers import compare_equal_buffers, make_scratch_buffer COPY_BATCH_SIZE = 4096 @@ -21,18 +20,6 @@ OVERLAP_WARNING_FILTER = "ignore:overlap_mode:UserWarning" -def uses_batch_entry_point() -> bool: - """Whether copy_batch reaches ``cuMemcpyBatchAsync`` on this system. - - True only with cuda.core built against CUDA 13, cuda.bindings 13.0+, - and a driver reporting CUDA 13.0 or newer. Delegates to the - implementation's own dispatch predicate rather than re-deriving it, so - the tests cannot drift from it. ``copy_batch`` itself works either way - -- only non-default ``CopyOptions`` need the batched entry point. - """ - return _batch_entry_point_in_use() - - def assert_managed_holds(dev, buf, value, *, stream): """Assert a managed buffer holds ``value``. diff --git a/cuda_core/tests/memory/conftest.py b/cuda_core/tests/memory/conftest.py index aa83720fbc1..c0683f9aff3 100644 --- a/cuda_core/tests/memory/conftest.py +++ b/cuda_core/tests/memory/conftest.py @@ -5,18 +5,14 @@ Provides the device, stream and buffer fixtures shared by ``test_copy_batch.py`` (data movement) and ``test_copy_batch_options.py`` -(options and validation). Constants and helper functions that tests -import by name live in ``helpers.copy_batch``. +(options and validation). """ import pytest -from helpers.copy_batch import ( - COPY_BATCH_COUNT, - COPY_BATCH_SIZE, - uses_batch_entry_point, -) +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 @@ -39,7 +35,7 @@ def requires_copy_options(): unavailable, ``copy_batch`` falls back to per-copy ``cuMemcpyAsync``, which cannot convey attributes, so non-default options are rejected. """ - if not uses_batch_entry_point(): + 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" diff --git a/cuda_core/tests/memory/test_copy_batch.py b/cuda_core/tests/memory/test_copy_batch.py index eb6b89dd045..04c7cc41af9 100644 --- a/cuda_core/tests/memory/test_copy_batch.py +++ b/cuda_core/tests/memory/test_copy_batch.py @@ -5,8 +5,7 @@ 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. Options handling and argument -validation live in ``test_copy_batch_options.py``. +batch is correctly ordered on its stream. """ import pytest @@ -237,35 +236,6 @@ def test_ordered_between_prior_and_later_stream_work(self, device_bufs, copy_str for src in srcs: assert compare_buffer_to_constant(src, after) - @pytest.mark.agent_authored(model="Claude Opus 5") - def test_visible_on_other_stream_after_explicit_wait(self, copy_batch_device, device_bufs, copy_stream): - """A batch on one stream is not ordered against an unrelated stream. - - The second stream must be made to wait explicitly; once it does, - it observes the copied bytes. - """ - srcs, dsts = device_bufs - other = copy_batch_device.create_stream() - try: - for src in srcs: - src.fill(33, stream=copy_stream) - copy_batch(copy_stream, srcs, dsts) - - # Explicit cross-stream dependency, then observe from `other`. - other.wait(copy_stream) - probes = [copy_batch_device.memory_resource.allocate(COPY_BATCH_SIZE, stream=other) for _ in dsts] - for dst, probe in zip(dsts, probes): - dst.copy_to(probe, stream=other) - other.sync() - - for probe in probes: - assert compare_buffer_to_constant(probe, 33) - for probe in probes: - probe.close(other) - other.sync() - finally: - other.close() - @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. diff --git a/cuda_core/tests/memory/test_copy_batch_options.py b/cuda_core/tests/memory/test_copy_batch_options.py index 41d718612ed..eefd73c2276 100644 --- a/cuda_core/tests/memory/test_copy_batch_options.py +++ b/cuda_core/tests/memory/test_copy_batch_options.py @@ -4,8 +4,7 @@ """``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. The data-movement -behaviour itself lives in ``test_copy_batch.py``. +option field behaves, and every rejection path. """ import warnings @@ -16,17 +15,16 @@ # 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, compare_equal_buffers, set_buffer +from helpers.buffers import compare_buffer_to_constant, set_buffer from helpers.copy_batch import ( COPY_BATCH_SIZE, OVERLAP_WARNING_FILTER, assert_managed_holds, - uses_batch_entry_point, ) from cuda.core import Device, Host, LegacyPinnedMemoryResource from cuda.core._memory._copy_enums import _attr_run_starts -from cuda.core._utils.version import binding_version +from cuda.core._memory._copy_ops import _batch_entry_point_in_use, _normalize_copy_options from cuda.core.utils import ( CopyOptions, MemcpyOverlapMode, @@ -35,17 +33,39 @@ ) -class TestAttrRunStarts: - """Unit tests for the attrsIdxs run-length encoding. +class TestOptionsEncoding: + """How ``options`` becomes the driver's ``attrs`` / ``attrsIdxs`` pair. - Pure logic, no CUDA: ``attrs[k]`` applies to the copies in - ``[starts[k], starts[k + 1])``. + 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_broadcast_collapses_to_one_run(self): - attrs = [CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY)] * 4 - assert _attr_run_starts(attrs) == [0] + 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): @@ -111,62 +131,14 @@ def test_per_copy_options(self, requires_copy_options, h2d_bufs, copy_stream): CopyOptions(src_access_order=MemcpySrcAccessOrder.DURING_API_CALL), CopyOptions(src_access_order=MemcpySrcAccessOrder.STREAM), ] - assert _attr_run_starts(per_copy_options) == [0, 1, 2, 3] - + # 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_scalar_options_broadcast(self, requires_copy_options, copy_batch_device, h2d_bufs, copy_stream): - """A scalar option must apply to every copy. - - Verified three ways: the scalar collapses to a single driver - attribute, a scalar and an equivalent explicit per-copy list give - identical bytes, and a short list is *not* silently broadcast. - """ - srcs, scalar_dsts = h2d_bufs - device_mr = copy_batch_device.memory_resource - pinned_mr = LegacyPinnedMemoryResource() - n = len(srcs) - scalar_option = CopyOptions(src_access_order=MemcpySrcAccessOrder.ANY) - - for i, src in enumerate(srcs): - set_buffer(src, i + 95) - - # A scalar is expanded internally to n copies of one option, which - # the encoder then collapses back to a single driver entry. - assert _attr_run_starts([scalar_option] * n) == [0] - - copy_batch(copy_stream, srcs, scalar_dsts, options=scalar_option) - - listed_dsts = [device_mr.allocate(COPY_BATCH_SIZE, stream=copy_stream) for _ in srcs] - copy_batch(copy_stream, srcs, listed_dsts, options=[scalar_option] * n) - copy_stream.sync() - - # Every copy received the option, and both spellings agree. - for i, (scalar_dst, listed_dst) in enumerate(zip(scalar_dsts, listed_dsts)): - assert compare_buffer_to_constant(scalar_dst, i + 95) - scalar_host = pinned_mr.allocate(COPY_BATCH_SIZE) - listed_host = pinned_mr.allocate(COPY_BATCH_SIZE) - scalar_dst.copy_to(scalar_host, stream=copy_stream) - listed_dst.copy_to(listed_host, stream=copy_stream) - copy_stream.sync() - assert compare_equal_buffers(scalar_host, listed_host) - scalar_host.close(copy_stream) - listed_host.close(copy_stream) - - # A sequence is paired by index and never broadcast, so a - # one-element list is a length error rather than a scalar. - with pytest.raises(ValueError, match="options length"): - copy_batch(copy_stream, srcs, listed_dsts, options=[scalar_option]) - - for buf in listed_dsts: - buf.close(copy_stream) - copy_stream.sync() - @pytest.mark.agent_authored(model="Claude Opus 5") def test_location_hints(self, requires_copy_options, copy_batch_device, copy_stream): dev = copy_batch_device @@ -195,7 +167,12 @@ def test_location_hints(self, requires_copy_options, copy_batch_device, copy_str @pytest.mark.agent_authored(model="Claude Opus 5") def test_host_numa_location_hint(self, requires_copy_options, copy_batch_device, copy_stream): - """NUMA host hints round-trip on CUDA 13 and are rejected on CUDA 12.""" + """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)] @@ -203,16 +180,10 @@ def test_host_numa_location_hint(self, requires_copy_options, copy_batch_device, for i, src in enumerate(srcs): src.fill(i + 85, stream=copy_stream) - options = CopyOptions(dst_location_hint=Host(numa_id=0)) - - if binding_version() < (13, 0, 0): - with pytest.raises(TypeError, match="CUDA 13"): - copy_batch(copy_stream, srcs, dsts, options=options) - else: - copy_batch(copy_stream, srcs, dsts, options=options) - copy_stream.sync() - for i, dst in enumerate(dsts): - assert_managed_holds(dev, dst, 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) @@ -256,7 +227,6 @@ def test_overlap_mode_warns_only_on_discrete_gpu( copy_stream.sync() @pytest.mark.agent_authored(model="Claude Opus 5") - @pytest.mark.filterwarnings(OVERLAP_WARNING_FILTER) def test_default_overlap_mode_does_not_warn(self, h2d_bufs, copy_stream): srcs, dsts = h2d_bufs with warnings.catch_warnings(): @@ -378,7 +348,7 @@ class TestPerCopyFallback: @pytest.fixture(autouse=True) def _skip_if_batched(self): - if uses_batch_entry_point(): + if _batch_entry_point_in_use(): pytest.skip("cuMemcpyBatchAsync is in use; fallback path not exercised") @pytest.mark.agent_authored(model="Claude Opus 5") From 9e14ce11463ae6606de1ae784359148187181b4c Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Mon, 10 Aug 2026 16:57:43 -0700 Subject: [PATCH 6/6] Deduplicate _to_cumemlocation --- cuda_core/cuda/core/_memory/_copy_ops.pyx | 32 ++----------------- .../cuda/core/_memory/_managed_memory_ops.pyx | 32 +++++-------------- .../tests/memory/test_copy_batch_options.py | 10 +++++- 3 files changed, 20 insertions(+), 54 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_copy_ops.pyx b/cuda_core/cuda/core/_memory/_copy_ops.pyx index 11964b99f8e..0b005f2f5b5 100644 --- a/cuda_core/cuda/core/_memory/_copy_ops.pyx +++ b/cuda_core/cuda/core/_memory/_copy_ops.pyx @@ -15,6 +15,7 @@ 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 @@ -99,33 +100,6 @@ def _normalize_copy_options( ) -IF CUDA_CORE_BUILD_MAJOR >= 13: - cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc): - """Convert a _LocSpec dataclass to a cydriver.CUmemLocation struct.""" - 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: - cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc): - raise NotImplementedError( - "_to_cumemlocation requires cuda.core built against CUDA 13 headers" - ) - - cdef cydriver.CUmemcpyAttributes _to_cu_memcpy_attributes(object attr): """Convert a CopyOptions to a cydriver.CUmemcpyAttributes struct.""" cdef cydriver.CUmemcpyAttributes cu_attr @@ -137,9 +111,9 @@ cdef cydriver.CUmemcpyAttributes _to_cu_memcpy_attributes(object attr): 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) + cu_attr.srcLocHint = to_cumemlocation(src_loc) if dst_loc is not None: - cu_attr.dstLocHint = _to_cumemlocation(dst_loc) + cu_attr.dstLocHint = to_cumemlocation(dst_loc) return cu_attr diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx index 73175af87b1..dd9b21bf150 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx @@ -12,6 +12,10 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: from cuda.bindings cimport cydriver 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 @@ -76,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 @@ -213,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: @@ -277,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: @@ -350,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/tests/memory/test_copy_batch_options.py b/cuda_core/tests/memory/test_copy_batch_options.py index eefd73c2276..7b6065cccaf 100644 --- a/cuda_core/tests/memory/test_copy_batch_options.py +++ b/cuda_core/tests/memory/test_copy_batch_options.py @@ -140,7 +140,15 @@ def test_per_copy_options(self, requires_copy_options, h2d_bufs, copy_stream): assert compare_buffer_to_constant(dst, i + 40) @pytest.mark.agent_authored(model="Claude Opus 5") - def test_location_hints(self, requires_copy_options, copy_batch_device, copy_stream): + 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)]