Coverage for cuda/core/_memory/_managed_memory_ops.pyx: 71.34%

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 

8from typing import TYPE_CHECKING 

9  

10IF CUDA_CORE_BUILD_MAJOR >= 13: 

11 from libcpp.vector cimport vector 

12  

13from cuda.bindings cimport cydriver 

14from cuda.core._memory._buffer cimport Buffer, Buffer_check_open, Buffer_coerce_batch # no-cython-lint 

15  

16# to_cumemlocation / cumemlocation_from_id are referenced only from CUDA 13 

17# branches. cython-lint does not evaluate compile-time IF blocks, so they 

18# need a pragma to be seen as used. 

19from cuda.core._memory._location cimport cumemlocation_from_id # no-cython-lint 

20from cuda.core._memory._location cimport to_cumemlocation # no-cython-lint 

21from cuda.core._rt cimport as_cu 

22from cuda.core._stream cimport Stream, Stream_accept 

23from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

24  

25from cuda.core._host import Host 

26from cuda.core._utils.cuda_utils import driver 

27from cuda.core._memory._managed_location import _coerce_location 

28  

29if TYPE_CHECKING: 

30 from cuda.core._graph import GraphBuilder 

31 from cuda.core._device import Device 

32  

33cdef frozenset _ALL_LOCATION_TYPES = frozenset(("device", "host", "host_numa", "host_numa_current")) 

34cdef frozenset _DEVICE_HOST_NUMA = frozenset(("device", "host", "host_numa")) 

35cdef frozenset _DEVICE_HOST_ONLY = frozenset(("device", "host")) 

36  

37cdef set _ADVICE_IGNORES_LOCATION = { 

38 driver.CUmem_advise.CU_MEM_ADVISE_SET_READ_MOSTLY, 

39 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_READ_MOSTLY, 

40 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION, 

41} 

42  

43cdef dict _ADVICE_ALLOWED_LOCTYPES = { 

44 driver.CUmem_advise.CU_MEM_ADVISE_SET_READ_MOSTLY: _DEVICE_HOST_NUMA, 

45 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_READ_MOSTLY: _DEVICE_HOST_NUMA, 

46 driver.CUmem_advise.CU_MEM_ADVISE_SET_PREFERRED_LOCATION: _ALL_LOCATION_TYPES, 

47 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION: _DEVICE_HOST_NUMA, 

48 driver.CUmem_advise.CU_MEM_ADVISE_SET_ACCESSED_BY: _DEVICE_HOST_ONLY, 

49 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_ACCESSED_BY: _DEVICE_HOST_ONLY, 

50} 

51  

52  

53cdef void _require_managed_buffer(Buffer self, str what): 

54 # Buffer.is_managed handles both pointer-attribute and memory-resource 

55 # paths (e.g. pool-allocated managed memory whose pointer attribute 

56 # does not advertise CU_POINTER_ATTRIBUTE_IS_MANAGED). 

57 if not self.is_managed: 23 ctx1caopqrsnhiblmjkvwgftdez

58 raise ValueError(f"{what} requires a managed-memory allocation") 1 ctx1z

59  

60  

61_SINGLE_MANAGED_HINT = "the ManagedBuffer instance method" 

62  

63  

64cdef inline tuple _coerce_batch_buffers(object buffers, str what): 

65 """Coerce ``buffers`` to a tuple[Buffer, ...]; rejects a single Buffer. 

66  

67 For single-buffer operations, use the corresponding ManagedBuffer 

68 instance method instead. 

69 """ 

70 return Buffer_coerce_batch(buffers, what, _SINGLE_MANAGED_HINT) 9 ctx1cadexyBCD

71  

72  

73cdef tuple _broadcast_locations(object location, Py_ssize_t n, bint allow_none, str what): 

74 if isinstance(location, Sequence): 6 ctx1cadexy

75 if len(location) != n: 3 ctx1dxy

