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

177 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-29 01:38 +0000

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 

15from cuda.core._resource_handles cimport as_cu 

16from cuda.core._stream cimport Stream, Stream_accept 

17from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

18  

19from cuda.core._host import Host 

20from cuda.core._utils.cuda_utils import driver 

21from cuda.core._memory._managed_location import _coerce_location 

22  

23if TYPE_CHECKING: 

24 from cuda.core._graph import GraphBuilder 

25 from cuda.core._device import Device 

26  

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

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

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

30  

31cdef set _ADVICE_IGNORES_LOCATION = { 

32 driver.CUmem_advise.CU_MEM_ADVISE_SET_READ_MOSTLY, 

33 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_READ_MOSTLY, 

34 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION, 

35} 

36  

37cdef dict _ADVICE_ALLOWED_LOCTYPES = { 

38 driver.CUmem_advise.CU_MEM_ADVISE_SET_READ_MOSTLY: _DEVICE_HOST_NUMA, 

39 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_READ_MOSTLY: _DEVICE_HOST_NUMA, 

40 driver.CUmem_advise.CU_MEM_ADVISE_SET_PREFERRED_LOCATION: _ALL_LOCATION_TYPES, 

41 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION: _DEVICE_HOST_NUMA, 

42 driver.CUmem_advise.CU_MEM_ADVISE_SET_ACCESSED_BY: _DEVICE_HOST_ONLY, 

43 driver.CUmem_advise.CU_MEM_ADVISE_UNSET_ACCESSED_BY: _DEVICE_HOST_ONLY, 

44} 

45  

46  

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

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

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

50 # does not advertise CU_POINTER_ATTRIBUTE_IS_MANAGED). 

51 if not self.is_managed: 1abmnlokijestuvhfpcgw

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

53  

54  

55cdef tuple _coerce_batch_buffers(object buffers, str what): 

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

57  

58 For single-buffer operations, use the corresponding ManagedBuffer 

59 instance method instead. 

60 """ 

61 cdef Buffer buf 

62 cdef list out 

63 if isinstance(buffers, Buffer): 1abcgqrxyz

64 raise TypeError( 1xyz

65 f"{what}: pass a sequence of Buffers; for a single buffer use " 1xyz

66 f"the ManagedBuffer instance method" 

67 ) 

68 if isinstance(buffers, Sequence): 1abcgqr

69 if not buffers: 1abcgqr

70 raise ValueError(f"{what}: empty buffers sequence") 

71 out = [] 1abcgqr

72 for t in buffers: 1abcgqr

73 buf = <Buffer?>t 1abcgqr

74 out.append(buf) 1abcgqr

75 return tuple(out) 1abcgqr

76 raise TypeError( 

77 f"{what}: buffers must be a sequence of Buffer, " 

78 f"got {type(buffers).__name__}" 

79 ) 

80  

81  

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

83 if isinstance(location, Sequence): 1abcgqr

84 if len(location) != n: 1cqr

85 raise ValueError( 1qr

86 f"{what}: location length {len(location)} does not match " 1qr

87 f"targets length {n}" 1qr

88 ) 

89 return tuple(_coerce_location(loc, allow_none=allow_none) for loc in location) 1c

90 cdef object coerced = _coerce_location(location, allow_none=allow_none) 1abg

91 return tuple([coerced] * n) 1abg

92  

93  

94IF CUDA_CORE_BUILD_MAJOR >= 13: 

95 # Convert a _LocSpec dataclass to a cydriver.CUmemLocation struct. 

96 cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc): 

97 cdef str kind = loc.kind 1abmnlokijesthfcg

98 if kind == "device": 1abmnlokijesthfcg

99 return cydriver.CUmemLocation( 

100 type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 1abmnlokijestfcg

101 id=<int>loc.id) 1abmnlokijestfcg

102 elif kind == "host": 1blkehfc

103 return cydriver.CUmemLocation( 

104 type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1blkefc

105 id=0) 

106 elif kind == "host_numa": 1h

107 return cydriver.CUmemLocation( 

108 type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1h

109 id=<int>loc.id) 1h

110 else: # host_numa_current 

111 return cydriver.CUmemLocation( 

112 type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, 

113 id=0) 

114ELSE: 

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

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

117 cdef str kind = loc.kind 

118 if kind == "device": 

119 return <int>loc.id 

120 if kind == "host": 

121 return -1 

122 raise RuntimeError( 

123 "Host(numa_id=...) / Host.numa_current() require both cuda-bindings 13.0+ " 

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

125 ) 

126  

127  

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

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

130  

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

132 :meth:`ManagedBuffer.discard` instead. 

133  

134 Parameters 

135 ---------- 

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

137 Stream for the asynchronous discard. First positional, required 

138 (mirrors :func:`launch`). 

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

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

141 released without prefetching new contents; subsequent access is 

142 satisfied by lazy migration. 

143  

144 Raises 

145 ------ 

146 NotImplementedError 

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

148 """ 

