Skip to content

Commit 250a442

Browse files
Andy-JostAndy-Jost
committed
Merge origin/main into ajost/closed-resource-validation
2 parents bd21d3c + d6afb93 commit 250a442

26 files changed

Lines changed: 1581 additions & 343 deletions

‎cuda_core/cuda/core/_cpp/resource_handles.cpp‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,13 @@ decltype(&cuDevSmResourceSplit) p_cuDevSmResourceSplit = nullptr;
111111
void* p_cuDevSmResourceSplit = nullptr;
112112
#endif
113113

114+
// cuMemcpyWithAttributesAsync (13.2+ — may be null on older drivers/bindings)
115+
#if CUDA_VERSION >= 13020
116+
decltype(&cuMemcpyWithAttributesAsync) p_cuMemcpyWithAttributesAsync = nullptr;
117+
#else
118+
void* p_cuMemcpyWithAttributesAsync = nullptr;
119+
#endif
120+
114121
// NVRTC function pointers
115122
decltype(&nvrtcDestroyProgram) p_nvrtcDestroyProgram = nullptr;
116123

@@ -2834,4 +2841,25 @@ bool has_sm_resource_split() noexcept {
28342841
return p_cuDevSmResourceSplit != nullptr;
28352842
}
28362843

2844+
// ============================================================================
2845+
// cuMemcpyWithAttributesAsync wrapper
2846+
// ============================================================================
2847+
2848+
CUresult memcpy_with_attributes_async(CUdeviceptr dst, CUdeviceptr src, size_t size,
2849+
void* attr, CUstream hStream) {
2850+
#if CUDA_VERSION >= 13020
2851+
if (!p_cuMemcpyWithAttributesAsync) {
2852+
return CUDA_ERROR_NOT_SUPPORTED;
2853+
}
2854+
return p_cuMemcpyWithAttributesAsync(
2855+
dst, src, size, static_cast<CUmemcpyAttributes*>(attr), hStream);
2856+
#else
2857+
return CUDA_ERROR_NOT_SUPPORTED;
2858+
#endif
2859+
}
2860+
2861+
bool has_memcpy_with_attributes_async() noexcept {
2862+
return p_cuMemcpyWithAttributesAsync != nullptr;
2863+
}
2864+
28372865
} // namespace cuda_core

‎cuda_core/cuda/core/_cpp/resource_handles.hpp‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,15 @@ extern decltype(&cuDevSmResourceSplit) p_cuDevSmResourceSplit;
146146
extern void* p_cuDevSmResourceSplit;
147147
#endif
148148

149+
// cuMemcpyWithAttributesAsync (13.2+ — may be null on older drivers/bindings)
150+
#if CUDA_VERSION >= 13020
151+
extern decltype(&cuMemcpyWithAttributesAsync) p_cuMemcpyWithAttributesAsync;
152+
#else
153+
// cuMemcpyWithAttributesAsync doesn't exist in CUDA < 13.2 headers, so use a
154+
// void* placeholder. The pointer is always null when built against older CUDA.
155+
extern void* p_cuMemcpyWithAttributesAsync;
156+
#endif
157+
149158
// ============================================================================
150159
// NVRTC function pointers
151160
//
@@ -1110,4 +1119,21 @@ CUresult sm_resource_split(CUdevResource* result, unsigned int nbGroups,
11101119
// Returns true if the cuDevSmResourceSplit function pointer is available.
11111120
bool has_sm_resource_split() noexcept;
11121121

1122+
// ============================================================================
1123+
// cuMemcpyWithAttributesAsync wrapper (13.2+)
1124+
//
1125+
// Calls through p_cuMemcpyWithAttributesAsync if available, otherwise returns
1126+
// CUDA_ERROR_NOT_SUPPORTED. This avoids a direct Cython cimport of the
1127+
// cydriver cdef function, which would fail at module init on cuda-bindings
1128+
// < 13.2 (see https://github.com/NVIDIA/cuda-python/issues/2063).
1129+
// ============================================================================
1130+
1131+
// attr is void* so the Cython declaration doesn't reference CUmemcpyAttributes
1132+
// (absent from cuda-bindings built against CUDA < 12.8). The C++ side casts it.
1133+
CUresult memcpy_with_attributes_async(CUdeviceptr dst, CUdeviceptr src, size_t size,
1134+
void* attr, CUstream hStream);
1135+
1136+
// Returns true if the cuMemcpyWithAttributesAsync function pointer is available.
1137+
bool has_memcpy_with_attributes_async() noexcept;
1138+
11131139
} // namespace cuda_core