76 raise ValueError( 2 ctx1xy

77 f"{what}: location length {len(location)} does not match " 2 ctx1xy

78 f"targets length {n}" 2 ctx1xy

79 ) 

80 return tuple(_coerce_location(loc, allow_none=allow_none) for loc in location) 1 ctx1d

81 cdef object coerced = _coerce_location(location, allow_none=allow_none) 3 ctx1cae

82 return tuple([coerced] * n) 3 ctx1cae

83  

84  

85IF CUDA_CORE_BUILD_MAJOR < 13: 

86 # CUDA 12 cuMemPrefetchAsync takes a device ordinal (-1 = host). 

87 cdef inline int _to_legacy_device(object loc) except? -2: 

88 cdef str kind = loc.kind 

89 if kind == "device": 

90 return <int>loc.id 

91 if kind == "host": 

92 return -1 

93 raise RuntimeError( 

94 "Host(numa_id=...) / Host.numa_current() require the CUDA 13 build of cuda.core " 

95 "and a CUDA 13+ runtime driver; use Host() instead" 

96 ) 

97  

98  

99def discard_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer]) -> None: 

100 """Discard a batch of managed-memory ranges. 

101  

102 Requires CUDA 13+. For a single buffer, use 

103 :meth:`ManagedBuffer.discard` instead. 

104  

105 Parameters 

106 ---------- 

107 stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder` 

108 Stream for the asynchronous discard. First positional, required 

109 (mirrors :func:`launch`). 

110 buffers : Sequence[:class:`Buffer`] 

111 Two or more managed allocations to discard. Resident pages are 

112 released without prefetching new contents; subsequent access is 

113 satisfied by lazy migration. 

114  

115 Raises 

116 ------ 

117 NotImplementedError 

118 On a CUDA 12 build of ``cuda.core``. 

119 """ 

120 cdef tuple bufs = _coerce_batch_buffers(buffers, "discard_batch") 2 ctx1cB

121 cdef Stream s = Stream_accept(stream) 1 ctx1c

122  

123 cdef Buffer buf 

124 for buf in bufs: 1 ctx1c

125 _require_managed_buffer(buf, "discard_batch") 1 ctx1c

126  

127 _do_batch_discard(bufs, s) 1 ctx1c

128  

129  

130def _do_single_discard_py(Buffer buf, stream: Stream | GraphBuilder | None) -> None: 

131 """Internal: single-buffer discard for ManagedBuffer.discard().""" 

132 _require_managed_buffer(buf, "discard") 2 ctx1hi

133 cdef Stream s = Stream_accept(stream) 2 ctx1hi

134 # No single-range cuMemDiscard exists; route through the batched call 

135 # with count=1. 

136 cdef tuple bufs = (buf,) 2 ctx1hi

137 _do_batch_discard(bufs, s) 2 ctx1hi

138  

139  

140cdef void _do_batch_discard(tuple bufs, Stream s): 

141 IF CUDA_CORE_BUILD_MAJOR >= 13: 

142 cdef Py_ssize_t n = len(bufs) 3 ctx1chi

143 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 3 ctx1chi

144 cdef vector[cydriver.CUdeviceptr] ptrs 

145 cdef vector[size_t] sizes 

146 ptrs.resize(n) 3 ctx1chi

147 sizes.resize(n) 3 ctx1chi

148 cdef Buffer buf 

149 cdef Py_ssize_t i 

150 for i in range(n): 3 ctx1chi

151 buf = <Buffer>bufs[i] 3 ctx1chi

152 ptrs[i] = as_cu(buf._h_ptr) 3 ctx1chi

153 sizes[i] = buf._size 3 ctx1chi

154 with nogil: 3 ctx1chi

155 HANDLE_RETURN(cydriver.cuMemDiscardBatchAsync( 3 ctx1chi

156 ptrs.data(), sizes.data(), <size_t>n, 0, hstream, 

157 )) 

158 ELSE: 

