Skip to content

Commit d835ecf

Browse files
committed
Ocuda.core: introduce copy options for Buffer.copy_{to/from}
1 parent 0eb23c4 commit d835ecf

4 files changed

Lines changed: 483 additions & 11 deletions

File tree

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

Lines changed: 15 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
@@ -135,7 +136,7 @@ class Buffer:
135136
def __exit__(self, exc_type, exc_val, exc_tb):
136137
...
137138

138-
def copy_to(self, dst: Buffer | None=None, *, stream: Stream | GraphBuilder) -> Buffer:
139+
def copy_to(self, dst: Buffer | None=None, *, stream: Stream | GraphBuilder, options: CopyOptions | None=None) -> Buffer:
139140
"""Copy from this buffer to the dst buffer asynchronously on the given stream.
140141
141142
Copies the data from this buffer to the provided dst buffer.
@@ -150,10 +151,16 @@ class Buffer:
150151
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
151152
Keyword argument specifying the stream for the
152153
asynchronous copy
154+
options : :class:`~utils.CopyOptions`, optional
155+
Transfer hints (source access order, location hints, overlap mode).
156+
Honored only when both cuda.bindings and the driver are CUDA 13.2+
157+
and the stream is not under graph capture; otherwise a
158+
:class:`UserWarning` is emitted and the copy falls back to
159+
``cuMemcpyAsync``.
153160
154161
"""
155162

156-
def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder) -> None:
163+
def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder, options: CopyOptions | None=None) -> None:
157164
"""Copy from the src buffer to this buffer asynchronously on the given stream.
158165
159166
Parameters
@@ -163,6 +170,12 @@ class Buffer:
163170
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
164171
Keyword argument specifying the stream for the
165172
asynchronous copy
173+
options : :class:`~utils.CopyOptions`, optional
174+
Transfer hints (source access order, location hints, overlap mode).
175+
Honored only when both cuda.bindings and the driver are CUDA 13.2+
176+
and the stream is not under graph capture; otherwise a
177+
:class:`UserWarning` is emitted and the copy falls back to
178+
``cuMemcpyAsync``.
166179
167180
"""
168181

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

Lines changed: 78 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,17 @@ 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+
from cuda.core._stream cimport Stream, Stream_accept, Stream_is_default_token, default_stream
3133
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value
3234

3335
import sys
36+
import warnings
3437
from collections.abc import Sequence
3538
from typing import TYPE_CHECKING
3639

40+
from cuda.core._memory._copy_enums import CopyOptions
3741
from cuda.core._utils.pycompat import BufferProtocol
3842
from cuda.core._dlpack import classify_dl_device, make_py_capsule
3943
from cuda.core._device import Device
@@ -159,6 +163,29 @@ cdef inline void _init_memory_attrs(Buffer self):
159163
self._mem_attrs_inited.store(True, memory_order_release)
160164

161165

166+
cdef bint _stream_is_capturing(Stream s):
167+
cdef cydriver.CUstreamCaptureStatus cap_status
168+
IF CUDA_CORE_BUILD_MAJOR >= 13:
169+
HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &cap_status,
170+
NULL, NULL, NULL, NULL, NULL))
171+
ELSE:
172+
HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &cap_status,
173+
NULL, NULL, NULL, NULL))
174+
return cap_status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE
175+
176+
177+
cdef void _do_copy_with_attributes(
178+
cydriver.CUdeviceptr dst, cydriver.CUdeviceptr src, size_t nbytes,
179+
object options, cydriver.CUstream hstream,
180+
):
181+
IF CUDA_CORE_BUILD_MAJOR >= 13:
182+
cdef cydriver.CUmemcpyAttributes cu_attr = _to_cu_memcpy_attributes(options)
183+
with nogil:
184+
HANDLE_RETURN(cydriver.cuMemcpyWithAttributesAsync(dst, src, nbytes, &cu_attr, hstream))
185+
ELSE:
186+
pass # unreachable: _with_attributes_available() is always False on CUDA 12
187+
188+
162189
cdef class Buffer:
163190
"""Represent a handle to allocated memory.
164191
@@ -357,7 +384,8 @@ cdef class Buffer:
357384
self.close()
358385
return False
359386

360-
def copy_to(self, dst: Buffer | None = None, *, stream: Stream | GraphBuilder) -> Buffer:
387+
def copy_to(self, dst: Buffer | None = None, *, stream: Stream | GraphBuilder,
388+
options: CopyOptions | None = None) -> Buffer:
361389
"""Copy from this buffer to the dst buffer asynchronously on the given stream.
362390

363391
Copies the data from this buffer to the provided dst buffer.
@@ -372,6 +400,12 @@ cdef class Buffer:
372400
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
373401
Keyword argument specifying the stream for the
374402
asynchronous copy
403+
options : :class:`~utils.CopyOptions`, optional
404+
Transfer hints (source access order, location hints, overlap mode).
405+
Honored only when both cuda.bindings and the driver are CUDA 13.2+
406+
and the stream is not under graph capture; otherwise a
407+
:class:`UserWarning` is emitted and the copy falls back to
408+
``cuMemcpyAsync``.
375409