149 cdef tuple bufs = _coerce_batch_buffers(buffers, "discard_batch") 1ax

150 cdef Stream s = Stream_accept(stream) 1a

151  

152 cdef Buffer buf 

153 for buf in bufs: 1a

154 _require_managed_buffer(buf, "discard_batch") 1a

155  

156 _do_batch_discard(bufs, s) 1a

157  

158  

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

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

161 _require_managed_buffer(buf, "discard") 1ij

162 cdef Stream s = Stream_accept(stream) 1ij

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

164 # with count=1. 

165 cdef tuple bufs = (buf,) 1ij

166 _do_batch_discard(bufs, s) 1ij

167  

168  

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

170 IF CUDA_CORE_BUILD_MAJOR >= 13: 

171 cdef Py_ssize_t n = len(bufs) 1aij

172 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 1aij

173 cdef vector[cydriver.CUdeviceptr] ptrs 

174 cdef vector[size_t] sizes 

175 ptrs.resize(n) 1aij

176 sizes.resize(n) 1aij

177 cdef Buffer buf 

178 cdef Py_ssize_t i 

179 for i in range(n): 1aij

180 buf = <Buffer>bufs[i] 1aij

181 ptrs[i] = as_cu(buf._h_ptr) 1aij

182 sizes[i] = buf._size 1aij

183 with nogil: 1aij

184 HANDLE_RETURN(cydriver.cuMemDiscardBatchAsync( 1aij

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

186 )) 

187 ELSE: 

188 raise NotImplementedError( 

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

190 ) 

191  

192  

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

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

195  

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

197 public API. 

198 """ 

199 _require_managed_buffer(buf, "advise") 1mnlokuvhfpw

200 if not isinstance(advice, driver.CUmem_advise): 1mnlokuvhfp

201 raise TypeError( 

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

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

204 ) 

205 cdef frozenset allowed_kinds = _ADVICE_ALLOWED_LOCTYPES.get(advice) 1mnlokuvhfp

206 if allowed_kinds is None: 1mnlokuvhfp

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

208 cdef bint allow_none = advice in _ADVICE_IGNORES_LOCATION 1mnlokuvhfp

209 cdef object loc = _coerce_location(location, allow_none=allow_none) 1mnlokuvhfp

210 if loc is not None and loc.kind not in allowed_kinds: 1mnlokuvhfp

211 raise ValueError( 1kuv

212 f"advise {advice.name} does not support location_type='{loc.kind}'" 1kuv

213 ) 

214 _do_single_advise(buf, advice, loc, allow_none) 1mnlokhfp

215  

216  

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

218 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 1mnlokhfp

219 cdef size_t nbytes = buf._size 1mnlokhfp

220 cdef cydriver.CUmem_advise advice_enum = <cydriver.CUmem_advise>(<int>int(advice_value)) 1mnlokhfp

221 IF CUDA_CORE_BUILD_MAJOR >= 13: 

222 cdef cydriver.CUmemLocation cu_loc 

223 if loc is None: 1mnlokhfp

224 # Driver ignores location for read_mostly / unset_preferred_location 

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

226 # host placeholder. 

227 cu_loc = cydriver.CUmemLocation( 

228 type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1kfp

229 id=0) 

230 else: 

231 cu_loc = _to_cumemlocation(loc) 1mnlokhf

232 with nogil: 1mnlokhfp

233 HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, cu_loc)) 1mnlokhfp

234 ELSE: 

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

236 with nogil: 

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

238  

239  

240def prefetch_batch( 

241 stream: Stream | GraphBuilder, 

242 buffers: Sequence[Buffer], 

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

244) -> None: 

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

246  

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

248 :meth:`ManagedBuffer.prefetch` instead. 

249  

250 Parameters 

251 ---------- 

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

253 Stream for the asynchronous prefetch. First positional, required 

254 (mirrors :func:`launch`). 

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

256 Two or more managed allocations to operate on. 

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

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

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

260  

261 Notes 

262 ----- 

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

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

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

266 """ 

