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

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,162 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from cuda.core._utils.cuda_utils import driver
+from cuda.core._utils.pycompat import StrEnum
+from cuda.core._utils.version import binding_version
+
+if TYPE_CHECKING:
+    from cuda.core._device import Device
+    from cuda.core._host import Host
+
+
+__all__ = ["CopyOptions", "MemcpyOverlapMode", "MemcpySrcAccessOrder"]
+
+
+class MemcpySrcAccessOrder(StrEnum):
+    """Source access order hint for batched memcpy operations.
+
+    Maps to ``CUmemcpySrcAccessOrder``. The ``INVALID`` and ``MAX``
+    sentinel values from the driver enum are excluded from the public
+    Python surface.
+    """
+
+    STREAM = "stream"
+    DURING_API_CALL = "during_api_call"
+    ANY = "any"
+
+
+class MemcpyOverlapMode(StrEnum):
+    """Overlap mode hint for batched memcpy operations.
+
+    Maps to ``CUmemcpyFlags``. Renamed from "flags" to "overlap_mode"
+    for clarity; the only non-default flag is CE/compute overlap
+    (Tegra).
+    """
+
+    DEFAULT = "default"
+    PREFER_OVERLAP_WITH_COMPUTE = "prefer_overlap_with_compute"
+
+
+@dataclasses.dataclass(frozen=True)
+class CopyOptions:
+    """Attribute bundle for a single copy within a batched memcpy.
+
+    Parameters
+    ----------
+    src_access_order : :class:`MemcpySrcAccessOrder` or str
+        Hint describing how the source will be accessed.
+        Default is ``"stream"`` (stream-ordered access).
+    src_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None
+        Hint for the source memory location. ``None`` means no hint.
+    dst_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None
+        Hint for the destination memory location. ``None`` means no hint.
+    overlap_mode : :class:`MemcpyOverlapMode` or str
+        Hint for copy-engine / compute overlap.  Only meaningful on
+        integrated (Tegra) GPUs; on discrete GPUs the driver silently
+        ignores it and a :class:`UserWarning` is emitted.
+        Default is ``"default"``.
+    """
+
+    src_access_order: MemcpySrcAccessOrder | str = "stream"
+    src_location_hint: Device | Host | None = None
+    dst_location_hint: Device | Host | None = None
+    overlap_mode: MemcpyOverlapMode | str = "default"
+
+    def __post_init__(self):
+        # Frozen, unlike the other *Options dataclasses in cuda.core, because
+        # the batched-API contract agreed in NVIDIA/cuda-python#1775 specifies
+        # immutable per-call options:
+        # https://github.com/NVIDIA/cuda-python/pull/1775#issuecomment-4355502334
+        #
+        # Normalizing str -> StrEnum therefore has to go through
+        # object.__setattr__; a plain assignment would raise
+        # FrozenInstanceError. Done here rather than at use so that a typo
+        # fails at construction and the field always holds the enum.
+        if isinstance(self.src_access_order, str):
+            try:
+                object.__setattr__(
+                    self,
+                    "src_access_order",
+                    MemcpySrcAccessOrder(self.src_access_order),
+                )
+            except ValueError as exc:
+                raise ValueError(f"invalid src_access_order: {self.src_access_order!r}") from exc
+        if isinstance(self.overlap_mode, str):
+            try:
+                object.__setattr__(
+                    self,
+                    "overlap_mode",
+                    MemcpyOverlapMode(self.overlap_mode),
+                )
+            except ValueError as exc:
+                raise ValueError(f"invalid overlap_mode: {self.overlap_mode!r}") from exc
+
+    def _to_driver_enum(self) -> int:
+        """Return the driver CUmemcpySrcAccessOrder value."""
+        if not _SRC_ACCESS_ORDER_TO_DRIVER:
+            raise NotImplementedError(_CUDA13_REQUIRED)
+        return _SRC_ACCESS_ORDER_TO_DRIVER[MemcpySrcAccessOrder(self.src_access_order)]
+
+    def _to_driver_flags(self) -> int:
+        """Return the driver CUmemcpyFlags value."""
+        if not _OVERLAP_MODE_TO_DRIVER:
+            raise NotImplementedError(_CUDA13_REQUIRED)
+        return _OVERLAP_MODE_TO_DRIVER[MemcpyOverlapMode(self.overlap_mode)]
+
+
+_CUDA13_REQUIRED = "copy attributes require cuda.bindings 13.0 or newer"
+
+# CUmemcpySrcAccessOrder and CUmemcpyFlags are exposed by cuda.bindings 13.0+,
+# so these maps are empty when it is older. Nothing reaches them there:
+# copy_batch refuses non-default CopyOptions when the batched entry point is
+# unavailable.
+#
+# Keyed by ``str``: under ``python_version = "3.10"`` mypy resolves StrEnum to
+# the unstubbed backports shim and so infers the members as plain ``str``.
+# StrEnum members are ``str`` instances, so this holds on every version. The
+# values are wrapped in ``int()`` because the driver enums are untyped.
+_SRC_ACCESS_ORDER_TO_DRIVER: dict[str, int]
+_OVERLAP_MODE_TO_DRIVER: dict[str, int]
+
+if binding_version() >= (13, 0, 0):
+    _src_order = driver.CUmemcpySrcAccessOrder
+    _flags = driver.CUmemcpyFlags
+    _SRC_ACCESS_ORDER_TO_DRIVER = {
+        MemcpySrcAccessOrder.STREAM: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_STREAM),
+        MemcpySrcAccessOrder.DURING_API_CALL: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL),
+        MemcpySrcAccessOrder.ANY: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_ANY),
+    }
+    _OVERLAP_MODE_TO_DRIVER = {
+        MemcpyOverlapMode.DEFAULT: int(_flags.CU_MEMCPY_FLAG_DEFAULT),
+        MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE: int(_flags.CU_MEMCPY_FLAG_PREFER_OVERLAP_WITH_COMPUTE),
+    }
+    del _src_order, _flags
+else:
+    _SRC_ACCESS_ORDER_TO_DRIVER = {}
+    _OVERLAP_MODE_TO_DRIVER = {}
+
+
+def _attr_run_starts(attrs: Sequence[CopyOptions]) -> list[int]:
+    """Return the start index of each maximal run of equal attributes.
+
+    This mirrors the ``attrsIdxs`` indirection that ``cuMemcpyBatchAsync``
+    expects: ``attrs[k]`` applies to the copies in
+    ``[starts[k], starts[k + 1])``. Collapsing equal neighbours means a
+    broadcast attribute is passed to the driver once (``numAttrs == 1``)
+    rather than repeated per copy.
+    """
+    starts: list[int] = []
+    prev: CopyOptions | None = None
+    for i, attr in enumerate(attrs):
+        if i == 0 or attr != prev:
+            starts.append(i)
+        prev = attr
+    return starts

