Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions cuda_core/cuda/core/_cpp/resource_handles.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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
26 changes: 26 additions & 0 deletions cuda_core/cuda/core/_cpp/resource_handles.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +149 to +156

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have concerns about (1) minor-version-specific version gating, and (2) scope-creep inherent in using resource_handles for things besides resource lifetime management. I see this merely extends an existing pattern, so neither of these should block merge; but I'd like to follow up.


// ============================================================================
// NVRTC function pointers
//
Expand Down Expand Up @@ -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
48 changes: 46 additions & 2 deletions cuda_core/cuda/core/_memory/_buffer.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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:
Expand Down
137 changes: 128 additions & 9 deletions cuda_core/cuda/core/_memory/_buffer.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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.
Expand Down
13 changes: 12 additions & 1 deletion cuda_core/cuda/core/_memory/_copy_attributes.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading