Coverage for cuda/core/_memory/_managed_buffer.py: 44.44%

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4from __future__ import annotations 

5 

6from collections.abc import Iterable, Iterator, MutableSet 

7from typing import TYPE_CHECKING, Any 

8 

9from cuda.core._device import Device 

10from cuda.core._host import Host 

11from cuda.core._memory._buffer import Buffer 

12from cuda.core._memory._managed_location import _coerce_location 

13from cuda.core._memory._managed_memory_ops import ( 

14 _advise_one, 

15 _do_single_discard_prefetch_py, 

16 _do_single_discard_py, 

17 _do_single_prefetch_py, 

18 _read_last_prefetch_location_v2, 

19 _read_preferred_location_v2, 

20) 

21from cuda.core._utils.cuda_utils import driver, handle_return 

22from cuda.core._utils.version import BUILD_CUDA_MAJOR, driver_version 

23 

24if TYPE_CHECKING: 

25 from cuda.core._memory._buffer import MemoryResource 

26 from cuda.core._stream import Stream 

27 from cuda.core.graph import GraphBuilder 

28 

29__all__ = ["ManagedBuffer"] 

30 

31 

32_INT_SIZE = 4 

33 

34# Enum aliases — referenced once per property write, so cache the lookup. 

35_ADV = driver.CUmem_advise 

36_SET_READ_MOSTLY = _ADV.CU_MEM_ADVISE_SET_READ_MOSTLY 

37_UNSET_READ_MOSTLY = _ADV.CU_MEM_ADVISE_UNSET_READ_MOSTLY 

38_SET_PREFERRED = _ADV.CU_MEM_ADVISE_SET_PREFERRED_LOCATION 

39_UNSET_PREFERRED = _ADV.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION 

40_SET_ACCESSED_BY = _ADV.CU_MEM_ADVISE_SET_ACCESSED_BY 

41_UNSET_ACCESSED_BY = _ADV.CU_MEM_ADVISE_UNSET_ACCESSED_BY 

42 

43_RANGE = driver.CUmem_range_attribute 

44_ATTR_READ_MOSTLY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_READ_MOSTLY 

45_ATTR_PREFERRED = _RANGE.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION 

46_ATTR_ACCESSED_BY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_ACCESSED_BY 

47_ATTR_LAST_PREFETCH = _RANGE.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION 

48 

49 

50def _check_open(buf: Buffer) -> None: 

51 if buf.is_closed: 10 ctx1edabcgfhij

52 raise RuntimeError("Buffer has been closed") 1 ctx1f

53 

54 

55def _get_int_attr(buf: Buffer, attribute: Any) -> int: 

56 _check_open(buf) 2 ctx1fj

57 return int(handle_return(driver.cuMemRangeGetAttribute(_INT_SIZE, attribute, buf.handle, buf.size))) 1 ctx1j

58 

59 

60def _query_accessed_by(buf: Buffer) -> list[Device | Host]: 

61 """Read the live ``CU_MEM_RANGE_ATTRIBUTE_ACCESSED_BY`` list. 

62 

63 Driver fills an int32 array: device id, ``-1`` = host, ``-2`` = empty. 

64 Sized to ``cuDeviceGetCount() + 1`` (every visible device plus host). 

65 """ 

66 _check_open(buf) 6 ctx1edabcf

67 num_devices = handle_return(driver.cuDeviceGetCount()) 5 ctx1edabc

68 n = num_devices + 1 5 ctx1edabc

69 raw = handle_return(driver.cuMemRangeGetAttribute(n * _INT_SIZE, _ATTR_ACCESSED_BY, buf.handle, buf.size)) 5 ctx1edabc

70 return [Host() if v == -1 else Device(v) for v in raw if v != -2] 5 ctx1edabc

71 

72 

73class AccessedBySetProxy(MutableSet[Device | Host]): 