376410
"""
377411
cdef Stream s = Stream_accept(stream)
@@ -388,12 +422,27 @@ cdef class Buffer:
388422
raise ValueError( "buffer sizes mismatch between src and dst (sizes "
389423
f"are: src={src_size}, dst={dst_size})"
390424
)
391-
with nogil:
392-
HANDLE_RETURN(cydriver.cuMemcpyAsync(
393-
as_cu(dst._h_ptr), as_cu(self._h_ptr), src_size, as_cu(s._h_stream)))
425+
if options is None:
426+
with nogil:
427+
HANDLE_RETURN(cydriver.cuMemcpyAsync(
428+
as_cu(dst._h_ptr), as_cu(self._h_ptr), src_size, as_cu(s._h_stream)))
429+
elif _with_attributes_available() and not Stream_is_default_token(s) and not _stream_is_capturing(s):
430+
_do_copy_with_attributes(
431+
as_cu(dst._h_ptr), as_cu(self._h_ptr), src_size, options, as_cu(s._h_stream))
432+
else:
433+
warnings.warn(
434+
"copy_to: CopyOptions are not honored (requires CUDA 13.2+ driver and "
435+
"cuda.bindings, and a non-capturing, non-default stream); falling back to cuMemcpyAsync",
436+
UserWarning,
437+
stacklevel=2,
438+
)
439+
with nogil:
440+
HANDLE_RETURN(cydriver.cuMemcpyAsync(
441+
as_cu(dst._h_ptr), as_cu(self._h_ptr), src_size, as_cu(s._h_stream)))
394442
return dst
395443

396-
def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder) -> None:
444+
def copy_from(self, src: Buffer, *, stream: Stream | GraphBuilder,
445+
options: CopyOptions | None = None) -> None:
397446
"""Copy from the src buffer to this buffer asynchronously on the given stream.
398447

399448
Parameters
@@ -403,6 +452,12 @@ cdef class Buffer:
403452
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
404453
Keyword argument specifying the stream for the
405454
asynchronous copy
455+
options : :class:`~utils.CopyOptions`, optional
456+
Transfer hints (source access order, location hints, overlap mode).
457+
Honored only when both cuda.bindings and the driver are CUDA 13.2+
458+
and the stream is not under graph capture; otherwise a
459+
:class:`UserWarning` is emitted and the copy falls back to
460+
``cuMemcpyAsync``.
406461

407462
"""
408463
cdef Stream s = Stream_accept(stream)
@@ -413,9 +468,23 @@ cdef class Buffer:
413468
raise ValueError( "buffer sizes mismatch between src and dst (sizes "
414469
f"are: src={src_size}, dst={dst_size})"
415470
)
416-
with nogil:
417-
HANDLE_RETURN(cydriver.cuMemcpyAsync(
418-
as_cu(self._h_ptr), as_cu(src._h_ptr), dst_size, as_cu(s._h_stream)))
471+
if options is None:
472+
with nogil:
473+
HANDLE_RETURN(cydriver.cuMemcpyAsync(
474+
as_cu(self._h_ptr), as_cu(src._h_ptr), dst_size, as_cu(s._h_stream)))
475+
elif _with_attributes_available() and not Stream_is_default_token(s) and not _stream_is_capturing(s):
476+
_do_copy_with_attributes(
477+
as_cu(self._h_ptr), as_cu(src._h_ptr), dst_size, options, as_cu(s._h_stream))
478+
else:
479+
warnings.warn(
480+
"copy_from: CopyOptions are not honored (requires CUDA 13.2+ driver and "
481+
"cuda.bindings, and a non-capturing, non-default stream); falling back to cuMemcpyAsync",
482+
UserWarning,
483+
stacklevel=2,
484+
)
485+
with nogil:
486+
HANDLE_RETURN(cydriver.cuMemcpyAsync(
487+
as_cu(self._h_ptr), as_cu(src._h_ptr), dst_size, as_cu(s._h_stream)))
419488

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

‎cuda_core/docs/source/release/1.2.0-notes.rst‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,16 @@ New features
1818
are rejected. Copies within a batch must not alias.
1919
(`#1333 <https://github.com/NVIDIA/cuda-python/issues/1333>`__)
2020

21+
- :meth:`Buffer.copy_to` and :meth:`Buffer.copy_from` now accept an optional
22+
``options`` keyword argument (:class:`~utils.CopyOptions`). When both
23+
``cuda.bindings`` and the driver are CUDA 13.2 or newer and the stream is
24+
not under graph capture, the copy is submitted via
25+
``cuMemcpyWithAttributesAsync``. On older installs, or when the stream is
26+
capturing, a :class:`UserWarning` is emitted and the copy falls back to
27+
``cuMemcpyAsync``; ``options=None`` (the default) always uses
28+
``cuMemcpyAsync`` with no warning.
29+
(`#2365 <https://github.com/NVIDIA/cuda-python/issues/2365>`__)
30+
2131
Fixes and enhancements
2232
----------------------
2333

0 commit comments

Comments
 (0)