cuda.core: Add copy_batch to cuda.core.utils (#2593)

* cuda.core: Add copy_batch to cuda.core.utils

* fallback for CUDA 12 and type annotations

* be more precise about CUDA requirements

* skip tests on Windows that require managed memory

* rework some tests

* Deduplicate _to_cumemlocation

* add missing file

* address review feedback

* review feedback: don't assume NUMA capabilities

* review feedback: clarify buffer requirements for async batched copies

* review feedback: explicitly reject special default streams

* review feedback: explicitly reject capturing streams

* review feedback: drop warning about unsupported PREFER_OVERLAP_WITH_COMPUTE hint

* review feedback: add missing descriptions for copy options values

* review feedback: align CopyOptions validation with existing practice

* review feedback: drop conditional imports for type checking

* account for CUDA 12/13 driver differences

* CUDA 12: drop rejection of unsupported copy options

* simplify tests

Author: juenglin

cuda_core/AGENTS.md

@@ -151,6 +151,13 @@ a `StrEnum` is accepted as an argument, a `str` should also be acceptable.  An
 invalid value should raise an exception.  When a function returns a `str` drawn
 from a small number of values, return a `StrEnum` subclass instead.
 
+For `__post_init__` validation in frozen dataclasses, use the
+`not isinstance(value, EnumType) → try EnumType(value) except (ValueError,
+TypeError)` pattern (modelled on `_normalize_enum` in
+`cuda/core/texture/_texture.pyx`). This accepts the enum itself or a valid
+string, and raises `ValueError` eagerly for any other type rather than
+silently storing it.
+
 ### Exception handling
 
 Raising exceptions is preferred over a C-style return code that must be checked

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)

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

cuda_core/cuda/core/_memory/_copy_enums.py

