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

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from __future__ import annotations 

6  

7from collections.abc import Sequence (empty)

8from typing import TYPE_CHECKING (empty)

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._resource_handles 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 (empty)

26from cuda.core._utils.cuda_utils import driver (empty)

27from cuda.core._memory._managed_location import _coerce_location (empty)

28  

29if TYPE_CHECKING: (empty)

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")) (empty)

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

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

36  

37cdef set _ADVICE_IGNORES_LOCATION = { 

38 driver.CUmem_advise.CU_MEM_ADVISE_SET_READ_MOSTLY, (empty)

39 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_READ_MOSTLY, (empty)

40 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION, (empty)

41} 

42  

43cdef dict _ADVICE_ALLOWED_LOCTYPES = { 

44 driver.CUmem_advise.CU_MEM_ADVISE_SET_READ_MOSTLY: _DEVICE_HOST_NUMA, (empty)

45 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_READ_MOSTLY: _DEVICE_HOST_NUMA, (empty)

46 driver.CUmem_advise.CU_MEM_ADVISE_SET_PREFERRED_LOCATION: _ALL_LOCATION_TYPES, (empty)

47 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION: _DEVICE_HOST_NUMA, (empty)

48 driver.CUmem_advise.CU_MEM_ADVISE_SET_ACCESSED_BY: _DEVICE_HOST_ONLY, (empty)

49 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_ACCESSED_BY: _DEVICE_HOST_ONLY, (empty)

50} 

51  

52  

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

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: 20 ctx1aclmnokijdstqrfepghw

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

59  

60  

61_SINGLE_MANAGED_HINT = "the ManagedBuffer instance method" (empty)

62  

63  

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

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 ctx1acghuvxyz

71  

72  

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

74 if isinstance(location, Sequence): 6 ctx1acghuv

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

76 raise ValueError( 2 ctx1uv

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

78 f"targets length {n}" 3 ctx1buv

79 ) 

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

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

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

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 both cuda-bindings 13.0+ " 

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

96 ) 

97  

98  

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

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 ctx1ax

121 cdef Stream s = Stream_accept(stream) 1 ctx1a

122  

123 cdef Buffer buf 

124 for buf in bufs: 1 ctx1a

125 _require_managed_buffer(buf, "discard_batch") 1 ctx1a

126  

127 _do_batch_discard(bufs, s) 1 ctx1a

128  

129  

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

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

132 _require_managed_buffer(buf, "discard") 2 ctx1ij

133 cdef Stream s = Stream_accept(stream) 3 ctx1bij

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

135 # with count=1. 

136 cdef tuple bufs = (buf,) 2 ctx1ij

137 _do_batch_discard(bufs, s) 2 ctx1ij

138  

139  

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

141 IF CUDA_CORE_BUILD_MAJOR >= 13: 

142 cdef Py_ssize_t n = len(bufs) 3 ctx1aij

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

144 cdef vector[cydriver.CUdeviceptr] ptrs 

145 cdef vector[size_t] sizes 

146 ptrs.resize(n) 3 ctx1aij

147 sizes.resize(n) 3 ctx1aij

148 cdef Buffer buf 

149 cdef Py_ssize_t i 

150 for i in range(n): 3 ctx1aij

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

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

153 sizes[i] = buf._size 3 ctx1aij

154 with nogil: 3 ctx1aij