74 """Live driver-backed view of ``set_accessed_by`` advice for a managed buffer. 

75 

76 Reads (``__contains__``, ``__iter__``, ``len(...)``) call 

77 ``cuMemRangeGetAttribute``; writes (``add``, ``discard``) call 

78 ``cuMemAdvise``. There is no in-memory mirror, so the view always 

79 reflects the current driver state. 

80 

81 Note 

82 ---- 

83 The driver returns integer device ordinals (``-1`` for host); host 

84 NUMA distinctions applied via ``Host(numa_id=...)`` collapse to a 

85 generic ``Host()`` when iterating this set. 

86 """ 

87 

88 __slots__ = ("_buf",) 

89 

90 def __init__(self, buf: ManagedBuffer): 

91 self._buf = buf 9 ctx1edabcgfhi

92 

93 # Operators such as &|^ produce a plain set, not another proxy. 

94 @classmethod 

95 def _from_iterable(cls, it: Iterable[Any]) -> set[Device | Host]: # type: ignore[override] 

96 return set(it) 1 ctx1d

97 

98 # --- abstract methods required by MutableSet --- 

99 

100 def __contains__(self, location: object) -> bool: 

101 if not isinstance(location, (Device, Host)): 5 ctx1edabc

102 return False 1 ctx1a

103 return location in _query_accessed_by(self._buf) 4 ctx1edbc

104 

105 def __iter__(self) -> Iterator[Device | Host]: 

106 return iter(_query_accessed_by(self._buf)) 2 ctx1da

107 

108 def __len__(self) -> int: 

109 return len(_query_accessed_by(self._buf)) 2 ctx1df

110 

111 def add(self, location: Device | Host) -> None: 

112 """Apply ``set_accessed_by`` advice for ``location``.""" 

113 if not isinstance(location, (Device, Host)): 6 ctx1edaghi

114 raise TypeError(f"expected Device or Host, got {type(location).__name__}") 1 ctx1a

115 _advise_one(self._buf, _SET_ACCESSED_BY, location) 5 ctx1edghi

116 

117 def discard(self, location: Device | Host) -> None: 

118 """Apply ``unset_accessed_by`` advice for ``location``. 

119 

120 Per the ``MutableSet`` contract, ``discard`` is a no-op for elements 

121 not in the set. ``set_accessed_by`` only accepts ``Device`` and the 

122 generic ``Host()`` — NUMA-aware host variants (``Host(numa_id=...)``, 

123 ``Host.numa_current()``) can never enter the set, so discarding them 

124 is silently ignored rather than forwarded to the driver. 

125 """ 

126 if not isinstance(location, (Device, Host)): 3 ctx1eda

127 return 1 ctx1a

128 if isinstance(location, Host) and (location.numa_id is not None or location.is_numa_current): 2 ctx1ed

129 return 1 ctx1d

130 _advise_one(self._buf, _UNSET_ACCESSED_BY, location) 2 ctx1ed

131 

132 def __repr__(self) -> str: 

133 return f"AccessedBySetProxy({set(_query_accessed_by(self._buf))!r})" 1 ctx1d

134 

135 

136class ManagedBuffer(Buffer): 

137 """Managed (unified) memory buffer with a property-style advice API. 

138 

139 Returned by :meth:`ManagedMemoryResource.allocate`, or wrap an 

140 existing managed-memory pointer with :meth:`ManagedBuffer.from_handle`. 

141 

142 Examples 

143 -------- 

144 >>> buf = mr.allocate(size) 

145 >>> buf.read_mostly = True 

146 >>> buf.preferred_location = Device(0) 

147 >>> buf.accessed_by.add(Device(1)) 

148 >>> buf.prefetch(Device(0), stream=stream) 

149 

150 Note 

151 ---- 

152 On CUDA 13 builds, ``preferred_location`` and 

153 ``last_prefetch_location`` round-trip full NUMA information. On CUDA 12 

154 builds, ``Host(numa_id=...)`` and ``Host.numa_current()`` are rejected 

155 with ``TypeError`` at the call boundary — only ``Device(...)`` and the 

156 generic ``Host()`` are accepted. Use ``Host()`` to target the host on 

157 CUDA 12. 

158 """ 