159 raise NotImplementedError( 

160 "discard requires a CUDA 13 build of cuda.core" 

161 ) 

162  

163  

164def _advise_one(Buffer buf, advice: driver.CUmem_advise, location: Device | Host | None) -> None: 

165 """Internal: apply managed-memory advice to a single buffer. 

166  

167 Used by :class:`ManagedBuffer` property setters. Not part of the 

168 public API. 

169 """ 

170 _require_managed_buffer(buf, "advise") 12 ctx1opqrsnvwgftz

171 if not isinstance(advice, driver.CUmem_advise): 11 ctx1opqrsnvwgft

172 raise TypeError( 

173 f"advice must be a cuda.bindings.driver.CUmem_advise value, " 

174 f"got {type(advice).__name__}" 

175 ) 

176 cdef frozenset allowed_kinds = _ADVICE_ALLOWED_LOCTYPES.get(advice) 11 ctx1opqrsnvwgft

177 if allowed_kinds is None: 11 ctx1opqrsnvwgft

178 raise ValueError(f"Unsupported advice value: {advice!r}") 

179 cdef bint allow_none = advice in _ADVICE_IGNORES_LOCATION 11 ctx1opqrsnvwgft

180 cdef object loc = _coerce_location(location, allow_none=allow_none) 11 ctx1opqrsnvwgft

181 if loc is not None and loc.kind not in allowed_kinds: 11 ctx1opqrsnvwgft

182 raise ValueError( 3 ctx1nvw

183 f"advise {advice.name} does not support location_type='{loc.kind}'" 3 ctx1nvw

184 ) 

185 _do_single_advise(buf, advice, loc, allow_none) 9 ctx1opqrsngft

186  

187  

188cdef void _do_single_advise(Buffer buf, object advice_value, object loc, bint allow_none): 

189 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 9 ctx1opqrsngft

190 cdef size_t nbytes = buf._size 9 ctx1opqrsngft

191 cdef cydriver.CUmem_advise advice_enum = <cydriver.CUmem_advise>(<int>int(advice_value)) 9 ctx1opqrsngft

192 IF CUDA_CORE_BUILD_MAJOR >= 13: 

193 cdef cydriver.CUmemLocation cu_loc 

194 if loc is None: 9 ctx1opqrsngft

195 # Driver ignores location for read_mostly / unset_preferred_location 

196 # advice values but still validates the CUmemLocation; pass a 

197 # host placeholder. 

198 cu_loc = cumemlocation_from_id( 3 ctx1nft

199 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 0) 

200 else: 

201 cu_loc = to_cumemlocation(loc.kind, loc.id) 8 ctx1opqrsngf

202 with nogil: 9 ctx1opqrsngft

203 HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, cu_loc)) 9 ctx1opqrsngft

204 ELSE: 

205 cdef int dev_int = -1 if loc is None else _to_legacy_device(loc) 

206 with nogil: 

207 HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, dev_int)) 

208  

209  

210def prefetch_batch( 

211 stream: Stream | GraphBuilder, 

212 buffers: Sequence[Buffer], 

213 locations: Device | Host | Sequence[Device | Host], 

214) -> None: 

215 """Prefetch a batch of managed-memory ranges to target locations. 

216  

217 Requires CUDA 13+. For a single buffer, use 

218 :meth:`ManagedBuffer.prefetch` instead. 

219  

220 Parameters 

221 ---------- 

222 stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder` 

223 Stream for the asynchronous prefetch. First positional, required 

224 (mirrors :func:`launch`). 

225 buffers : Sequence[:class:`Buffer`] 

226 Two or more managed allocations to operate on. 

227 locations : :class:`~cuda.core.Device` | :class:`~cuda.core.Host` | Sequence[...] 

228 Target location(s). A single location applies to all buffers; a 

229 sequence must match ``len(buffers)``. 

230  

231 Notes 

232 ----- 

233 On a CUDA 12 build, falls back to a Python-level loop calling 

234 ``cuMemPrefetchAsync`` per buffer (no batched driver entry point on 

235 CUDA 12). CUDA 13 builds use ``cuMemPrefetchBatchAsync`` directly. 

236 """ 