@@ -0,0 +1,178 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Sequence
+
+from cuda.core._device import Device
+from cuda.core._host import Host
+from cuda.core._utils.cuda_utils import driver
+from cuda.core._utils.pycompat import StrEnum
+from cuda.core._utils.version import binding_version
+
+__all__ = ["CopyOptions", "MemcpyOverlapMode", "MemcpySrcAccessOrder"]
+
+
+class MemcpySrcAccessOrder(StrEnum):
+    """Source access order hint for batched memcpy operations.
+
+    Maps to ``CUmemcpySrcAccessOrder``.
+
+    ``STREAM``
+        Source reads follow stream order. Earlier stream work may still be
+        accessing the source when the copy is enqueued.
+    ``DURING_API_CALL``
+        The driver may read the source out of stream order, but all reads
+        are complete before :func:`copy_batch` returns. No earlier stream
+        work may be accessing the source at the time of the call.
+    ``ANY``
+        The driver may read the source after the call returns. The caller
+        must keep the source unchanged until the copy completes in stream
+        order. No earlier stream work may be accessing the source.
+    """
+
+    STREAM = "stream"
+    DURING_API_CALL = "during_api_call"
+    ANY = "any"
+
+
+class MemcpyOverlapMode(StrEnum):
+    """Overlap mode hint for batched memcpy operations.
+
+    Maps to ``CUmemcpyFlags``.
+
+    ``DEFAULT``
+        No overlap preference; the driver uses its default scheduling.
+    ``PREFER_OVERLAP_WITH_COMPUTE``
+        Hint that the copy should preferably overlap with concurrent
+        compute work. This is advisory and may be ignored depending on
+        the platform and copy parameters.
+    """
+
+    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. Honored only for managed
+        memory on devices with concurrent managed access and for
+        system-allocated pageable memory on devices with pageable memory
+        access; ignored for all other memory types. Does not prefetch
+        memory and does not set persistent memory advice.
+        ``None`` means no hint.
+    dst_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None
+        Hint for the destination memory location. Same semantics and
+        restrictions as ``src_location_hint``. ``None`` means no hint.
+    overlap_mode : :class:`MemcpyOverlapMode` or str
+        Hint requesting that the copy overlap with concurrent compute work.
+        This is advisory; it has an effect only on devices that support it.
+        Default is ``"default"``.
+    """
+
+    src_access_order: MemcpySrcAccessOrder | str = "stream"
+    src_location_hint: Device | Host | None = None
+    dst_location_hint: Device | Host | None = None
+    overlap_mode: MemcpyOverlapMode | str = "default"
+
+    def __post_init__(self):
+        # Frozen, unlike the other *Options dataclasses in cuda.core, because
+        # the batched-API contract agreed in NVIDIA/cuda-python#1775 specifies
+        # immutable per-call options:
+        # https://github.com/NVIDIA/cuda-python/pull/1775#issuecomment-4355502334
+        #
+        # Normalizing str -> StrEnum therefore has to go through
+        # object.__setattr__; a plain assignment would raise
+        # FrozenInstanceError. Done here rather than at use so that a typo
+        # fails at construction and the field always holds the enum.
+        if not isinstance(self.src_access_order, MemcpySrcAccessOrder):
+            try:
+                object.__setattr__(
+                    self,
+                    "src_access_order",
+                    MemcpySrcAccessOrder(self.src_access_order),
+                )
+            except (ValueError, TypeError) as exc:
+                raise ValueError(f"invalid src_access_order: {self.src_access_order!r}") from exc
+        if not isinstance(self.overlap_mode, MemcpyOverlapMode):
+            try:
+                object.__setattr__(
+                    self,
+                    "overlap_mode",
+                    MemcpyOverlapMode(self.overlap_mode),
+                )
+            except (ValueError, TypeError) as exc:
+                raise ValueError(f"invalid overlap_mode: {self.overlap_mode!r}") from exc
+
+    def _to_driver_enum(self) -> int:
+        """Return the driver CUmemcpySrcAccessOrder value."""
+        if not _SRC_ACCESS_ORDER_TO_DRIVER:
+            raise NotImplementedError(_CUDA13_REQUIRED)
+        return _SRC_ACCESS_ORDER_TO_DRIVER[MemcpySrcAccessOrder(self.src_access_order)]
+
+    def _to_driver_flags(self) -> int:
+        """Return the driver CUmemcpyFlags value."""
+        if not _OVERLAP_MODE_TO_DRIVER:
+            raise NotImplementedError(_CUDA13_REQUIRED)
+        return _OVERLAP_MODE_TO_DRIVER[MemcpyOverlapMode(self.overlap_mode)]
+
+
+_CUDA13_REQUIRED = "copy attributes require cuda.bindings 13.0 or newer"
+
+# CUmemcpySrcAccessOrder and CUmemcpyFlags are exposed by cuda.bindings 13.0+,
+# so these maps are empty when it is older. Nothing reaches them there:
+# copy_batch refuses non-default CopyOptions when the batched entry point is
+# unavailable.
+#
+# Keyed by ``str``: under ``python_version = "3.10"`` mypy resolves StrEnum to
+# the unstubbed backports shim and so infers the members as plain ``str``.
+# StrEnum members are ``str`` instances, so this holds on every version. The
+# values are wrapped in ``int()`` because the driver enums are untyped.
+_SRC_ACCESS_ORDER_TO_DRIVER: dict[str, int]
+_OVERLAP_MODE_TO_DRIVER: dict[str, int]
+
+if binding_version() >= (13, 0, 0):
+    _src_order = driver.CUmemcpySrcAccessOrder
+    _flags = driver.CUmemcpyFlags
+    _SRC_ACCESS_ORDER_TO_DRIVER = {
+        MemcpySrcAccessOrder.STREAM: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_STREAM),
+        MemcpySrcAccessOrder.DURING_API_CALL: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL),
+        MemcpySrcAccessOrder.ANY: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_ANY),
+    }
+    _OVERLAP_MODE_TO_DRIVER = {
+        MemcpyOverlapMode.DEFAULT: int(_flags.CU_MEMCPY_FLAG_DEFAULT),
+        MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE: int(_flags.CU_MEMCPY_FLAG_PREFER_OVERLAP_WITH_COMPUTE),
+    }
+    del _src_order, _flags
+else:
+    _SRC_ACCESS_ORDER_TO_DRIVER = {}
+    _OVERLAP_MODE_TO_DRIVER = {}
+
+
+def _attr_run_starts(attrs: Sequence[CopyOptions]) -> list[int]:
+    """Return the start index of each maximal run of equal attributes.
+
+    This mirrors the ``attrsIdxs`` indirection that ``cuMemcpyBatchAsync``
+    expects: ``attrs[k]`` applies to the copies in
+    ``[starts[k], starts[k + 1])``. Collapsing equal neighbours means a
+    broadcast attribute is passed to the driver once (``numAttrs == 1``)
+    rather than repeated per copy.
+    """
+    starts: list[int] = []
+    prev: CopyOptions | None = None
+    for i, attr in enumerate(attrs):
+        if i == 0 or attr != prev:
+            starts.append(i)
+        prev = attr
+    return starts

cuda_core/cuda/core/_memory/_copy_ops.pyi

@@ -0,0 +1,85 @@
+# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_memory/_copy_ops.pyx
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from cuda.core._memory._buffer import Buffer
+from cuda.core._memory._copy_enums import CopyOptions
+from cuda.core._stream import Stream
+
+_SINGLE_COPY_HINT = 'Buffer.copy_to / Buffer.copy_from'
+
+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.
+
+    Source buffer and destination buffer sizes 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. Both passing a
+    :class:`~graph.GraphBuilder` and passing its underlying
+    :attr:`~graph.GraphBuilder.stream` while capture is active are
+    rejected. Build graph copies with
+    :meth:`graph.GraphNode.memcpy` or per-buffer :meth:`Buffer.copy_to`.
+
+    Parameters
+    ----------
+    stream : :class:`~_stream.Stream`
+        Stream for the asynchronous copy. First positional and required
+        (mirrors :func:`launch`). Does not accept a capturing stream
+        (including a :class:`~graph.GraphBuilder`'s underlying stream); use
+        :meth:`graph.GraphNode.memcpy` or per-buffer
+        :meth:`Buffer.copy_to` to build copies into a graph.
+    srcs : Sequence[:class:`Buffer`]
+        Source buffers. Must be a sequence, not a single Buffer.
+    dsts : Sequence[:class:`Buffer`]
+        Destination buffers. Must match ``len(srcs)``.
+    options : :class:`CopyOptions` | Sequence[:class:`CopyOptions`] | None
+        Per-copy options. A single value applies to every copy; a
+        sequence pairs by index and must match ``len(srcs)``. ``None``
+        uses stream-ordered defaults.
+
+    Raises
+    ------
+    ValueError
+        If lengths or sizes mismatch.
+    TypeError
+        If a single Buffer is passed instead of a sequence, if a
+        default-stream token (``LEGACY_DEFAULT_STREAM`` /
+        ``PER_THREAD_DEFAULT_STREAM``) is passed, or if the stream is
+        currently in graph capture mode.
+
+    Notes
+    -----
+    Batching through ``cuMemcpyBatchAsync`` requires all three of:
+    ``cuda.core`` built against CUDA 13 headers, ``cuda.bindings`` 13.0 or
+    newer, and a driver reporting CUDA 13.0 or newer
+    (``cuDriverGetVersion() >= 13000``). ``cuda.bindings`` binds only the
+    CUDA 13.0 revision of the entry point, so a driver that predates it is
+    refused even where it implements the earlier CUDA 12.8 signature.
+
+    The driver may execute batch items concurrently and in any order.
+    A batch must therefore not contain copies where the source range of
+    one copy overlaps the destination range of another; such aliasing
+    produces undefined results. Detecting overlaps at runtime is
+    impractical; callers are responsible for ensuring no aliasing exists.
+
+    On pre-CUDA 13 installs the copies fall back to a Python-level loop
+    over ``cuMemcpyAsync``, so the potential performance benefit of
+    asynchronous batched copies is not realized. :class:`CopyOptions` are
+    silently ignored on the fallback path.
+
+    """