fallback for CUDA 12 and type annotations

Author: juenglin

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

@@ -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]:

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,26 +28,41 @@ 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``
         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.
+    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.