cuda.core: introduce copy options for Buffer.copy_{to/from} (#2636)

cuda.core: introduce copy options for Buffer.copy_to/copy_from

Add an optional `options: CopyOptions` keyword argument to
`Buffer.copy_to` and `Buffer.copy_from`, exposing the same dataclass
introduced by `copy_batch` on the per-buffer path. When set, the copy
is submitted via `cuMemcpyWithAttributesAsync` if both cuda.bindings
and the driver are CUDA 13.2+, the stream isn't `LEGACY_DEFAULT_STREAM`,
and it isn't currently capturing; `PER_THREAD_DEFAULT_STREAM` is
accepted like any other stream. `options=None` is unaffected and
always uses the existing `cuMemcpyAsync` path.

Passing `options` together with `LEGACY_DEFAULT_STREAM` or a capturing
stream raises `TypeError`, matching `copy_batch`; a graph cannot
represent these attributes, so `GraphNode.memcpy` (a plain,
non-attributed copy) is the only way to get a copy into a graph today.

On an older cuda.bindings/driver, `src_access_order` values of
`STREAM` and `ANY` fall back to `cuMemcpyAsync` silently, since
stream-ordered access already satisfies both. `DURING_API_CALL`
promises all source reads complete before the call returns; a
stream-ordered fallback can't honor that, so it raises `RuntimeError`
instead of silently downgrading the guarantee, which could otherwise
let a caller overwrite a source buffer before the real read happens.

While aligning the two APIs, this also fixes two bugs in the existing
`copy_batch`: it previously rejected `PER_THREAD_DEFAULT_STREAM`
outright even though the driver accepts it, and its pre-CUDA-13
fallback loop silently ignored `DURING_API_CALL` despite a stale
comment claiming that case was already rejected. Both now match the
per-buffer behavior via a shared `_reject_unsupported_during_api_call`
helper in `_copy_enums.py`.

`cuMemcpyWithAttributesAsync` is absent from cuda.bindings older than
13.2, so it's routed through a small C++ function-pointer shim
(`_cpp/resource_handles.{cpp,hpp}`) resolved at runtime, avoiding a
hard Cython cimport that would break older-bindings builds.

Tests: new `tests/memory/test_copy_single_options.py` covers data
correctness across all `CopyOptions` fields, the `TypeError`/
`RuntimeError` rejection paths, default-stream-token and graph-capture
behavior (including that `options=None` is unaffected by either), and
`dst=None` auto-allocation. `test_copy_batch.py`/
`test_copy_batch_options.py` gain matching coverage for the
`copy_batch` fixes, plus direct unit tests of the shared
`DURING_API_CALL` guard. `test_memory.py` adds previously-missing
size-mismatch rejection tests for `copy_to`/`copy_from`.

Closes #2365.

Author: juenglin

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -111,6 +111,13 @@ decltype(&cuDevSmResourceSplit) p_cuDevSmResourceSplit = nullptr;
 void* p_cuDevSmResourceSplit = nullptr;
 #endif
 
+// cuMemcpyWithAttributesAsync (13.2+ — may be null on older drivers/bindings)
+#if CUDA_VERSION >= 13020
+decltype(&cuMemcpyWithAttributesAsync) p_cuMemcpyWithAttributesAsync = nullptr;
+#else
+void* p_cuMemcpyWithAttributesAsync = nullptr;
+#endif
+
 // NVRTC function pointers
 decltype(&nvrtcDestroyProgram) p_nvrtcDestroyProgram = nullptr;
 
@@ -2834,4 +2841,25 @@ bool has_sm_resource_split() noexcept {
     return p_cuDevSmResourceSplit != nullptr;
 }
 
+// ============================================================================
+// cuMemcpyWithAttributesAsync wrapper
+// ============================================================================
+
+CUresult memcpy_with_attributes_async(CUdeviceptr dst, CUdeviceptr src, size_t size,
+                                       void* attr, CUstream hStream) {
+#if CUDA_VERSION >= 13020
+    if (!p_cuMemcpyWithAttributesAsync) {
+        return CUDA_ERROR_NOT_SUPPORTED;
+    }
+    return p_cuMemcpyWithAttributesAsync(
+        dst, src, size, static_cast<CUmemcpyAttributes*>(attr), hStream);
+#else
+    return CUDA_ERROR_NOT_SUPPORTED;
+#endif
+}
+
+bool has_memcpy_with_attributes_async() noexcept {
+    return p_cuMemcpyWithAttributesAsync != nullptr;
+}
+
 }  // namespace cuda_core

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -146,6 +146,15 @@ extern decltype(&cuDevSmResourceSplit) p_cuDevSmResourceSplit;
 extern void* p_cuDevSmResourceSplit;
 #endif
 
