Coverage for cuda/core/_memory/_copy_ops.pyx: 70.45%

1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from __future__ import annotations 

6  

7from collections.abc import Sequence 

8  

9IF CUDA_CORE_BUILD_MAJOR >= 13: 

10 from libcpp.vector cimport vector 

11  

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._rt 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 

18  

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 

23  

24from cuda.core._memory._copy_enums import ( 

25 CopyOptions, 

26 _attr_run_starts, # no-cython-lint 

27 _reject_unsupported_during_api_call, 

28) 

29  

30_SINGLE_COPY_HINT = "Buffer.copy_to / Buffer.copy_from" 

31  

32  

33cdef inline bint _batch_entry_point_available(): 

34 """Whether cuMemcpyBatchAsync can actually be called here. 

35  

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). 

39  

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) 20 ctx1klmnopqrstbcdefghaij

50 ELSE: 

51 return False 

52  

53  

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. 

59  

60 ``None`` and a scalar broadcast; a sequence pairs by index and must 

61 already have length ``n``. 

62  

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: 26 ctx1klmnopqrstbcdefghaijvuwEAB

69 return (CopyOptions(),) * n 11 ctx1klmnopqrstE

70 if isinstance(options, CopyOptions): 15 ctx1bcdefghaijvuwAB

71 return (options,) * n 10 ctx1bcdefghijA

72 if isinstance(options, Sequence): 6 ctx1avuwAB

73 if len(options) != n: 5 ctx1avuAB

74 raise ValueError( 2 ctx1vB

75 f"copy_batch: options length {len(options)} does not match " 2 ctx1vB

76 f"buffers length {n}" 2 ctx1vB

77 ) 

78 for a in options: 3 ctx1auA

79 if not isinstance(a, CopyOptions): 3 ctx1auA

80 raise TypeError( 1 ctx1u

81 f"copy_batch: each options element must be CopyOptions, " 1 ctx1u

82 f"got {type(a).__name__}" 1 ctx1u

83 ) 

84 return tuple(options) 2 ctx1aA

85 raise TypeError( 1 ctx1w

86 f"copy_batch: options must be CopyOptions or a sequence of " 1 ctx1w

87 f"CopyOptions, got {type(options).__name__}" 1 ctx1w

88 ) 

89  

90  

91def copy_batch( 

92 Stream stream not None: 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. 

99  

100 Source buffer and destination buffer sizes must match. For a single 

101 buffer, use :meth:`Buffer.copy_to` or :meth:`Buffer.copy_from`. 

102  

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`. 

109  

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. 

129  

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. 

143  

144 Notes 

145 ----- 

146 Batching through ``cuMemcpyBatchAsync`` requires both the CUDA 13 build of 

147 ``cuda.core`` and a driver that reports CUDA 13.0 or newer 

148 (``cuDriverGetVersion() >= 13000``). ``cuda.bindings`` binds only the 

149 CUDA 13.0 revision of the entry point, so a driver that predates it is 

150 refused even where it implements the earlier CUDA 12.8 signature. 

151  

152 The driver may execute batch items concurrently and in any order. 

153 A batch must therefore not contain copies where the source range of 

154 one copy overlaps the destination range of another; such aliasing 

155 produces undefined results. Detecting overlaps at runtime is 

156 impractical; callers are responsible for ensuring no aliasing exists. 

157  

158 On pre-CUDA 13 installs the copies fall back to a Python-level loop 

159 over ``cuMemcpyAsync``, so the potential performance benefit of 

160 asynchronous batched copies is not realized. ``src_access_order`` values 

161 of ``STREAM`` and ``ANY`` are silently ignored on the fallback path 

162 (stream-ordered access already satisfies both); ``DURING_API_CALL`` 

163 raises ``RuntimeError`` instead, since silently downgrading it to 

164 stream-ordered access would let a caller reuse the source buffer before 

165 the real read happens. 

166  

167 """ 

168 cdef tuple src_bufs = Buffer_coerce_batch(srcs, "copy_batch", _SINGLE_COPY_HINT) 32 ctx1klmnopqrzCstbcdefghaijDvuwFGHIxy

169 cdef tuple dst_bufs = Buffer_coerce_batch(dsts, "copy_batch", _SINGLE_COPY_HINT) 32 ctx1klmnopqrzCstbcdefghaijDvuwFGHIxy

170 cdef Py_ssize_t n = len(src_bufs) 28 ctx1klmnopqrzCstbcdefghaijDvuwxy

171  