cuda_core/cuda/core/_memory/_copy_ops.pyi

@@ -0,0 +1,86 @@
+# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_memory/_copy_ops.pyx
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from cuda.core._memory._buffer import Buffer
+from cuda.core._memory._copy_enums import CopyOptions
+from cuda.core._stream import Stream
+
+_SINGLE_COPY_HINT = 'Buffer.copy_to / Buffer.copy_from'
+_DEFAULT_COPY_OPTIONS = CopyOptions()
+
+def _batch_entry_point_in_use() -> bool:
+    """Internal: expose the dispatch predicate so tests can gate on it."""
+
+def _normalize_copy_options(options: CopyOptions | Sequence[CopyOptions] | None, n: int) -> tuple[CopyOptions, ...]:
+    """Expand ``options`` to exactly one :class:`CopyOptions` per copy.
+
+    ``None`` and a scalar broadcast; a sequence pairs by index and must
+    already have length ``n``.
+
+    Internal, but deliberately importable: options are hints that change
+    how the driver stages a transfer and never the bytes it produces, so
+    this expansion (and the run encoding applied to it) is the only
+    observable evidence that a scalar reached every copy.
+    """
+
+def copy_batch(stream: Stream, srcs: Sequence[Buffer], dsts: Sequence[Buffer], *, options: CopyOptions | Sequence[CopyOptions] | None=None) -> None:
+    """Copy a batch of buffers asynchronously.
+
+    Sizes are taken from the source buffers and each destination must
+    match. For a single buffer, use :meth:`Buffer.copy_to` or
+    :meth:`Buffer.copy_from`.
+
+    The driver provides no graph-node form of ``cuMemcpyBatchAsync``, so
+    this cannot be captured into a graph. Build graph copies with
+    :meth:`graph.GraphNode.memcpy` or per-buffer :meth:`Buffer.copy_to`.
+
+    Parameters
+    ----------
+    stream : :class:`~_stream.Stream`
+        Stream for the asynchronous copy. First positional and required
+        (mirrors :func:`launch`). Unlike most stream-taking APIs this does
+        not accept a :class:`~graph.GraphBuilder`; one is rejected with
+        ``TypeError`` because the copy cannot be captured.
+    srcs : Sequence[:class:`Buffer`]
+        Source buffers. Must be a sequence, not a single Buffer.
+    dsts : Sequence[:class:`Buffer`]
+        Destination buffers. Must match ``len(srcs)``.
+    options : :class:`CopyOptions` | Sequence[:class:`CopyOptions`] | None
+        Per-copy options. A single value applies to every copy; a
+        sequence pairs by index and must match ``len(srcs)``. ``None``
+        uses stream-ordered defaults.
+
+    Raises
+    ------
+    ValueError
+        If lengths or sizes mismatch.
+    TypeError
+        If a single Buffer is passed instead of a sequence.
+    NotImplementedError
+        If non-default ``options`` are given where
+        ``cuMemcpyBatchAsync`` is unavailable (see Notes).
+
+    Notes
+    -----
+    Batching through ``cuMemcpyBatchAsync`` requires all three of:
+    ``cuda.core`` built against CUDA 13 headers, ``cuda.bindings`` 13.0 or
+    newer, and a driver reporting CUDA 13.0 or newer
+    (``cuDriverGetVersion() >= 13000``). ``cuda.bindings`` binds only the
+    CUDA 13.0 revision of the entry point, so a driver that predates it is
+    refused even where it implements the earlier CUDA 12.8 signature.
+
+    Short of that, the copies fall back to a Python-level loop over
+    ``cuMemcpyAsync``, which is semantically equivalent but does not
+    amortize launch overhead. The fallback has no way to convey
+    :class:`CopyOptions` to the driver, so non-default options raise
+    :class:`NotImplementedError` there rather than being silently ignored.
+
+    Warns
+    -----
+    UserWarning
+        If ``overlap_mode='prefer_overlap_with_compute'`` is requested
+        on a non-integrated (discrete) GPU.
+    """