‎cuda_core/cuda/core/_memory/_buffer.pyi‎

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from __future__ import annotations
44

55
import cython
6+
from cuda.core._memory._copy_enums import CopyOptions
67
from cuda.core._memory._device_memory_resource import DeviceMemoryResource
78
from cuda.core._memory._ipc import IPCBufferDescriptor
89
from cuda.core._memory._pinned_memory_resource import PinnedMemoryResource
@@ -170,7 +171,7 @@ class Buffer:
170171
def __exit__(self, exc_type, exc_val, exc_tb):
171172
...
172173

173-
def copy_to(self, dst: Buffer | None=None, *, stream: Stream | GraphBuilder) -> Buffer:
174+
def copy_to(self, dst: Buffer | None=None, *, stream: Stream | GraphBuilder, options: CopyOptions | None=None) -> Buffer:
174175
"""Copy from this buffer to the dst buffer asynchronously on the given stream.
175176
176177
Copies the data from this buffer to the provided dst buffer.
@@ -185,10 +186,32 @@ class Buffer:
185186
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
186187
Keyword argument specifying the stream for the
187188
asynchronous copy
189+
options : :class:`~utils.CopyOptions`, optional
190+
Transfer hints (source access order, location hints, overlap mode).
191+
Honored when cuda.bindings and the driver are both CUDA 13.2 or
192+
newer. Not accepted with ``LEGACY_DEFAULT_STREAM``; use
193+
``PER_THREAD_DEFAULT_STREAM`` instead. Not accepted with a
194+
capturing stream either, since a graph cannot represent these
195+
attributes; use :meth:`graph.GraphNode.memcpy` for a plain,
196+
non-attributed copy node, or pass ``options=None``. On an older
197+
cuda.bindings/driver, ``src_access_order`` values of ``STREAM``
198+
and ``ANY`` are silently ignored; ``DURING_API_CALL`` raises
199+
instead of silently downgrading its guarantee.
200+
201+
Raises
202+
------
203+
TypeError
204+
If ``options`` is not a :class:`~utils.CopyOptions` instance, or
205+
if ``options`` is given together with ``LEGACY_DEFAULT_STREAM``
206+
or a stream currently in graph capture mode.
207+
RuntimeError
208+
If ``options.src_access_order`` is ``DURING_API_CALL`` and
209+
cuda.bindings or the driver is older than CUDA 13.2: falling
210+
back to a plain copy cannot honor that guarantee.
188211
189212
"""
190213

191-
def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder) -> None:
214+
def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder, options: CopyOptions | None=None) -> None:
192215
"""Copy from the src buffer to this buffer asynchronously on the given stream.
193216
194217
Parameters
@@ -198,7 +221,28 @@ class Buffer:
198221
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
199222
Keyword argument specifying the stream for the
200223
asynchronous copy
224+
options : :class:`~utils.CopyOptions`, optional
225+
Transfer hints (source access order, location hints, overlap mode).
226+
Honored when cuda.bindings and the driver are both CUDA 13.2 or
227+
newer. Not accepted with ``LEGACY_DEFAULT_STREAM``; use
228+
``PER_THREAD_DEFAULT_STREAM`` instead. Not accepted with a
229+
capturing stream either, since a graph cannot represent these
230+
attributes; use :meth:`graph.GraphNode.memcpy` for a plain,
231+
non-attributed copy node, or pass ``options=None``. On an older
232+
cuda.bindings/driver, ``src_access_order`` values of ``STREAM``
233+
and ``ANY`` are silently ignored; ``DURING_API_CALL`` raises
234+
instead of silently downgrading its guarantee.
201235
236+
Raises
237+
------
238+
TypeError
239+
If ``options`` is not a :class:`~utils.CopyOptions` instance, or
240+
if ``options`` is given together with ``LEGACY_DEFAULT_STREAM``
241+
or a stream currently in graph capture mode.
242+
RuntimeError
243+
If ``options.src_access_order`` is ``DURING_API_CALL`` and
244+
cuda.bindings or the driver is older than CUDA 13.2: falling
245+
back to a plain copy cannot honor that guarantee.
202246
"""
203247

