Coverage for cuda/core/_memory/_copy_ops.pyx: 84.09%
88 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 02:41 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5from __future__ import annotations
7from collections.abc import Sequence
9IF CUDA_CORE_BUILD_MAJOR >= 13:
10 from libcpp.vector cimport vector
12from cuda.bindings cimport cydriver
13from cuda.core._memory._buffer cimport Buffer, Buffer_coerce_batch
14from cuda.core._memory._copy_attributes cimport _to_cu_memcpy_attributes # no-cython-lint
15from cuda.core._resource_handles cimport as_cu
16from cuda.core._stream cimport Stream, Stream_accept, Stream_is_legacy_default_token
17from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
19# cy_driver_version and _attr_run_starts are referenced only from CUDA 13
20# branches. cython-lint does not evaluate compile-time IF blocks, so they need
21# a pragma to be seen as used.
22from cuda.core._utils.version cimport cy_driver_version # no-cython-lint
24from cuda.core._memory._copy_enums import (
25 CopyOptions,
26 _attr_run_starts, # no-cython-lint
27 _reject_unsupported_during_api_call,
28)
30_SINGLE_COPY_HINT = "Buffer.copy_to / Buffer.copy_from"
33cdef inline bint _batch_entry_point_available():
34 """Whether cuMemcpyBatchAsync can actually be called here.
36 Requires ``cuda.core`` built against CUDA 13 headers (compile time) and
37 a driver reporting CUDA 13.0 or newer, i.e.
38 ``cuDriverGetVersion() >= 13000`` (run time).
40 The run-time bound is set by the binding layer, not by when the driver
41 gained the feature. CUDA 12.8 already exposed a ``cuMemcpyBatchAsync``,
42 but its signature carried a ``failIdx`` out-parameter that CUDA 13.0
43 dropped. ``cuda.bindings`` resolves only the 13.0 revision, via
44 ``cuGetProcAddress_v2('cuMemcpyBatchAsync', ..., 13000, ...)``, so an
45 older driver yields a NULL pointer even though it may implement the
46 earlier entry point.
47 """
48 IF CUDA_CORE_BUILD_MAJOR >= 13:
49 return cy_driver_version() >= (13, 0, 0) 1klmnopqrstbcdefghaij
50 ELSE:
51 return False
54def _normalize_copy_options(
55 options: CopyOptions | Sequence[CopyOptions] | None,
56 Py_ssize_t n,
57) -> tuple[CopyOptions, ...]:
58 """Expand ``options`` to exactly one :class:`CopyOptions` per copy.
60 ``None`` and a scalar broadcast; a sequence pairs by index and must
61 already have length ``n``.
63 Internal, but deliberately importable: options are hints that change
64 how the driver stages a transfer and never the bytes it produces, so
65 this expansion (and the run encoding applied to it) is the only
66 observable evidence that a scalar reached every copy.
67 """
68 if options is None: 1klmnopqrstbcdefghaijvuwFBC
69 return (CopyOptions(),) * n 1klmnopqrstF
70 if isinstance(options, CopyOptions): 1bcdefghaijvuwBC
71 return (options,) * n 1bcdefghijB
72 if isinstance(options, Sequence): 1avuwBC
73 if len(options) != n: 1avuBC
74 raise ValueError( 1vC
75 f"copy_batch: options length {len(options)} does not match " 1vC
76 f"buffers length {n}" 1vC
77 )
78 for a in options: 1zauB
79 if not isinstance(a, CopyOptions): 1auB
80 raise TypeError( 1u
81 f"copy_batch: each options element must be CopyOptions, " 1u
82 f"got {type(a).__name__}" 1u
83 )
84 return tuple(options) 1aB
85 raise TypeError( 1w
86 f"copy_batch: options must be CopyOptions or a sequence of " 1w
87 f"CopyOptions, got {type(options).__name__}" 1w
88 )
91def copy_batch(
92 stream: Stream,
93 srcs: Sequence[Buffer],
94 dsts: Sequence[Buffer],
95 *,
96 options: CopyOptions | Sequence[CopyOptions] | None = None,
97) -> None:
98 """Copy a batch of buffers asynchronously.
100 Source buffer and destination buffer sizes must match. For a single
101 buffer, use :meth:`Buffer.copy_to` or :meth:`Buffer.copy_from`.
103 The driver provides no graph-node form of ``cuMemcpyBatchAsync``, so
104 this cannot be captured into a graph. Both passing a
105 :class:`~graph.GraphBuilder` and passing its underlying
106 :attr:`~graph.GraphBuilder.stream` while capture is active are
107 rejected. Build graph copies with
108 :meth:`graph.GraphNode.memcpy` or per-buffer :meth:`Buffer.copy_to`.
110 Parameters
111 ----------
112 stream : :class:`~_stream.Stream`
113 Stream for the asynchronous copy. First positional and required
114 (mirrors :func:`launch`). Does not accept a capturing stream
115 (including a :class:`~graph.GraphBuilder`\'s underlying stream); use
116 :meth:`graph.GraphNode.memcpy` or per-buffer
117 :meth:`Buffer.copy_to` to build copies into a graph. Does not accept
118 ``LEGACY_DEFAULT_STREAM``, which ``cuMemcpyBatchAsync`` rejects
119 outright; ``PER_THREAD_DEFAULT_STREAM`` is a real stream to the
120 driver and is accepted.
121 srcs : Sequence[:class:`Buffer`]
122 Source buffers. Must be a sequence, not a single Buffer.
123 dsts : Sequence[:class:`Buffer`]
124 Destination buffers. Must match ``len(srcs)``.
125 options : :class:`CopyOptions` | Sequence[:class:`CopyOptions`] | None
126 Per-copy options. A single value applies to every copy; a
127 sequence pairs by index and must match ``len(srcs)``. ``None``
128 uses stream-ordered defaults.
130 Raises
131 ------
132 ValueError
133 If lengths or sizes mismatch.
134 TypeError
135 If a single Buffer is passed instead of a sequence, if
136 ``LEGACY_DEFAULT_STREAM`` is passed, or if the stream is currently
137 in graph capture mode.
138 RuntimeError
139 If any copy requests ``src_access_order=DURING_API_CALL`` and the
140 native ``cuMemcpyBatchAsync`` path is unavailable (see Notes): the
141 per-copy ``cuMemcpyAsync`` fallback reads the source in stream
142 order only, which cannot honor that guarantee.
144 Notes
145 -----
146 Batching through ``cuMemcpyBatchAsync`` requires all three of:
147 ``cuda.core`` built against CUDA 13 headers, ``cuda.bindings`` 13.0 or
148 newer, and a driver reporting CUDA 13.0 or newer
149 (``cuDriverGetVersion() >= 13000``). ``cuda.bindings`` binds only the
150 CUDA 13.0 revision of the entry point, so a driver that predates it is
151 refused even where it implements the earlier CUDA 12.8 signature.
153 The driver may execute batch items concurrently and in any order.
154 A batch must therefore not contain copies where the source range of
155 one copy overlaps the destination range of another; such aliasing
156 produces undefined results. Detecting overlaps at runtime is
157 impractical; callers are responsible for ensuring no aliasing exists.
159 On pre-CUDA 13 installs the copies fall back to a Python-level loop
160 over ``cuMemcpyAsync``, so the potential performance benefit of
161 asynchronous batched copies is not realized. ``src_access_order`` values
162 of ``STREAM`` and ``ANY`` are silently ignored on the fallback path
163 (stream-ordered access already satisfies both); ``DURING_API_CALL``
164 raises ``RuntimeError`` instead, since silently downgrading it to
165 stream-ordered access would let a caller reuse the source buffer before
166 the real read happens.
168 """
169 cdef tuple src_bufs = Buffer_coerce_batch(srcs, "copy_batch", _SINGLE_COPY_HINT) 1klmnopqrADstbcdefghaijEvuwGHIJxy
170 cdef tuple dst_bufs = Buffer_coerce_batch(dsts, "copy_batch", _SINGLE_COPY_HINT) 1klmnopqrADstbcdefghaijEvuwGHIJxy
171 cdef Py_ssize_t n = len(src_bufs) 1klmnopqrADstbcdefghaijEvuwxy
173 if len(dst_bufs) != n: 1klmnopqrADstbcdefghaijEvuwxy
174 raise ValueError( 1E
175 f"copy_batch: srcs length {n} does not match dsts length {len(dst_bufs)}" 1E
176 )
178 cdef Stream s = Stream_accept(stream) 1klmnopqrADstbcdefghaijvuwxy
180 if Stream_is_legacy_default_token(s): 1klmnopqrADstbcdefghaijvuwxy
181 raise TypeError( 1D
182 "copy_batch does not accept LEGACY_DEFAULT_STREAM; cuMemcpyBatchAsync "
183 "rejects it outright, unlike PER_THREAD_DEFAULT_STREAM, which is a real "
184 "stream to the driver and is accepted. Pass an explicit stream or "
185 "PER_THREAD_DEFAULT_STREAM."
186 )
188 cdef cydriver.CUstreamCaptureStatus _cap_status
189 IF CUDA_CORE_BUILD_MAJOR >= 13:
190 HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &_cap_status, 1klmnopqrAstbcdefghaijvuwxy
191 NULL, NULL, NULL, NULL, NULL))
192 ELSE:
193 HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &_cap_status,
194 NULL, NULL, NULL, NULL))
195 if _cap_status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE: 1klmnopqrAstbcdefghaijvuwxy
196 raise TypeError( 1A
197 "copy_batch does not support graph capture; "
198 "use GraphNode.memcpy or per-buffer Buffer.copy_to instead"
199 )
201 cdef Buffer src_buf
202 cdef Buffer dst_buf
203 cdef Py_ssize_t i
205 for i in range(n): 1klmnopqrstbcdefghaijvuwxy
206 src_buf = <Buffer>src_bufs[i] 1klmnopqrstbcdefghaijvuwxy
207 dst_buf = <Buffer>dst_bufs[i] 1klmnopqrstbcdefghaijvuwxy
208 if src_buf.size != dst_buf.size: 1klmnopqrstbcdefghaijvuwxy
209 raise ValueError( 1xy
210 f"copy_batch: buffer size mismatch at index {i} " 1xy
211 f"(src={src_buf.size}, dst={dst_buf.size})" 1xy
212 )
214 cdef tuple attr_tuple = _normalize_copy_options(options, n) 1klmnopqrstbcdefghaijvuw
216 _do_copy_batch(src_bufs, dst_bufs, s, attr_tuple) 1klmnopqrstbcdefghaij
219cdef void _do_copy_batch(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tuple):
220 IF CUDA_CORE_BUILD_MAJOR >= 13:
221 # Building against CUDA 13 headers says nothing about the installed
222 # driver, so the run-time version still has to be checked before
223 # calling a 13.0-only entry point (see PRs #2054 / #2064).
224 if _batch_entry_point_available(): 1klmnopqrstbcdefghaij
225 _do_copy_batch_native(src_bufs, dst_bufs, s, attr_tuple) 1klmnopqrstbcdefghaij
226 else:
227 _reject_during_api_call_fallback(attr_tuple)
228 _do_copy_batch_loop(src_bufs, dst_bufs, s)
229 ELSE:
230 _reject_during_api_call_fallback(attr_tuple)
231 _do_copy_batch_loop(src_bufs, dst_bufs, s)
234cdef void _reject_during_api_call_fallback(tuple attr_tuple):
235 """Raise before the per-copy cuMemcpyAsync loop if any copy needs
236 DURING_API_CALL, which that fallback cannot honor (see
237 _reject_unsupported_during_api_call for why this must raise rather than
238 silently ignore the option, unlike STREAM and ANY).
239 """
240 cdef Py_ssize_t i
241 for i in range(len(attr_tuple)):
242 _reject_unsupported_during_api_call(
243 (<object>attr_tuple[i]).src_access_order,
244 "cuda.core built against CUDA 13 headers and cuda.bindings/driver "
245 "13.0 or newer (cuMemcpyBatchAsync is unavailable here)",
246 index=i,
247 )
250cdef void _do_copy_batch_loop(tuple src_bufs, tuple dst_bufs, Stream s):
251 """Per-copy cuMemcpyAsync fallback where the batch entry point is absent.
253 Issues copies one at a time, so the performance benefit of batching is
254 not realized. STREAM and ANY are silently ignored here (satisfied by
255 stream-ordered cuMemcpyAsync regardless); DURING_API_CALL is rejected by
256 _reject_during_api_call_fallback before this is ever called.
257 """
258 cdef Py_ssize_t n = len(src_bufs)
259 cdef Py_ssize_t i
260 cdef Buffer src_buf
261 cdef Buffer dst_buf
262 cdef size_t nbytes
263 cdef cydriver.CUstream hstream = as_cu(s._h_stream)
265 for i in range(n):
266 src_buf = <Buffer>src_bufs[i]
267 dst_buf = <Buffer>dst_bufs[i]
268 nbytes = src_buf._size
269 with nogil:
270 HANDLE_RETURN(cydriver.cuMemcpyAsync(
271 as_cu(dst_buf._h_ptr), as_cu(src_buf._h_ptr), nbytes, hstream))
274IF CUDA_CORE_BUILD_MAJOR >= 13:
275 cdef void _do_copy_batch_native(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tuple):
276 cdef Py_ssize_t n = len(src_bufs) 1klmnopqrstbcdefghaij
277 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 1klmnopqrstbcdefghaij
278 cdef vector[cydriver.CUdeviceptr] dst_ptrs
279 cdef vector[cydriver.CUdeviceptr] src_ptrs
280 cdef vector[size_t] sizes
281 cdef vector[size_t] attrs_idxs
282 dst_ptrs.resize(n) 1klmnopqrstbcdefghaij
283 src_ptrs.resize(n) 1klmnopqrstbcdefghaij
284 sizes.resize(n) 1klmnopqrstbcdefghaij
286 cdef Buffer src_buf
287 cdef Buffer dst_buf
288 cdef Py_ssize_t i
290 # Collapse equal neighbouring attributes into runs so a broadcast
291 # attribute reaches the driver once (numAttrs == 1) instead of being
292 # repeated per copy. attrs[k] applies to [attrsIdxs[k], attrsIdxs[k+1]).
293 cdef list run_starts = _attr_run_starts(attr_tuple) 1klmnopqrstbcdefghaij
294 cdef vector[cydriver.CUmemcpyAttributes] cu_attrs
295 cdef size_t num_attrs = <size_t>len(run_starts) 1klmnopqrstbcdefghaij
296 cu_attrs.reserve(num_attrs) 1klmnopqrstbcdefghaij
297 attrs_idxs.reserve(num_attrs) 1klmnopqrstbcdefghaij
298 for i in run_starts: 1klmnopqrstbcdefghaij
299 cu_attrs.push_back(_to_cu_memcpy_attributes(attr_tuple[i])) 1klmnopqrstbcdefghaij
300 attrs_idxs.push_back(<size_t>i) 1klmnopqrstbcdefghaij
302 for i in range(n): 1klmnopqrstbcdefghaij
303 src_buf = <Buffer>src_bufs[i] 1klmnopqrstbcdefghaij
304 dst_buf = <Buffer>dst_bufs[i] 1klmnopqrstbcdefghaij
305 src_ptrs[i] = as_cu(src_buf._h_ptr) 1klmnopqrstbcdefghaij
306 dst_ptrs[i] = as_cu(dst_buf._h_ptr) 1klmnopqrstbcdefghaij
307 sizes[i] = src_buf.size 1klmnopqrstbcdefghaij
309 with nogil: 1klmnopqrstbcdefghaij
310 HANDLE_RETURN(cydriver.cuMemcpyBatchAsync( 1klmnopqrstbcdefghaij
311 dst_ptrs.data(),
312 src_ptrs.data(),
313 sizes.data(),
314 <size_t>n,
315 cu_attrs.data(),
316 attrs_idxs.data(),
317 num_attrs,
318 hstream,
319 ))