267 cdef tuple bufs = _coerce_batch_buffers(buffers, "prefetch_batch") 1abcgrz

268 cdef Py_ssize_t n = len(bufs) 1abcgr

269 cdef tuple locs = _broadcast_locations(locations, n, False, "prefetch_batch") 1abcgr

270 cdef Stream s = Stream_accept(stream) 1abcg

271  

272 cdef Buffer buf 

273 for buf in bufs: 1abcg

274 _require_managed_buffer(buf, "prefetch_batch") 1abcg

275  

276 _do_batch_prefetch(bufs, locs, s) 1abcg

277  

278  

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

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

281  

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

283 """ 

284 _require_managed_buffer(buf, "prefetch") 1ijestuvw

285 cdef object loc = _coerce_location(location, allow_none=False) 1ijestuv

286 cdef Stream s = Stream_accept(stream) 1ijest

287 _do_single_prefetch(buf, loc, s) 1ijest

288  

289  

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

291 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 1ijest

292 cdef size_t nbytes = buf._size 1ijest

293 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 1ijest

294 IF CUDA_CORE_BUILD_MAJOR >= 13: 

295 cdef cydriver.CUmemLocation cu_loc = _to_cumemlocation(loc) 1ijest

296 with nogil: 1ijest

297 HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, cu_loc, 0, hstream)) 1ijest

298 ELSE: 

299 cdef int dev_int = _to_legacy_device(loc) 

300 with nogil: 

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

302  

303  

304IF CUDA_CORE_BUILD_MAJOR >= 13: 

305 # Function-pointer type for cuMemPrefetchBatchAsync / 

306 # cuMemDiscardAndPrefetchBatchAsync; both have identical signatures. 

307 ctypedef cydriver.CUresult (*_BatchPrefetchFn)( 

308 cydriver.CUdeviceptr*, size_t*, size_t, 

309 cydriver.CUmemLocation*, size_t*, size_t, 

310 unsigned long long, cydriver.CUstream, 

311 ) except ?cydriver.CUDA_ERROR_NOT_FOUND nogil 

312  

313  

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

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

316  

317 Bypasses cuda.bindings.driver.cuMemRangeGetAttribute (whose 

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

319 attributes) by calling cydriver directly. 

320  

321 Returns Device | Host | None. 

322 """ 

323 cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) 1hf

324 cdef size_t nbytes = buf._size 1hf

325 cdef int loc_type = 0 1hf

326 cdef int loc_id = 0 1hf

327 with nogil: 1hf

328 HANDLE_RETURN(cydriver.cuMemRangeGetAttribute( 1hf

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

330 cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE, 

331 cu_ptr, nbytes, 

332 )) 

333 HANDLE_RETURN(cydriver.cuMemRangeGetAttribute( 1hf

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

335 cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID, 

336 cu_ptr, nbytes, 

337 )) 

338 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE: 1hf