237 cdef tuple bufs = _coerce_batch_buffers(buffers, "prefetch_batch") 6 ctx1cadeyD

238 cdef Py_ssize_t n = len(bufs) 5 ctx1cadey

239 cdef tuple locs = _broadcast_locations(locations, n, False, "prefetch_batch") 5 ctx1cadey

240 cdef Stream s = Stream_accept(stream) 4 ctx1cade

241  

242 cdef Buffer buf 

243 for buf in bufs: 4 ctx1cade

244 _require_managed_buffer(buf, "prefetch_batch") 4 ctx1cade

245  

246 _do_batch_prefetch(bufs, locs, s) 4 ctx1cade

247  

248  

249def _do_single_prefetch_py(Buffer buf, location: Device | Host | None, stream: Stream | GraphBuilder | None) -> None: 

250 """Internal: single-buffer prefetch for ManagedBuffer.prefetch(). 

251  

252 Uses cuMemPrefetchAsync (works on CUDA 12 and 13). 

253 """ 

254 _require_managed_buffer(buf, "prefetch") 10 ctx1hiblmjkvwz

255 cdef object loc = _coerce_location(location, allow_none=False) 9 ctx1hiblmjkvw

256 cdef Stream s = Stream_accept(stream) 7 ctx1hiblmjk

257 _do_single_prefetch(buf, loc, s) 7 ctx1hiblmjk

258  

259  

260cdef void _do_single_prefetch(Buffer buf, object loc, Stream s): 

261 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 7 ctx1hiblmjk

262 cdef size_t nbytes = buf._size 7 ctx1hiblmjk

263 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 7 ctx1hiblmjk

264 IF CUDA_CORE_BUILD_MAJOR >= 13: 

265 cdef cydriver.CUmemLocation cu_loc = to_cumemlocation(loc.kind, loc.id) 7 ctx1hiblmjk

266 with nogil: 7 ctx1hiblmjk

267 HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, cu_loc, 0, hstream)) 7 ctx1hiblmjk

268 ELSE: 

269 cdef int dev_int = _to_legacy_device(loc) 

270 with nogil: 

271 HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, dev_int, hstream)) 

272  

273  

274IF CUDA_CORE_BUILD_MAJOR >= 13: 

275 # Function-pointer type for cuMemPrefetchBatchAsync / 

276 # cuMemDiscardAndPrefetchBatchAsync; both have identical signatures. 

277 ctypedef cydriver.CUresult (*_BatchPrefetchFn)( 

278 cydriver.CUdeviceptr*, size_t*, size_t, 

279 cydriver.CUmemLocation*, size_t*, size_t, 

280 unsigned long long, cydriver.CUstream, 

281 ) except ?cydriver.CUDA_ERROR_NOT_FOUND nogil 

282  

283  

284 cdef object _read_location_v2( 

285 Buffer buf, 

286 cydriver.CUmem_range_attribute type_attribute, 

287 cydriver.CUmem_range_attribute id_attribute, 

288 ): 

289 # cuda.bindings.driver.cuMemRangeGetAttribute does not yet accept the 

290 # CUDA 13 _TYPE / _ID attributes, so query them through cydriver. 

291 Buffer_check_open(buf) 12 ctx1ablmujkAgfde

292 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 11 ctx1ablmujkgfde

293 cdef size_t nbytes = buf._size 11 ctx1ablmujkgfde

294 cdef int loc_type = 0 11 ctx1ablmujkgfde

295 cdef int loc_id = 0 11 ctx1ablmujkgfde

296 with nogil: 11 ctx1ablmujkgfde

