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
« 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
5from __future__ import annotations
7from collections.abc import Sequence
8from typing import TYPE_CHECKING
10IF CUDA_CORE_BUILD_MAJOR >= 13:
11 from libcpp.vector cimport vector
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
19from cuda.core._host import Host
20from cuda.core._utils.cuda_utils import driver
21from cuda.core._memory._managed_location import _coerce_location
23if TYPE_CHECKING:
24 from cuda.core._graph import GraphBuilder
25 from cuda.core._device import Device
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"))
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}
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}
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
55cdef tuple _coerce_batch_buffers(object buffers, str what):
56 """Coerce ``buffers`` to a tuple[Buffer, ...]; rejects a single Buffer.
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 )
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
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 )
128def discard_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer]) -> None:
129 """Discard a batch of managed-memory ranges.
131 Requires CUDA 13+. For a single buffer, use
132 :meth:`ManagedBuffer.discard` instead.
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.
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
152 cdef Buffer buf
153 for buf in bufs: 1a
154 _require_managed_buffer(buf, "discard_batch") 1a
156 _do_batch_discard(bufs, s) 1a
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
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 )
193def _advise_one(Buffer buf, advice: driver.CUmem_advise, location: Device | Host | None) -> None:
194 """Internal: apply managed-memory advice to a single buffer.
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
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))
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.
247 Requires CUDA 13+. For a single buffer, use
248 :meth:`ManagedBuffer.prefetch` instead.
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)``.
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
272 cdef Buffer buf
273 for buf in bufs: 1abcg
274 _require_managed_buffer(buf, "prefetch_batch") 1abcg
276 _do_batch_prefetch(bufs, locs, s) 1abcg
279def _do_single_prefetch_py(Buffer buf, location: Device | Host | None, stream: Stream | GraphBuilder | None) -> None:
280 """Internal: single-buffer prefetch for ManagedBuffer.prefetch().
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
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))
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
314 def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
315 """Internal: read preferred_location with full NUMA detail.
317 Bypasses cuda.bindings.driver.cuMemRangeGetAttribute (whose
318 attribute allowlist doesn't yet include the cu13 _TYPE / _ID
319 attributes) by calling cydriver directly.
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
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 )
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)
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.
408 Requires CUDA 13+. For a single buffer, use
409 :meth:`ManagedBuffer.discard_prefetch` instead.
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)``.
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
432 cdef Buffer buf
433 for buf in bufs: 1b
434 _require_managed_buffer(buf, "discard_prefetch_batch") 1b
436 _do_batch_discard_prefetch(bufs, locs, s) 1b
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
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 )