339 from cuda.core._device import Device 1f

340 return Device(loc_id) 1f

341 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST: 1hf

342 return Host() 1f

343 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA: 1hf

344 return Host(numa_id=loc_id) 

345 if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT: 1hf

346 return Host.numa_current() 

347 return None # CU_MEM_LOCATION_TYPE_INVALID — no preferred location 1hf

348  

349  

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

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

352 cdef Py_ssize_t n = len(bufs) 1abecg

353 cdef cydriver.CUstream hstream = as_cu(s._h_stream) 1abecg

354 cdef vector[cydriver.CUdeviceptr] ptrs 

355 cdef vector[size_t] sizes 

356 cdef vector[cydriver.CUmemLocation] loc_arr 

357 cdef vector[size_t] loc_indices 

358 ptrs.resize(n) 1abecg

359 sizes.resize(n) 1abecg

360 loc_arr.resize(n) 1abecg

361 loc_indices.resize(n) 1abecg

362 cdef Buffer buf 

363 cdef Py_ssize_t i 

364 for i in range(n): 1abecg

365 buf = <Buffer>bufs[i] 1abecg

366 ptrs[i] = as_cu(buf._h_ptr) 1abecg

367 sizes[i] = buf._size 1abecg

368 loc_arr[i] = _to_cumemlocation(locs[i]) 1abecg

369 loc_indices[i] = <size_t>i 1abecg

370 with nogil: 1abecg

371 HANDLE_RETURN(fn( 1abecg

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

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

374 0, hstream, 

375 )) 

376ELSE: 

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

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

379 # _read_preferred_location_v2` unconditionally at module top. 

380 # `ManagedBuffer.preferred_location` gates on both 

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

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

383 raise NotImplementedError( 

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

385 ) 

386  

387  

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

389 IF CUDA_CORE_BUILD_MAJOR >= 13: 

390 _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemPrefetchBatchAsync) 1abcg

391 ELSE: 

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

393 cdef Buffer buf 

394 cdef Py_ssize_t i 

395 cdef Py_ssize_t n = len(bufs) 

396 for i in range(n): 

397 buf = <Buffer>bufs[i] 

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

399  

400  

401def discard_prefetch_batch( 

402 stream: Stream | GraphBuilder, 

403 buffers: Sequence[Buffer], 

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

405) -> None: 

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

407  

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

409 :meth:`ManagedBuffer.discard_prefetch` instead. 

410  

411 Parameters 

412 ---------- 

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

414 Stream for the asynchronous operation. First positional, required 

415 (mirrors :func:`launch`). 

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

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

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

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

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

421  

422 Raises 

423 ------ 

424 NotImplementedError 

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

426 """ 

427 cdef tuple bufs = _coerce_batch_buffers(buffers, "discard_prefetch_batch") 1bqy

428 cdef Py_ssize_t n = len(bufs) 1bq

429 cdef tuple locs = _broadcast_locations(locations, n, False, "discard_prefetch_batch") 1bq

430 cdef Stream s = Stream_accept(stream) 1b

431  

432 cdef Buffer buf 

433 for buf in bufs: 1b

434 _require_managed_buffer(buf, "discard_prefetch_batch") 1b

435  

436 _do_batch_discard_prefetch(bufs, locs, s) 1b

437  

438  

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

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

441 ManagedBuffer.discard_prefetch().""" 

442 _require_managed_buffer(buf, "discard_prefetch") 1ew

443 cdef object loc = _coerce_location(location, allow_none=False) 1e

444 cdef Stream s = Stream_accept(stream) 1e

445 cdef tuple bufs = (buf,) 1e

446 cdef tuple locs = (loc,) 1e

447 _do_batch_discard_prefetch(bufs, locs, s) 1e

448  

449  

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

451 IF CUDA_CORE_BUILD_MAJOR >= 13: 

452 _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemDiscardAndPrefetchBatchAsync) 1be

453 ELSE: 

454 raise NotImplementedError( 

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

456 )