172 if len(dst_bufs) != n: 28 ctx1klmnopqrzCstbcdefghaijDvuwxy

173 raise ValueError( 1 ctx1D

174 f"copy_batch: srcs length {n} does not match dsts length {len(dst_bufs)}" 1 ctx1D

175 ) 

176  

177 cdef Stream s = Stream_accept(stream) 27 ctx1klmnopqrzCstbcdefghaijvuwxy

178  

179 if Stream_is_legacy_default_token(s): 27 ctx1klmnopqrzCstbcdefghaijvuwxy

180 raise TypeError( 1 ctx1C

181 "copy_batch does not accept LEGACY_DEFAULT_STREAM; cuMemcpyBatchAsync " 

182 "rejects it outright, unlike PER_THREAD_DEFAULT_STREAM, which is a real " 

183 "stream to the driver and is accepted. Pass an explicit stream or " 

184 "PER_THREAD_DEFAULT_STREAM." 

185 ) 

186  

187 cdef cydriver.CUstreamCaptureStatus _cap_status 

188 IF CUDA_CORE_BUILD_MAJOR >= 13: 

189 HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &_cap_status, 26 ctx1klmnopqrzstbcdefghaijvuwxy

190 NULL, NULL, NULL, NULL, NULL)) 

191 ELSE: 

192 HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(as_cu(s._h_stream), &_cap_status, 

193 NULL, NULL, NULL, NULL)) 

194 if _cap_status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE: 26 ctx1klmnopqrzstbcdefghaijvuwxy

195 raise TypeError( 1 ctx1z

196 "copy_batch does not support graph capture; " 

197 "use GraphNode.memcpy or per-buffer Buffer.copy_to instead" 

198 ) 

199  

200 cdef Buffer src_buf 

201 cdef Buffer dst_buf 

202 cdef Py_ssize_t i 

203  

204 for i in range(n): 25 ctx1klmnopqrstbcdefghaijvuwxy

205 src_buf = <Buffer>src_bufs[i] 25 ctx1klmnopqrstbcdefghaijvuwxy

206 dst_buf = <Buffer>dst_bufs[i] 25 ctx1klmnopqrstbcdefghaijvuwxy

207 if src_buf.size != dst_buf.size: 25 ctx1klmnopqrstbcdefghaijvuwxy

208 raise ValueError( 2 ctx1xy

209 f"copy_batch: buffer size mismatch at index {i} " 2 ctx1xy

210 f"(src={src_buf.size}, dst={dst_buf.size})" 2 ctx1xy

211 ) 

212  

213 cdef tuple attr_tuple = _normalize_copy_options(options, n) 23 ctx1klmnopqrstbcdefghaijvuw

214  

215 _do_copy_batch(src_bufs, dst_bufs, s, attr_tuple) 20 ctx1klmnopqrstbcdefghaij

216  

217  

218cdef void _do_copy_batch(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tuple): 

219 IF CUDA_CORE_BUILD_MAJOR >= 13: 

220 # Building against CUDA 13 headers says nothing about the installed 

221 # driver, so the run-time version still has to be checked before 

222 # calling a 13.0-only entry point (see PRs #2054 / #2064). 

223 if _batch_entry_point_available(): 20 ctx1klmnopqrstbcdefghaij

224 _do_copy_batch_native(src_bufs, dst_bufs, s, attr_tuple) 20 ctx1klmnopqrstbcdefghaij

225 else: 

226 _reject_during_api_call_fallback(attr_tuple) 

227 _do_copy_batch_loop(src_bufs, dst_bufs, s) 

228 ELSE: 

229 _reject_during_api_call_fallback(attr_tuple) 

230 _do_copy_batch_loop(src_bufs, dst_bufs, s) 

231  

232  

233cdef void _reject_during_api_call_fallback(tuple attr_tuple): 

234 """Raise before the per-copy cuMemcpyAsync loop if any copy needs 

235 DURING_API_CALL, which that fallback cannot honor (see 

236 _reject_unsupported_during_api_call for why this must raise rather than 

237 silently ignore the option, unlike STREAM and ANY). 

238 """ 

239 cdef Py_ssize_t i 

240 for i in range(len(attr_tuple)): 

241 _reject_unsupported_during_api_call( 

242 (<object>attr_tuple[i]).src_access_order, 

243 "the CUDA 13 build of cuda.core and a driver that reports CUDA 13.0 or newer " 

244 "(cuMemcpyBatchAsync is unavailable here)", 

245 index=i, 

246 ) 