204248
def fill(self, value: int | BufferProtocol, *, stream: Stream | GraphBuilder) -> None:

‎cuda_core/cuda/core/_memory/_buffer.pyx‎

Lines changed: 128 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,20 @@ from cuda.core._resource_handles cimport (
2727
)
2828
from cuda.core.typing import DevicePointerType
2929

30-
from cuda.core._stream cimport Stream, Stream_accept, default_stream
30+
from cuda.core._memory._copy_attributes cimport _with_attributes_available
31+
from cuda.core._memory._copy_attributes cimport _to_cu_memcpy_attributes # no-cython-lint
32+
33+
IF CUDA_CORE_BUILD_MAJOR >= 13:
34+
from cuda.core._resource_handles cimport memcpy_with_attributes_async
35+
36+
from cuda.core._stream cimport Stream, Stream_accept, Stream_is_legacy_default_token, default_stream
3137
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value
3238

3339
import sys
3440
from collections.abc import Sequence
3541
from typing import TYPE_CHECKING
3642

43+
from cuda.core._memory._copy_enums import CopyOptions, _reject_unsupported_during_api_call
3744
from cuda.core._utils.pycompat import BufferProtocol
3845
from cuda.core._dlpack import classify_dl_device, make_py_capsule
3946
from cuda.core._device import Device
@@ -160,6 +167,75 @@ cdef inline void _init_memory_attrs(Buffer self):
160167
self._mem_attrs_inited.store(True, memory_order_release)
161168

162169