297 HANDLE_RETURN(cydriver.cuMemRangeGetAttribute( 11 ctx1ablmujkgfde

298 <void*>&loc_type, sizeof(int), 

299 type_attribute, 

300 cu_ptr, nbytes, 

301 )) 

302 HANDLE_RETURN(cydriver.cuMemRangeGetAttribute( 11 ctx1ablmujkgfde

303 <void*>&loc_id, sizeof(int), 

304 id_attribute, 

305 cu_ptr, nbytes, 

306 )) 

307 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE: 11 ctx1ablmujkgfde

308 from cuda.core._device import Device 7 ctx1ablmfde

309 return Device(loc_id) 7 ctx1ablmfde

310 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST: 6 ctx1ujkgfd

311 return Host() 2 ctx1fd

312 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA: 5 ctx1ujkgf

313 return Host(numa_id=loc_id) 2 ctx1jk

314 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT: 3 ctx1ugf

315 return Host.numa_current() 

316 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_INVALID: 3 ctx1ugf

317 return None 3 ctx1ugf

318 raise RuntimeError( 

319 f"cuMemRangeGetAttribute returned unexpected location type: {loc_type}" 

320 ) 

321  

322  

323 def _read_preferred_location_v2(Buffer buf) -> Device | Host | None: 

324 """Internal: read preferred_location with full NUMA detail.""" 

325 return _read_location_v2( 3 ctx1Agf

326 buf, 

327 cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE, 

328 cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID, 3 ctx1Agf

329 ) 

330  

331  

332 def _read_last_prefetch_location_v2(Buffer buf) -> Device | Host | None: 

333 """Internal: read last_prefetch_location with full NUMA detail.""" 

334 return _read_location_v2( 10 ctx1ablmujkAde

335 buf, 

336 cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION_TYPE, 

337 cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION_ID, 10 ctx1ablmujkAde

338 ) 

339  

340  

341 cdef void _do_batch_prefetch_op(tuple bufs, tuple locs, Stream s, _BatchPrefetchFn fn): 

342 """Shared body for batched prefetch / discard-and-prefetch.""" 

343 cdef Py_ssize_t n = len(bufs) 5 ctx1cabde

344 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 5 ctx1cabde

345 cdef vector[cydriver.CUdeviceptr] ptrs 

346 cdef vector[size_t] sizes 

347 cdef vector[cydriver.CUmemLocation] loc_arr 

348 cdef vector[size_t] loc_indices 

349 ptrs.resize(n) 5 ctx1cabde

350 sizes.resize(n) 5 ctx1cabde

351 loc_arr.resize(n) 5 ctx1cabde

352 loc_indices.resize(n) 5 ctx1cabde

353 cdef Buffer buf 

354 cdef Py_ssize_t i 

355 cdef object loc_spec 

356 for i in range(n): 5 ctx1cabde

357 buf = <Buffer>bufs[i] 5 ctx1cabde

358 ptrs[i] = as_cu(buf._h_ptr) 5 ctx1cabde

359 sizes[i] = buf._size 5 ctx1cabde

360 loc_spec = locs[i] 5 ctx1cabde

361 loc_arr[i] = to_cumemlocation(loc_spec.kind, loc_spec.id) 5 ctx1cabde

362 loc_indices[i] = <size_t>i 5 ctx1cabde

363 with nogil: 5 ctx1cabde

364 HANDLE_RETURN(fn( 5 ctx1cabde

365 ptrs.data(), sizes.data(), <size_t>n, 

366 loc_arr.data(), loc_indices.data(), <size_t>n, 

367 0, hstream, 

368 )) 

369ELSE: 

370 def _read_preferred_location_v2(Buffer buf) -> Device | Host | None: 

371 # Symbols exist so _managed_buffer.py can import the v2 readers 

372 # unconditionally. Their properties gate on the CUDA 13 build and 

373 # driver_version() >= (13, 0, 0), so these paths are unreachable on a 

374 # CUDA 12 build. 