155 HANDLE_RETURN(cydriver.cuMemDiscardBatchAsync( 3 ctx1aij

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: (empty)

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") 11 ctx1lmnokqrfepw

171 if not isinstance(advice, driver.CUmem_advise): 10 ctx1lmnokqrfep

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) 10 ctx1lmnokqrfep

177 if allowed_kinds is None: 10 ctx1lmnokqrfep

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

179 cdef bint allow_none = advice in _ADVICE_IGNORES_LOCATION 10 ctx1lmnokqrfep

180 cdef object loc = _coerce_location(location, allow_none=allow_none) 10 ctx1lmnokqrfep

181 if loc is not None and loc.kind not in allowed_kinds: 10 ctx1lmnokqrfep

182 raise ValueError( 3 ctx1kqr

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

184 ) 

185 _do_single_advise(buf, advice, loc, allow_none) 8 ctx1lmnokfep

186  

187  

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

189 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 8 ctx1lmnokfep

190 cdef size_t nbytes = buf._size 8 ctx1lmnokfep

191 cdef cydriver.CUmem_advise advice_enum = <cydriver.CUmem_advise>(<int>int(advice_value)) 8 ctx1lmnokfep

192 IF CUDA_CORE_BUILD_MAJOR >= 13: 

193 cdef cydriver.CUmemLocation cu_loc 

194 if loc is None: 8 ctx1lmnokfep

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 ctx1kep

199 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 0) 

200 else: 

201 cu_loc = to_cumemlocation(loc.kind, loc.id) 7 ctx1lmnokfe

202 with nogil: 8 ctx1lmnokfep

203 HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, cu_loc)) 8 ctx1lmnokfep

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( (empty)

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 ctx1acghvz

238 cdef Py_ssize_t n = len(bufs) 5 ctx1acghv

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

240 cdef Stream s = Stream_accept(stream) 4 ctx1acgh

241  

242 cdef Buffer buf 

243 for buf in bufs: 4 ctx1acgh

244 _require_managed_buffer(buf, "prefetch_batch") 4 ctx1acgh

245  

246 _do_batch_prefetch(bufs, locs, s) 4 ctx1acgh

247  

248  

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

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") 8 ctx1ijdstqrw

255 cdef object loc = _coerce_location(location, allow_none=False) 7 ctx1ijdstqr

256 cdef Stream s = Stream_accept(stream) 5 ctx1ijdst

257 _do_single_prefetch(buf, loc, s) 5 ctx1ijdst

258  

259  

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

261 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 5 ctx1ijdst

262 cdef size_t nbytes = buf._size 5 ctx1ijdst

263 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 5 ctx1ijdst

264 IF CUDA_CORE_BUILD_MAJOR >= 13: 

265 cdef cydriver.CUmemLocation cu_loc = to_cumemlocation(loc.kind, loc.id) 5 ctx1ijdst

266 with nogil: 5 ctx1ijdst

267 HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, cu_loc, 0, hstream)) 5 ctx1ijdst

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 def _read_preferred_location_v2(Buffer buf) -> Device | Host | None: (empty)

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

286  

287 Bypasses cuda.bindings.driver.cuMemRangeGetAttribute (whose 

288 attribute allowlist doesn't yet include the cu13 _TYPE / _ID 

289 attributes) by calling cydriver directly. 

290  

291 Returns Device | Host | None. 