170+
cdef bint _stream_is_capturing(Stream s):
171+
cdef cydriver.CUstreamCaptureStatus cap_status
172+
IF CUDA_CORE_BUILD_MAJOR >= 13:
173+
HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &cap_status,
174+
NULL, NULL, NULL, NULL, NULL))
175+
ELSE:
176+
HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &cap_status,
177+
NULL, NULL, NULL, NULL))
178+
return cap_status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE
179+
180+
181+
cdef void _do_copy_with_attributes(
182+
cydriver.CUdeviceptr dst, cydriver.CUdeviceptr src, size_t nbytes,
183+
object options, cydriver.CUstream hstream,
184+
):
185+
IF CUDA_CORE_BUILD_MAJOR >= 13:
186+
# Routed through the memcpy_with_attributes_async() C++ shim since
187+
# cydriver.cuMemcpyWithAttributesAsync is absent from cuda-bindings < 13.2.
188+
cdef cydriver.CUmemcpyAttributes cu_attr = _to_cu_memcpy_attributes(options)
189+
with nogil:
190+
HANDLE_RETURN(memcpy_with_attributes_async(dst, src, nbytes, <void*>&cu_attr, hstream))
191+
ELSE:
192+
pass # unreachable: _with_attributes_available() is always False on CUDA 12
193+
194+
195+
cdef void _dispatch_buffer_copy(
196+
cydriver.CUdeviceptr dst, cydriver.CUdeviceptr src, size_t nbytes,
197+
Stream s, object options, str method_name,
198+
):
199+
"""Submit a single copy, honoring CopyOptions when the attributes path is usable."""
200+
if options is None:
201+
with nogil:
202+
HANDLE_RETURN(cydriver.cuMemcpyAsync(dst, src, nbytes, as_cu(s._h_stream)))
203+
return
204+
if not isinstance(options, CopyOptions):
205+
raise TypeError(
206+
f"{method_name}: options must be CopyOptions, got {type(options).__name__}"
207+
)
208+
if Stream_is_legacy_default_token(s):
209+
raise TypeError(
210+
f"{method_name} does not accept LEGACY_DEFAULT_STREAM with options "
211+
"(matches copy_batch); cuMemcpyWithAttributesAsync rejects it outright, "
212+
"unlike PER_THREAD_DEFAULT_STREAM, which is a real stream to the driver "
213+
"and is accepted. Pass an explicit stream, PER_THREAD_DEFAULT_STREAM, "
214+
"or options=None."
215+
)
216+
if _stream_is_capturing(s):
217+
raise TypeError(
218+
f"{method_name} does not support graph capture with options "
219+
"(matches copy_batch); the driver has no graph-node form of "
220+
"cuMemcpyWithAttributesAsync, so options cannot be honored in a graph. "
221+
"Use GraphNode.memcpy for a plain (non-attributed) copy node, or pass "
222+
"options=None."
223+
)
224+
if _with_attributes_available():
225+
_do_copy_with_attributes(dst, src, nbytes, options, as_cu(s._h_stream))
226+
else:
227+
_reject_unsupported_during_api_call(
228+
options.src_access_order,
229+
"cuda.bindings and the driver to both report CUDA 13.2 or newer "
230+
"(cuMemcpyWithAttributesAsync is unavailable here)",
231+
)
232+
# STREAM and ANY never require access sooner than stream order, so
233+
# cuMemcpyAsync satisfies them; options are otherwise silently
234+
# ignored on this pre-CUDA-13.2 fallback path, matching copy_batch.
235+
with nogil:
236+
HANDLE_RETURN(cydriver.cuMemcpyAsync(dst, src, nbytes, as_cu(s._h_stream)))
237+
238+
163239
cdef class Buffer:
164240
"""Represent a handle to allocated memory.
165241
@@ -396,7 +472,8 @@ cdef class Buffer:
396472
self.close()
397473
return False
398474

399-
def copy_to(self, dst: Buffer | None = None, *, stream: Stream | GraphBuilder) -> Buffer:
475+
def copy_to(self, dst: Buffer | None = None, *, stream: Stream | GraphBuilder,
476+
options: CopyOptions | None = None) -> Buffer:
400477
"""Copy from this buffer to the dst buffer asynchronously on the given stream.
401478

402479
Copies the data from this buffer to the provided dst buffer.
@@ -411,6 +488,28 @@ cdef class Buffer:
411488
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
412489
Keyword argument specifying the stream for the
413490
asynchronous copy
491+
options : :class:`~utils.CopyOptions`, optional
492+
Transfer hints (source access order, location hints, overlap mode).
493+
Honored when cuda.bindings and the driver are both CUDA 13.2 or
494+
newer. Not accepted with ``LEGACY_DEFAULT_STREAM``; use
495+
``PER_THREAD_DEFAULT_STREAM`` instead. Not accepted with a
496+
capturing stream either, since a graph cannot represent these
497+
attributes; use :meth:`graph.GraphNode.memcpy` for a plain,
498+
non-attributed copy node, or pass ``options=None``. On an older
499+
cuda.bindings/driver, ``src_access_order`` values of ``STREAM``
500+
and ``ANY`` are silently ignored; ``DURING_API_CALL`` raises
501+
instead of silently downgrading its guarantee.
502+
503+
Raises
504+
------
505+
TypeError
506+
If ``options`` is not a :class:`~utils.CopyOptions` instance, or
507+
if ``options`` is given together with ``LEGACY_DEFAULT_STREAM``
508+
or a stream currently in graph capture mode.
509+
RuntimeError
510+
If ``options.src_access_order`` is ``DURING_API_CALL`` and
511+
cuda.bindings or the driver is older than CUDA 13.2: falling
512+
back to a plain copy cannot honor that guarantee.
414513