375 raise NotImplementedError( 

376 "_read_preferred_location_v2 requires a CUDA 13 build of cuda.core" 

377 ) 

378  

379  

380 def _read_last_prefetch_location_v2(Buffer buf) -> Device | Host | None: 

381 raise NotImplementedError( 

382 "_read_last_prefetch_location_v2 requires a CUDA 13 build of cuda.core" 

383 ) 

384  

385  

386cdef void _do_batch_prefetch(tuple bufs, tuple locs, Stream s): 

387 IF CUDA_CORE_BUILD_MAJOR >= 13: 

388 _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemPrefetchBatchAsync) 4 ctx1cade

389 ELSE: 

390 # cu12 has no cuMemPrefetchBatchAsync; loop per-range. 

391 cdef Buffer buf 

392 cdef Py_ssize_t i 

393 cdef Py_ssize_t n = len(bufs) 

394 for i in range(n): 

395 buf = <Buffer>bufs[i] 

396 _do_single_prefetch(buf, locs[i], s) 

397  

398  

399def discard_prefetch_batch( 

400 stream: Stream | GraphBuilder, 

401 buffers: Sequence[Buffer], 

402 locations: Device | Host | Sequence[Device | Host], 

403) -> None: 

404 """Discard a batch of managed-memory ranges and prefetch them to target locations. 

405  

406 Requires CUDA 13+. For a single buffer, use 

407 :meth:`ManagedBuffer.discard_prefetch` instead. 

408  

409 Parameters 

410 ---------- 

411 stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder` 

412 Stream for the asynchronous operation. First positional, required 

413 (mirrors :func:`launch`). 

414 buffers : Sequence[:class:`Buffer`] 

415 Two or more managed allocations to discard and re-prefetch. 

416 locations : :class:`~cuda.core.Device` | :class:`~cuda.core.Host` | Sequence[...] 

417 Target location(s). A single location applies to all buffers; 

418 a sequence must match ``len(buffers)``. 

419  

420 Raises 

421 ------ 

422 NotImplementedError 

423 On a CUDA 12 build of ``cuda.core``. 

424 """ 

425 cdef tuple bufs = _coerce_batch_buffers(buffers, "discard_prefetch_batch") 3 ctx1axC

426 cdef Py_ssize_t n = len(bufs) 2 ctx1ax

427 cdef tuple locs = _broadcast_locations(locations, n, False, "discard_prefetch_batch") 2 ctx1ax

428 cdef Stream s = Stream_accept(stream) 1 ctx1a

429  

430 cdef Buffer buf 

431 for buf in bufs: 1 ctx1a

432 _require_managed_buffer(buf, "discard_prefetch_batch") 1 ctx1a

433  

434 _do_batch_discard_prefetch(bufs, locs, s) 1 ctx1a

435  

436  

437def _do_single_discard_prefetch_py(Buffer buf, location: Device | Host | None, stream: Stream | GraphBuilder | None) -> None: 

438 """Internal: single-buffer discard+prefetch for 

439 ManagedBuffer.discard_prefetch().""" 

440 _require_managed_buffer(buf, "discard_prefetch") 2 ctx1bz

441 cdef object loc = _coerce_location(location, allow_none=False) 1 ctx1b

442 cdef Stream s = Stream_accept(stream) 1 ctx1b

443 cdef tuple bufs = (buf,) 1 ctx1b

444 cdef tuple locs = (loc,) 1 ctx1b

445 _do_batch_discard_prefetch(bufs, locs, s) 1 ctx1b

446  

447  

448cdef void _do_batch_discard_prefetch(tuple bufs, tuple locs, Stream s): 

449 IF CUDA_CORE_BUILD_MAJOR >= 13: 

450 _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemDiscardAndPrefetchBatchAsync) 2 ctx1ab

451 ELSE: 

452 raise NotImplementedError( 

453 "discard_prefetch requires a CUDA 13 build of cuda.core" 

454 )