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

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4from __future__ import annotations (empty)

5 

6from collections.abc import Iterable, Iterator, MutableSet (empty)

7from typing import TYPE_CHECKING, Any (empty)

8 

9from cuda.core._device import Device (empty)

10from cuda.core._host import Host (empty)

11from cuda.core._memory._buffer import Buffer (empty)

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

13from cuda.core._memory._managed_memory_ops import ( (empty)

14 _advise_one, 

15 _do_single_discard_prefetch_py, 

16 _do_single_discard_py, 

17 _do_single_prefetch_py, 

18 _read_preferred_location_v2, 

19) 

20from cuda.core._utils.cuda_utils import driver, handle_return (empty)

21from cuda.core._utils.version import binding_version, driver_version (empty)

22 

23if TYPE_CHECKING: (empty)

24 from cuda.core._memory._buffer import MemoryResource 

25 from cuda.core._stream import Stream 

26 from cuda.core.graph import GraphBuilder 

27 

28__all__ = ["ManagedBuffer"] (empty)

29 

30 

31_INT_SIZE = 4 (empty)

32 

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

34_ADV = driver.CUmem_advise (empty)

35_SET_READ_MOSTLY = _ADV.CU_MEM_ADVISE_SET_READ_MOSTLY (empty)

36_UNSET_READ_MOSTLY = _ADV.CU_MEM_ADVISE_UNSET_READ_MOSTLY (empty)

37_SET_PREFERRED = _ADV.CU_MEM_ADVISE_SET_PREFERRED_LOCATION (empty)

38_UNSET_PREFERRED = _ADV.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION (empty)

39_SET_ACCESSED_BY = _ADV.CU_MEM_ADVISE_SET_ACCESSED_BY (empty)

40_UNSET_ACCESSED_BY = _ADV.CU_MEM_ADVISE_UNSET_ACCESSED_BY (empty)

41 

42_RANGE = driver.CUmem_range_attribute (empty)

43_ATTR_READ_MOSTLY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_READ_MOSTLY (empty)

44_ATTR_PREFERRED = _RANGE.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION (empty)

45_ATTR_ACCESSED_BY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_ACCESSED_BY (empty)

46 

47 

48def _check_open(buf: Buffer) -> None: (empty)

49 if buf.is_closed: 14 ctx1pedbcfimnghjqr

50 raise RuntimeError("Buffer has been closed") 

51 

52 

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

54 _check_open(buf) 7 ctx1pimnjqr

55 return int(handle_return(driver.cuMemRangeGetAttribute(_INT_SIZE, attribute, buf.handle, buf.size))) 7 ctx1pimnjqr

56 

57 

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

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

60 

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

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

63 """ 

64 _check_open(buf) 4 ctx1edbc

65 num_devices = handle_return(driver.cuDeviceGetCount()) 4 ctx1edbc

66 n = num_devices + 1 4 ctx1edbc

67 raw = handle_return(driver.cuMemRangeGetAttribute(n * _INT_SIZE, _ATTR_ACCESSED_BY, buf.handle, buf.size)) 4 ctx1edbc

68 return [Host() if v == -1 else Device(v) for v in raw if v != -2] 4 ctx1edbc

69 

70 

71class AccessedBySetProxy(MutableSet[Device | Host]): (empty)

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

73 

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

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

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

77 reflects the current driver state. 

78 

79 Note 

80 ---- 

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

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

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

84 """ 

85 

86 __slots__ = ("_buf",) (empty)

87 

88 def __init__(self, buf: ManagedBuffer): (empty)

89 self._buf = buf 7 ctx1edbcfgh

90 

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

92 @classmethod (empty)

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

94 return set(it) 1 ctx1d

95 

96 # --- abstract methods required by MutableSet --- 

97 

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

99 if not isinstance(location, (Device, Host)): 4 ctx1edbc

100 return False 

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

102 

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

104 return iter(_query_accessed_by(self._buf)) 1 ctx1d

105 

106 def __len__(self) -> int: (empty)

107 return len(_query_accessed_by(self._buf)) 1 ctx1d

108 

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

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

111 if not isinstance(location, (Device, Host)): 5 ctx1edfgh

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

113 _advise_one(self._buf, _SET_ACCESSED_BY, location) 5 ctx1edfgh

114 

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

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

117 

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

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

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

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

122 is silently ignored rather than forwarded to the driver. 