+// cuMemcpyWithAttributesAsync (13.2+ — may be null on older drivers/bindings)
+#if CUDA_VERSION >= 13020
+extern decltype(&cuMemcpyWithAttributesAsync) p_cuMemcpyWithAttributesAsync;
+#else
+// cuMemcpyWithAttributesAsync doesn't exist in CUDA < 13.2 headers, so use a
+// void* placeholder. The pointer is always null when built against older CUDA.
+extern void* p_cuMemcpyWithAttributesAsync;
+#endif
+
 // ============================================================================
 // NVRTC function pointers
 //
@@ -1110,4 +1119,21 @@ CUresult sm_resource_split(CUdevResource* result, unsigned int nbGroups,
 // Returns true if the cuDevSmResourceSplit function pointer is available.
 bool has_sm_resource_split() noexcept;
 
+// ============================================================================
+// cuMemcpyWithAttributesAsync wrapper (13.2+)
+//
+// Calls through p_cuMemcpyWithAttributesAsync if available, otherwise returns
+// CUDA_ERROR_NOT_SUPPORTED. This avoids a direct Cython cimport of the
+// cydriver cdef function, which would fail at module init on cuda-bindings
+// < 13.2 (see https://github.com/NVIDIA/cuda-python/issues/2063).
+// ============================================================================
+
+// attr is void* so the Cython declaration doesn't reference CUmemcpyAttributes
+// (absent from cuda-bindings built against CUDA < 12.8). The C++ side casts it.
+CUresult memcpy_with_attributes_async(CUdeviceptr dst, CUdeviceptr src, size_t size,
+                                       void* attr, CUstream hStream);
+
+// Returns true if the cuMemcpyWithAttributesAsync function pointer is available.
+bool has_memcpy_with_attributes_async() noexcept;
+
 }  // namespace cuda_core

cuda_core/cuda/core/_memory/_buffer.pyi

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import cython
+from cuda.core._memory._copy_enums import CopyOptions
 from cuda.core._memory._device_memory_resource import DeviceMemoryResource
 from cuda.core._memory._ipc import IPCBufferDescriptor
 from cuda.core._memory._pinned_memory_resource import PinnedMemoryResource
@@ -170,7 +171,7 @@ class Buffer:
     def __exit__(self, exc_type, exc_val, exc_tb):
         ...
 
-    def copy_to(self, dst: Buffer | None=None, *, stream: Stream | GraphBuilder) -> Buffer:
+    def copy_to(self, dst: Buffer | None=None, *, stream: Stream | GraphBuilder, options: CopyOptions | None=None) -> Buffer:
         """Copy from this buffer to the dst buffer asynchronously on the given stream.
 
         Copies the data from this buffer to the provided dst buffer.
@@ -185,10 +186,32 @@ class Buffer:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Keyword argument specifying the stream for the
             asynchronous copy
+        options : :class:`~utils.CopyOptions`, optional
+            Transfer hints (source access order, location hints, overlap mode).
+            Honored when cuda.bindings and the driver are both CUDA 13.2 or
+            newer. Not accepted with ``LEGACY_DEFAULT_STREAM``; use
+            ``PER_THREAD_DEFAULT_STREAM`` instead. Not accepted with a
+            capturing stream either, since a graph cannot represent these
+            attributes; use :meth:`graph.GraphNode.memcpy` for a plain,
+            non-attributed copy node, or pass ``options=None``. On an older
+            cuda.bindings/driver, ``src_access_order`` values of ``STREAM``
+            and ``ANY`` are silently ignored; ``DURING_API_CALL`` raises
+            instead of silently downgrading its guarantee.
+
+        Raises
+        ------
+        TypeError
+            If ``options`` is not a :class:`~utils.CopyOptions` instance, or
+            if ``options`` is given together with ``LEGACY_DEFAULT_STREAM``
+            or a stream currently in graph capture mode.
+        RuntimeError
+            If ``options.src_access_order`` is ``DURING_API_CALL`` and
+            cuda.bindings or the driver is older than CUDA 13.2: falling
+            back to a plain copy cannot honor that guarantee.
 
         """
 