159 

160 @classmethod 

161 def from_handle( 

162 cls, 

163 ptr, 

164 size: int, 

165 mr: MemoryResource | None = None, 

166 owner: object | None = None, 

167 *, 

168 stream: Stream | GraphBuilder | None = None, 

169 ) -> Buffer: 

170 """Wrap an existing managed-memory pointer in a :class:`ManagedBuffer`. 

171 

172 Use this when you have an externally-allocated managed pointer 

173 and want the property-style advice API (:attr:`read_mostly`, 

174 :attr:`preferred_location`, :attr:`last_prefetch_location`, 

175 :attr:`accessed_by`). 

176 

177 Parameters 

178 ---------- 

179 ptr : :obj:`~_memory.DevicePointerT` 

180 Pointer to a managed allocation. 

181 size : int 

182 Allocation size in bytes. 

183 mr : :obj:`~_memory.MemoryResource`, optional 

184 Memory resource that owns ``ptr``. When provided, its 

185 ``deallocate`` is called when the buffer is closed. 

186 owner : object, optional 

187 An object that keeps the underlying allocation alive. 

188 ``owner`` and ``mr`` cannot both be specified. 

189 stream : Stream | GraphBuilder, optional 

190 Keyword-only. The stream used to order the buffer's deallocation 

191 when ``mr`` owns the pointer. Defaults to ``default_stream()``. 

192 Recording a default-stream token requires a CUDA context to be 

193 current. If the buffer may be freed from a different host thread, 

194 pass a stream other than the per-thread default stream, which 

195 refers to a different stream on each thread. 

196 """ 

197 return cls._init(ptr, size, mr=mr, owner=owner, stream=stream) 22 ctx1yzABCDEFGHIfJkKLmMNOPQ

198 

199 @property 

200 def read_mostly(self) -> bool: 

201 """Whether ``set_read_mostly`` advice is currently applied.""" 

202 return _get_int_attr(self, _ATTR_READ_MOSTLY) != 0 2 ctx1fj

203 

204 @read_mostly.setter 

205 def read_mostly(self, value: bool) -> None: 

206 _advise_one(self, _SET_READ_MOSTLY if value else _UNSET_READ_MOSTLY, None) 3 ctx1gjm

207 

208 @property 

209 def preferred_location(self) -> Device | Host | None: 

210 """Currently applied ``set_preferred_location`` target, or ``None``. 

211 

212 On CUDA 13 builds, fully round-trips ``Host(numa_id=N)``. On CUDA 12 

213 the legacy attribute carries only a device ordinal (or ``-1`` for 

214 host), so ``Host(numa_id=N)`` set via the setter round-trips back 

215 as ``Host()``. 

216 """ 

217 # The v2 path uses CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_{TYPE,ID}, 

218 # both added in CUDA 13. The path exists in the CUDA 13 build only, and 

219 # the runtime driver must be 13.0+. Otherwise fall back to the legacy 

220 # device-ordinal path. See PR #2054 / #2064 for prior regressions. 

221 if BUILD_CUDA_MAJOR >= 13 and driver_version() >= (13, 0, 0): 3 ctx1fkl

222 return _read_preferred_location_v2(self) 3 ctx1fkl

223 # CUDA 12 legacy path. No NUMA info is available. A CUDA 13 build also 

224 # takes this path when the runtime driver is still 12.x. 

225 loc_id = _get_int_attr(self, _ATTR_PREFERRED) 

226 if loc_id == -2: 

227 return None 

228 if loc_id == -1: 

229 return Host() 

230 return Device(loc_id) 

231 