292 """ 

293 Buffer_check_open(buf) 2 ctx1fe

294 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 2 ctx1fe

295 cdef size_t nbytes = buf._size 2 ctx1fe

296 cdef int loc_type = 0 2 ctx1fe

297 cdef int loc_id = 0 2 ctx1fe

298 with nogil: 2 ctx1fe

299 HANDLE_RETURN(cydriver.cuMemRangeGetAttribute( 2 ctx1fe

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

301 cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE, 

302 cu_ptr, nbytes, 

303 )) 

304 HANDLE_RETURN(cydriver.cuMemRangeGetAttribute( 2 ctx1fe

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

306 cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID, 

307 cu_ptr, nbytes, 

308 )) 

309 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE: 2 ctx1fe

310 from cuda.core._device import Device 1 ctx1e

311 return Device(loc_id) 1 ctx1e

312 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST: 2 ctx1fe

313 return Host() 1 ctx1e

314 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA: 2 ctx1fe

315 return Host(numa_id=loc_id) 

316 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT: 2 ctx1fe

317 return Host.numa_current() 

318 return None # CU_MEM_LOCATION_TYPE_INVALID — no preferred location 2 ctx1fe

319  

320  

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

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

323 cdef Py_ssize_t n = len(bufs) 5 ctx1acdgh

324 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 5 ctx1acdgh

325 cdef vector[cydriver.CUdeviceptr] ptrs 

326 cdef vector[size_t] sizes 

327 cdef vector[cydriver.CUmemLocation] loc_arr 

328 cdef vector[size_t] loc_indices 

329 ptrs.resize(n) 5 ctx1acdgh

330 sizes.resize(n) 5 ctx1acdgh

331 loc_arr.resize(n) 5 ctx1acdgh

332 loc_indices.resize(n) 5 ctx1acdgh

333 cdef Buffer buf 

334 cdef Py_ssize_t i 

335 cdef object loc_spec 

336 for i in range(n): 5 ctx1acdgh

337 buf = <Buffer>bufs[i] 5 ctx1acdgh

338 ptrs[i] = as_cu(buf._h_ptr) 5 ctx1acdgh

339 sizes[i] = buf._size 5 ctx1acdgh

340 loc_spec = locs[i] 5 ctx1acdgh

341 loc_arr[i] = to_cumemlocation(loc_spec.kind, loc_spec.id) 5 ctx1acdgh

342 loc_indices[i] = <size_t>i 5 ctx1acdgh

343 with nogil: 5 ctx1acdgh

344 HANDLE_RETURN(fn( 5 ctx1acdgh

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

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

347 0, hstream, 

348 )) 

349ELSE: 

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

351 # Symbol exists so _managed_buffer.py can `from ... import 

352 # _read_preferred_location_v2` unconditionally at module top. 

353 # `ManagedBuffer.preferred_location` gates on both 

354 # binding_version() and driver_version() >= (13, 0, 0) before 

355 # calling, so this path is unreachable on a cu12 build. 

356 raise NotImplementedError( 

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

358 ) 

359  

360  

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

362 IF CUDA_CORE_BUILD_MAJOR >= 13: 

363 _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemPrefetchBatchAsync) 4 ctx1acgh

364 ELSE: 

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

366 cdef Buffer buf 

367 cdef Py_ssize_t i 

368 cdef Py_ssize_t n = len(bufs) 

369 for i in range(n): 

370 buf = <Buffer>bufs[i] 

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

372  

373  

374def discard_prefetch_batch( (empty)

375 stream: Stream | GraphBuilder, 

376 buffers: Sequence[Buffer], 

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

378) -> None: 

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

380  

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

382 :meth:`ManagedBuffer.discard_prefetch` instead. 

383  

384 Parameters 

385 ---------- 

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

387 Stream for the asynchronous operation. First positional, required 

388 (mirrors :func:`launch`). 

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

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

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

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

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

394  

395 Raises 

396 ------ 

397 NotImplementedError 

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

399 """ 

400 cdef tuple bufs = _coerce_batch_buffers(buffers, "discard_prefetch_batch") 3 ctx1cuy

401 cdef Py_ssize_t n = len(bufs) 2 ctx1cu

402 cdef tuple locs = _broadcast_locations(locations, n, False, "discard_prefetch_batch") 2 ctx1cu

403 cdef Stream s = Stream_accept(stream) 1 ctx1c

404  

405 cdef Buffer buf 

406 for buf in bufs: 1 ctx1c

407 _require_managed_buffer(buf, "discard_prefetch_batch") 1 ctx1c

408  

409 _do_batch_discard_prefetch(bufs, locs, s) 1 ctx1c

410  

411  

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

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

414 ManagedBuffer.discard_prefetch().""" 

415 _require_managed_buffer(buf, "discard_prefetch") 2 ctx1dw

416 cdef object loc = _coerce_location(location, allow_none=False) 1 ctx1d

417 cdef Stream s = Stream_accept(stream) 1 ctx1d

418 cdef tuple bufs = (buf,) 1 ctx1d

419 cdef tuple locs = (loc,) 1 ctx1d

420 _do_batch_discard_prefetch(bufs, locs, s) 1 ctx1d

421  

422  

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

424 IF CUDA_CORE_BUILD_MAJOR >= 13: 

425 _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemDiscardAndPrefetchBatchAsync) 2 ctx1cd

426 ELSE: 

427 raise NotImplementedError( 

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

429 )