-    def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder) -> None:
+    def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder, options: CopyOptions | None=None) -> None:
         """Copy from the src buffer to this buffer asynchronously on the given stream.
 
         Parameters
@@ -198,7 +221,28 @@ class Buffer:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Keyword argument specifying the stream for the
             asynchronous copy
+        options : :class:`~utils.CopyOptions`, optional
+            Transfer hints (source access order, location hints, overlap mode).
+            Honored when cuda.bindings and the driver are both CUDA 13.2 or
+            newer. Not accepted with ``LEGACY_DEFAULT_STREAM``; use
+            ``PER_THREAD_DEFAULT_STREAM`` instead. Not accepted with a
+            capturing stream either, since a graph cannot represent these
+            attributes; use :meth:`graph.GraphNode.memcpy` for a plain,
+            non-attributed copy node, or pass ``options=None``. On an older
+            cuda.bindings/driver, ``src_access_order`` values of ``STREAM``
+            and ``ANY`` are silently ignored; ``DURING_API_CALL`` raises
+            instead of silently downgrading its guarantee.
 
+        Raises
+        ------
+        TypeError
+            If ``options`` is not a :class:`~utils.CopyOptions` instance, or
+            if ``options`` is given together with ``LEGACY_DEFAULT_STREAM``
+            or a stream currently in graph capture mode.
+        RuntimeError
+            If ``options.src_access_order`` is ``DURING_API_CALL`` and
+            cuda.bindings or the driver is older than CUDA 13.2: falling
+            back to a plain copy cannot honor that guarantee.
         """
 
     def fill(self, value: int | BufferProtocol, *, stream: Stream | GraphBuilder) -> None:

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -27,13 +27,20 @@ from cuda.core._resource_handles cimport (
 )
 from cuda.core.typing import DevicePointerType
 
-from cuda.core._stream cimport Stream, Stream_accept, default_stream
+from cuda.core._memory._copy_attributes cimport _with_attributes_available
+from cuda.core._memory._copy_attributes cimport _to_cu_memcpy_attributes  # no-cython-lint
+
+IF CUDA_CORE_BUILD_MAJOR >= 13:
+    from cuda.core._resource_handles cimport memcpy_with_attributes_async
+
+from cuda.core._stream cimport Stream, Stream_accept, Stream_is_legacy_default_token, 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._memory._copy_enums import CopyOptions, _reject_unsupported_during_api_call
 from cuda.core._utils.pycompat import BufferProtocol
 from cuda.core._dlpack import classify_dl_device, make_py_capsule
 from cuda.core._device import Device
@@ -159,6 +166,75 @@ cdef inline void _init_memory_attrs(Buffer self):
         self._mem_attrs_inited.store(True, memory_order_release)
 
 