123 """ 

124 if not isinstance(location, (Device, Host)): 2 ctx1ed

125 return 

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

127 return 1 ctx1d

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

129 

130 def __repr__(self) -> str: (empty)

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

132 

133 

134class ManagedBuffer(Buffer): (empty)

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

136 

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

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

139 

140 Examples 

141 -------- 

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

143 >>> buf.read_mostly = True 

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

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

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

147 

148 Note 

149 ---- 

150 On CUDA 13 builds, ``preferred_location`` round-trips full NUMA 

151 information. On CUDA 12 builds, ``Host(numa_id=...)`` and 

152 ``Host.numa_current()`` are rejected with ``TypeError`` at the call 

153 boundary — only ``Device(...)`` and the generic ``Host()`` are 

154 accepted. Use ``Host()`` to target the host on CUDA 12. 

155 """ 

156 

157 @classmethod (empty)

158 def from_handle( (empty)

159 cls, 

160 ptr, 

161 size: int, 

162 mr: MemoryResource | None = None, 

163 owner: object | None = None, 

164 *, 

165 stream: Stream | GraphBuilder | None = None, 

166 ) -> Buffer: 

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

168 

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

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

171 :attr:`preferred_location`, :attr:`accessed_by`). 

172 

173 Parameters 

174 ---------- 

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

176 Pointer to a managed allocation. 

177 size : int 

178 Allocation size in bytes. 

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

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

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

182 owner : object, optional 

183 An object that keeps the underlying allocation alive. 

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

185 stream : Stream | GraphBuilder, optional 

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

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

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

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

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

191 refers to a different stream on each thread. 

192 """ 

193 return cls._init(ptr, size, mr=mr, owner=owner, stream=stream) 18 ctx1uvwxyzABCkDEoFGHIJ

194 

195 @property (empty)

196 def read_mostly(self) -> bool: (empty)

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

198 return _get_int_attr(self, _ATTR_READ_MOSTLY) != 0 1 ctx1j

199 

200 @read_mostly.setter (empty)

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

202 _advise_one(self, _SET_READ_MOSTLY if value else _UNSET_READ_MOSTLY, None) 3 ctx1fjo

203 

204 @property (empty)

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

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

207 

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

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

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

211 as ``Host()``. 

212 """ 

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

214 # both added in CUDA 13. Require both bindings and the runtime driver 

215 # to be 13.0+; otherwise fall back to the legacy device-ordinal path. 

216 # See PR #2054 / #2064 for prior bindings-only-check regressions. 

217 if binding_version() >= (13, 0, 0) and driver_version() >= (13, 0, 0): 2 ctx1kl

218 return _read_preferred_location_v2(self) 2 ctx1kl

219 # CUDA 12 legacy path (no NUMA info available; also taken when 

220 # bindings are 13.x but the runtime driver is still 12.x). 

221 loc_id = _get_int_attr(self, _ATTR_PREFERRED) 

222 if loc_id == -2: 

223 return None 

224 if loc_id == -1: 

225 return Host() 

226 return Device(loc_id) 

227 

228 @preferred_location.setter (empty)

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

230 if value is None: 3 ctx1fkl

231 _advise_one(self, _UNSET_PREFERRED, None) 1 ctx1l

232 else: 

233 _advise_one(self, _SET_PREFERRED, value) 3 ctx1fkl

234 

235 @property (empty)

236 def accessed_by(self) -> AccessedBySetProxy: (empty)

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

238 _check_open(self) 7 ctx1edbcfgh

239 return AccessedBySetProxy(self) 7 ctx1edbcfgh

240 

241 @accessed_by.setter (empty)

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

243 _check_open(self) 2 ctx1bc

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

245 # element can't leave accessed_by partially mutated. 

246 target: set[Device | Host] = set() 2 ctx1bc

247 for loc in locations: 2 ctx1bc

248 if not isinstance(loc, (Device, Host)): 2 ctx1bc

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

250 target.add(loc) 2 ctx1bc

251 for loc in target: 2 ctx1bc

252 spec = _coerce_location(loc) 2 ctx1bc

253 assert spec is not None 2 ctx1bc

254 if spec.kind not in ("device", "host"): 2 ctx1bc

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

256 current = set(_query_accessed_by(self)) 2 ctx1bc

257 for loc in current - target: 2 ctx1bc

258 _advise_one(self, _UNSET_ACCESSED_BY, loc) 1 ctx1c

259 for loc in target - current: 2 ctx1bc

260 _advise_one(self, _SET_ACCESSED_BY, loc) 2 ctx1bc

261 

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

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

264 _do_single_prefetch_py(self, location, stream) 8 ctx1stimngho

265 

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

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

268 _do_single_discard_py(self, stream) 2 ctx1st

269 

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

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

272 _do_single_discard_prefetch_py(self, location, stream) 2 ctx1io