415514
"""
416515
Buffer_check_open(self)
@@ -430,12 +529,12 @@ cdef class Buffer:
430529
raise ValueError( "buffer sizes mismatch between src and dst (sizes "
431530
f"are: src={src_size}, dst={dst_size})"
432531
)
433-
with nogil:
434-
HANDLE_RETURN(cydriver.cuMemcpyAsync(
435-
as_cu(dst._h_ptr), as_cu(self._h_ptr), src_size, as_cu(s._h_stream)))
532+
_dispatch_buffer_copy(
533+
as_cu(dst._h_ptr), as_cu(self._h_ptr), src_size, s, options, "copy_to")
436534
return dst
437535

438-
def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder) -> None:
536+
def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder,
537+
options: CopyOptions | None = None) -> None:
439538
"""Copy from the src buffer to this buffer asynchronously on the given stream.
440539

441540
Parameters
@@ -445,7 +544,28 @@ cdef class Buffer:
445544
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
446545
Keyword argument specifying the stream for the
447546
asynchronous copy
547+
options : :class:`~utils.CopyOptions`, optional
548+
Transfer hints (source access order, location hints, overlap mode).
549+
Honored when cuda.bindings and the driver are both CUDA 13.2 or
550+
newer. Not accepted with ``LEGACY_DEFAULT_STREAM``; use
551+
``PER_THREAD_DEFAULT_STREAM`` instead. Not accepted with a
552+
capturing stream either, since a graph cannot represent these
553+
attributes; use :meth:`graph.GraphNode.memcpy` for a plain,
554+
non-attributed copy node, or pass ``options=None``. On an older
555+
cuda.bindings/driver, ``src_access_order`` values of ``STREAM``
556+
and ``ANY`` are silently ignored; ``DURING_API_CALL`` raises
557+
instead of silently downgrading its guarantee.
448558

559+
Raises
560+
------
561+
TypeError
562+
If ``options`` is not a :class:`~utils.CopyOptions` instance, or
563+
if ``options`` is given together with ``LEGACY_DEFAULT_STREAM``
564+
or a stream currently in graph capture mode.
565+
RuntimeError
566+
If ``options.src_access_order`` is ``DURING_API_CALL`` and
567+
cuda.bindings or the driver is older than CUDA 13.2: falling
568+
back to a plain copy cannot honor that guarantee.
449569
"""
450570
Buffer_check_open(self)
451571
Buffer_check_open(src)
@@ -457,9 +577,8 @@ cdef class Buffer:
457577
raise ValueError( "buffer sizes mismatch between src and dst (sizes "
458578
f"are: src={src_size}, dst={dst_size})"
459579
)
460-
with nogil:
461-
HANDLE_RETURN(cydriver.cuMemcpyAsync(
462-
as_cu(self._h_ptr), as_cu(src._h_ptr), dst_size, as_cu(s._h_stream)))
580+
_dispatch_buffer_copy(
581+
as_cu(self._h_ptr), as_cu(src._h_ptr), dst_size, s, options, "copy_from")
463582

464583
def fill(self, value: int | BufferProtocol, *, stream: Stream | GraphBuilder) -> None:
465584
"""Fill this buffer with a repeating byte pattern.

‎cuda_core/cuda/core/_memory/_copy_attributes.pxd‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,19 @@ from cuda.core._utils.version cimport cy_binding_version, cy_driver_version # n
1111

1212

1313
IF CUDA_CORE_BUILD_MAJOR >= 13:
14+
from cuda.core._resource_handles cimport has_memcpy_with_attributes_async
15+
1416
cdef inline bint _with_attributes_available():
15-
return cy_driver_version() >= (13, 2, 0) and cy_binding_version() >= (13, 2, 0)
17+
# has_memcpy_with_attributes_async() says whether the installed
18+
# cuda-bindings actually exports cuMemcpyWithAttributesAsync (13.2+);
19+
# the version checks alone are not sufficient, since cuda.core's build
20+
# can be paired with a cuda-bindings install older than what it built
21+
# against (see https://github.com/NVIDIA/cuda-python/issues/2063).
22+
return (
23+
has_memcpy_with_attributes_async()
24+
and cy_driver_version() >= (13, 2, 0)
25+
and cy_binding_version() >= (13, 2, 0)
26+
)
1627
ELSE:
1728
cdef inline bint _with_attributes_available():
1829
return False

0 commit comments

Comments
 (0)