+cdef bint _stream_is_capturing(Stream s):
+    cdef cydriver.CUstreamCaptureStatus cap_status
+    IF CUDA_CORE_BUILD_MAJOR >= 13:
+        HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &cap_status,
+                                                     NULL, NULL, NULL, NULL, NULL))
+    ELSE:
+        HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &cap_status,
+                                                     NULL, NULL, NULL, NULL))
+    return cap_status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE
+
+
+cdef void _do_copy_with_attributes(
+    cydriver.CUdeviceptr dst, cydriver.CUdeviceptr src, size_t nbytes,
+    object options, cydriver.CUstream hstream,
+):
+    IF CUDA_CORE_BUILD_MAJOR >= 13:
+        # Routed through the memcpy_with_attributes_async() C++ shim since
+        # cydriver.cuMemcpyWithAttributesAsync is absent from cuda-bindings < 13.2.
+        cdef cydriver.CUmemcpyAttributes cu_attr = _to_cu_memcpy_attributes(options)
+        with nogil:
+            HANDLE_RETURN(memcpy_with_attributes_async(dst, src, nbytes, <void*>&cu_attr, hstream))
+    ELSE:
+        pass  # unreachable: _with_attributes_available() is always False on CUDA 12
+
+
+cdef void _dispatch_buffer_copy(
+    cydriver.CUdeviceptr dst, cydriver.CUdeviceptr src, size_t nbytes,
+    Stream s, object options, str method_name,
+):
+    """Submit a single copy, honoring CopyOptions when the attributes path is usable."""
+    if options is None:
+        with nogil:
+            HANDLE_RETURN(cydriver.cuMemcpyAsync(dst, src, nbytes, as_cu(s._h_stream)))
+        return
+    if not isinstance(options, CopyOptions):
+        raise TypeError(
+            f"{method_name}: options must be CopyOptions, got {type(options).__name__}"
+        )
+    if Stream_is_legacy_default_token(s):
+        raise TypeError(
+            f"{method_name} does not accept LEGACY_DEFAULT_STREAM with options "
+            "(matches copy_batch); cuMemcpyWithAttributesAsync rejects it outright, "
+            "unlike PER_THREAD_DEFAULT_STREAM, which is a real stream to the driver "
+            "and is accepted. Pass an explicit stream, PER_THREAD_DEFAULT_STREAM, "
+            "or options=None."
+        )
+    if _stream_is_capturing(s):
+        raise TypeError(
+            f"{method_name} does not support graph capture with options "
+            "(matches copy_batch); the driver has no graph-node form of "
+            "cuMemcpyWithAttributesAsync, so options cannot be honored in a graph. "
+            "Use GraphNode.memcpy for a plain (non-attributed) copy node, or pass "
+            "options=None."
+        )
+    if _with_attributes_available():
+        _do_copy_with_attributes(dst, src, nbytes, options, as_cu(s._h_stream))
+    else:
+        _reject_unsupported_during_api_call(
+            options.src_access_order,
+            "cuda.bindings and the driver to both report CUDA 13.2 or newer "
+            "(cuMemcpyWithAttributesAsync is unavailable here)",
+        )
+        # STREAM and ANY never require access sooner than stream order, so
+        # cuMemcpyAsync satisfies them; options are otherwise silently
+        # ignored on this pre-CUDA-13.2 fallback path, matching copy_batch.
+        with nogil:
+            HANDLE_RETURN(cydriver.cuMemcpyAsync(dst, src, nbytes, as_cu(s._h_stream)))
+
+
 cdef class Buffer:
     """Represent a handle to allocated memory.
 
@@ -393,7 +469,8 @@ cdef class Buffer:
         self.close()
         return False
 
-    def copy_to(self, dst: Buffer | None = None, *, stream: Stream | GraphBuilder) -> Buffer:
+    def copy_to(self, dst: Buffer | None = None, *, stream: Stream | GraphBuilder,
+                options: CopyOptions | None = None) -> Buffer:
         """Copy from this buffer to the dst buffer asynchronously on the given stream.
 
         Copies the data from this buffer to the provided dst buffer.
@@ -408,6 +485,28 @@ cdef class Buffer:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Keyword argument specifying the stream for the
             asynchronous copy
+        options : :class:`~utils.CopyOptions`, optional
+            Transfer hints (source access order, location hints, overlap mode).
+            Honored when cuda.bindings and the driver are both CUDA 13.2 or
+            newer. Not accepted with ``LEGACY_DEFAULT_STREAM``; use
+            ``PER_THREAD_DEFAULT_STREAM`` instead. Not accepted with a
+            capturing stream either, since a graph cannot represent these
+            attributes; use :meth:`graph.GraphNode.memcpy` for a plain,
+            non-attributed copy node, or pass ``options=None``. On an older
+            cuda.bindings/driver, ``src_access_order`` values of ``STREAM``
+            and ``ANY`` are silently ignored; ``DURING_API_CALL`` raises
+            instead of silently downgrading its guarantee.
+
+        Raises
+        ------
+        TypeError
+            If ``options`` is not a :class:`~utils.CopyOptions` instance, or
+            if ``options`` is given together with ``LEGACY_DEFAULT_STREAM``
+            or a stream currently in graph capture mode.
+        RuntimeError
+            If ``options.src_access_order`` is ``DURING_API_CALL`` and
+            cuda.bindings or the driver is older than CUDA 13.2: falling
+            back to a plain copy cannot honor that guarantee.
 
         """
         cdef Stream s = Stream_accept(stream)
@@ -424,12 +523,12 @@ cdef class Buffer:
             raise ValueError( "buffer sizes mismatch between src and dst (sizes "
                              f"are: src={src_size}, dst={dst_size})"
             )
-        with nogil:
-            HANDLE_RETURN(cydriver.cuMemcpyAsync(
-                as_cu(dst._h_ptr), as_cu(self._h_ptr), src_size, as_cu(s._h_stream)))
+        _dispatch_buffer_copy(
+            as_cu(dst._h_ptr), as_cu(self._h_ptr), src_size, s, options, "copy_to")
         return dst
 
-    def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder) -> None:
+    def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder,
+                  options: CopyOptions | None = None) -> None:
         """Copy from the src buffer to this buffer asynchronously on the given stream.
 
         Parameters