247  

248  

249cdef void _do_copy_batch_loop(tuple src_bufs, tuple dst_bufs, Stream s): 

250 """Per-copy cuMemcpyAsync fallback where the batch entry point is absent. 

251  

252 Issues copies one at a time, so the performance benefit of batching is 

253 not realized. STREAM and ANY are silently ignored here (satisfied by 

254 stream-ordered cuMemcpyAsync regardless); DURING_API_CALL is rejected by 

255 _reject_during_api_call_fallback before this is ever called. 

256 """ 

257 cdef Py_ssize_t n = len(src_bufs) 

258 cdef Py_ssize_t i 

259 cdef Buffer src_buf 

260 cdef Buffer dst_buf 

261 cdef size_t nbytes 

262 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 

263  

264 for i in range(n): 

265 src_buf = <Buffer>src_bufs[i] 

266 dst_buf = <Buffer>dst_bufs[i] 

267 nbytes = src_buf._size 

268 with nogil: 

269 HANDLE_RETURN(cydriver.cuMemcpyAsync( 

270 as_cu(dst_buf._h_ptr), as_cu(src_buf._h_ptr), nbytes, hstream)) 

271  

272  

273IF CUDA_CORE_BUILD_MAJOR >= 13: 

274 cdef void _do_copy_batch_native(tuple src_bufs, tuple dst_bufs, Stream s, tuple attr_tuple): 

275 cdef Py_ssize_t n = len(src_bufs) 20 ctx1klmnopqrstbcdefghaij

276 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 20 ctx1klmnopqrstbcdefghaij

277 cdef vector[cydriver.CUdeviceptr] dst_ptrs 

278 cdef vector[cydriver.CUdeviceptr] src_ptrs 

279 cdef vector[size_t] sizes 

280 cdef vector[size_t] attrs_idxs 

281 dst_ptrs.resize(n) 20 ctx1klmnopqrstbcdefghaij

282 src_ptrs.resize(n) 20 ctx1klmnopqrstbcdefghaij

283 sizes.resize(n) 20 ctx1klmnopqrstbcdefghaij

284  

285 cdef Buffer src_buf 

286 cdef Buffer dst_buf 

287 cdef Py_ssize_t i 

288  

289 # Collapse equal neighbouring attributes into runs so a broadcast 

290 # attribute reaches the driver once (numAttrs == 1) instead of being 

291 # repeated per copy. attrs[k] applies to [attrsIdxs[k], attrsIdxs[k+1]). 

292 cdef list run_starts = _attr_run_starts(attr_tuple) 20 ctx1klmnopqrstbcdefghaij

293 cdef vector[cydriver.CUmemcpyAttributes] cu_attrs 

294 cdef size_t num_attrs = <size_t>len(run_starts) 20 ctx1klmnopqrstbcdefghaij

295 cu_attrs.reserve(num_attrs) 20 ctx1klmnopqrstbcdefghaij

296 attrs_idxs.reserve(num_attrs) 20 ctx1klmnopqrstbcdefghaij

297 for i in run_starts: 20 ctx1klmnopqrstbcdefghaij

298 cu_attrs.push_back(_to_cu_memcpy_attributes(attr_tuple[i])) 20 ctx1klmnopqrstbcdefghaij

299 attrs_idxs.push_back(<size_t>i) 20 ctx1klmnopqrstbcdefghaij

300  

301 for i in range(n): 20 ctx1klmnopqrstbcdefghaij

302 src_buf = <Buffer>src_bufs[i] 20 ctx1klmnopqrstbcdefghaij

303 dst_buf = <Buffer>dst_bufs[i] 20 ctx1klmnopqrstbcdefghaij

304 src_ptrs[i] = as_cu(src_buf._h_ptr) 20 ctx1klmnopqrstbcdefghaij

305 dst_ptrs[i] = as_cu(dst_buf._h_ptr) 20 ctx1klmnopqrstbcdefghaij

306 sizes[i] = src_buf.size 20 ctx1klmnopqrstbcdefghaij

307  

308 with nogil: 20 ctx1klmnopqrstbcdefghaij

309 HANDLE_RETURN(cydriver.cuMemcpyBatchAsync( 20 ctx1klmnopqrstbcdefghaij

310 dst_ptrs.data(), 

311 src_ptrs.data(), 

312 sizes.data(), 

313 <size_t>n, 

314 cu_attrs.data(), 

315 attrs_idxs.data(), 

316 num_attrs, 

317 hstream, 

318 ))