@@ -27,13 +27,17 @@ from cuda.core._resource_handles cimport (
2727)
2828from 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
3133from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value
3234
3335import sys
36+ import warnings
3437from collections.abc import Sequence
3538from typing import TYPE_CHECKING
3639
40+ from cuda.core._memory._copy_enums import CopyOptions
3741from cuda.core._utils.pycompat import BufferProtocol
3842from cuda.core._dlpack import classify_dl_device, make_py_capsule
3943from 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+
162189cdef 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.
0 commit comments