@@ -439,7 +538,28 @@ cdef class Buffer:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
             Keyword argument specifying the stream for the
             asynchronous copy
+        options : :class:`~utils.CopyOptions`, optional
+            Transfer hints (source access order, location hints, overlap mode).
+            Honored when cuda.bindings and the driver are both CUDA 13.2 or
+            newer. Not accepted with ``LEGACY_DEFAULT_STREAM``; use
+            ``PER_THREAD_DEFAULT_STREAM`` instead. Not accepted with a
+            capturing stream either, since a graph cannot represent these
+            attributes; use :meth:`graph.GraphNode.memcpy` for a plain,
+            non-attributed copy node, or pass ``options=None``. On an older
+            cuda.bindings/driver, ``src_access_order`` values of ``STREAM``
+            and ``ANY`` are silently ignored; ``DURING_API_CALL`` raises
+            instead of silently downgrading its guarantee.
 
+        Raises
+        ------
+        TypeError
+            If ``options`` is not a :class:`~utils.CopyOptions` instance, or
+            if ``options`` is given together with ``LEGACY_DEFAULT_STREAM``
+            or a stream currently in graph capture mode.
+        RuntimeError
+            If ``options.src_access_order`` is ``DURING_API_CALL`` and
+            cuda.bindings or the driver is older than CUDA 13.2: falling
+            back to a plain copy cannot honor that guarantee.
         """
         cdef Stream s = Stream_accept(stream)
         cdef size_t dst_size = self._size
@@ -449,9 +569,8 @@ cdef class Buffer:
             raise ValueError( "buffer sizes mismatch between src and dst (sizes "
                              f"are: src={src_size}, dst={dst_size})"
             )
-        with nogil:
-            HANDLE_RETURN(cydriver.cuMemcpyAsync(
-                as_cu(self._h_ptr), as_cu(src._h_ptr), dst_size, as_cu(s._h_stream)))
+        _dispatch_buffer_copy(
+            as_cu(self._h_ptr), as_cu(src._h_ptr), dst_size, s, options, "copy_from")
 
     def fill(self, value: int | BufferProtocol, *, stream: Stream | GraphBuilder) -> None:
         """Fill this buffer with a repeating byte pattern.

cuda_core/cuda/core/_memory/_copy_attributes.pxd

@@ -11,8 +11,19 @@ from cuda.core._utils.version cimport cy_binding_version, cy_driver_version  # n
 
 
 IF CUDA_CORE_BUILD_MAJOR >= 13:
+    from cuda.core._resource_handles cimport has_memcpy_with_attributes_async
+
     cdef inline bint _with_attributes_available():
-        return cy_driver_version() >= (13, 2, 0) and cy_binding_version() >= (13, 2, 0)
+        # has_memcpy_with_attributes_async() says whether the installed
+        # cuda-bindings actually exports cuMemcpyWithAttributesAsync (13.2+);
+        # the version checks alone are not sufficient, since cuda.core's build
+        # can be paired with a cuda-bindings install older than what it built
+        # against (see https://github.com/NVIDIA/cuda-python/issues/2063).
+        return (
+            has_memcpy_with_attributes_async()
+            and cy_driver_version() >= (13, 2, 0)
+            and cy_binding_version() >= (13, 2, 0)
+        )
 ELSE:
     cdef inline bint _with_attributes_available():
         return False