232 @preferred_location.setter 

233 def preferred_location(self, value: Device | Host | None) -> None: 

234 if value is None: 3 ctx1gkl

235 _advise_one(self, _UNSET_PREFERRED, None) 1 ctx1l

236 else: 

237 _advise_one(self, _SET_PREFERRED, value) 3 ctx1gkl

238 

239 @property 

240 def last_prefetch_location(self) -> Device | Host | None: 

241 """Location targeted by the most recent explicit prefetch. 

242 

243 Returns ``None`` if any page in the range has never been prefetched, 

244 or if the pages do not share one last-prefetch location. This reports 

245 the application's requested destination, not current residency or 

246 completion of the asynchronous prefetch operation. 

247 

248 On CUDA 13 builds, fully round-trips ``Host(numa_id=N)``. On CUDA 12, 

249 the legacy attribute carries only a device ordinal (or ``-1`` for 

250 host), so host NUMA details are unavailable. 

251 """ 

252 if BUILD_CUDA_MAJOR >= 13 and driver_version() >= (13, 0, 0): 10 ctx1snoptqrfuv

253 return _read_last_prefetch_location_v2(self) 10 ctx1snoptqrfuv

254 loc_id = _get_int_attr(self, _ATTR_LAST_PREFETCH) 

255 if loc_id == -2: 

256 return None 

257 if loc_id == -1: 

258 return Host() 

259 return Device(loc_id) 

260 

261 @property 

262 def accessed_by(self) -> AccessedBySetProxy: 

263 """Live set-like view of ``set_accessed_by`` locations.""" 

264 _check_open(self) 9 ctx1edabcgfhi

265 return AccessedBySetProxy(self) 9 ctx1edabcgfhi

266 

267 @accessed_by.setter 

268 def accessed_by(self, locations: Iterable[Device | Host]) -> None: 

269 _check_open(self) 3 ctx1abc

270 # Validate every target before issuing any cuMemAdvise so an invalid 

271 # element can't leave accessed_by partially mutated. 

272 target: set[Device | Host] = set() 3 ctx1abc

273 for loc in locations: 3 ctx1abc

274 if not isinstance(loc, (Device, Host)): 3 ctx1abc

275 raise TypeError(f"accessed_by entries must be Device or Host, got {type(loc).__name__}") 1 ctx1a

276 target.add(loc) 3 ctx1abc

277 for loc in target: 3 ctx1abc

278 spec = _coerce_location(loc) 3 ctx1abc

279 assert spec is not None 3 ctx1abc

280 if spec.kind not in ("device", "host"): 3 ctx1abc

281 raise ValueError(f"advise {_SET_ACCESSED_BY.name} does not support location_type='{spec.kind}'") 1 ctx1b

282 current = set(_query_accessed_by(self)) 3 ctx1abc

283 for loc in current - target: 3 ctx1abc

284 _advise_one(self, _UNSET_ACCESSED_BY, loc) 1 ctx1c

285 for loc in target - current: 3 ctx1abc

286 _advise_one(self, _SET_ACCESSED_BY, loc) 3 ctx1abc

287 

288 def prefetch(self, location: Device | Host, *, stream: Stream | GraphBuilder) -> None: 

289 """Prefetch this range to ``location`` on ``stream``.""" 

290 _do_single_prefetch_py(self, location, stream) 10 ctx1wxnopqrhim

291 

292 def discard(self, *, stream: Stream | GraphBuilder) -> None: 

293 """Discard this range's resident pages on ``stream`` (CUDA 13+).""" 

294 _do_single_discard_py(self, stream) 2 ctx1wx

295 

296 def discard_prefetch(self, location: Device | Host, *, stream: Stream | GraphBuilder) -> None: 

297 """Discard this range and prefetch to ``location`` on ``stream`` (CUDA 13+).""" 

298 _do_single_discard_prefetch_py(self, location, stream) 2 ctx1nm