Coverage for cuda/core/_memory/_managed_memory_resource.pyx: 93.16%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from __future__ import annotations 

6  

7from cuda.bindings cimport cydriver 

8  

9from cuda.core._memory._memory_pool cimport _MemPool, _MP_allocate, MP_check_open 

10from cuda.core._memory._memory_pool cimport MP_init_create_pool, MP_init_current_pool # no-cython-lint 

11from cuda.core._stream cimport Stream, Stream_accept 

12from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

13from cuda.core._utils.cuda_utils cimport check_or_create_options # no-cython-lint 

14from cuda.core._utils.cuda_utils import CUDAError # no-cython-lint (empty)

15  

16import cython 

17from dataclasses import dataclass (empty)

18import threading (empty)

19from typing import TYPE_CHECKING (empty)

20import warnings (empty)

21  

22from cuda.core._memory._managed_buffer import ManagedBuffer (empty)

23from cuda.core.typing import ManagedMemoryLocationType (empty)

24  

25IF CUDA_CORE_BUILD_MAJOR >= 13: 

26 from cuda.core._utils.validators import check_str_enum (empty)

27  

28if TYPE_CHECKING: (empty)

29 from cuda.core.graph import GraphBuilder 

30  

31__all__ = ['ManagedMemoryResource', 'ManagedMemoryResourceOptions'] (empty)

32  

33  

34@dataclass (empty)

35cdef class ManagedMemoryResourceOptions: 

36 """Customizable :obj:`~_memory.ManagedMemoryResource` options. 

37  

38 Attributes 

39 ---------- 

40 preferred_location : int | None, optional 

41 A location identifier (device ordinal or NUMA node ID) whose 

42 meaning depends on ``preferred_location_type``. 

43 (Default to ``None``) 

44  

45 preferred_location_type : ManagedMemoryLocationType | str | None, optional 

46 Controls how ``preferred_location`` is interpreted. 

47  

48 When set to ``None`` (the default), legacy behavior is used: 

49 ``preferred_location`` is interpreted as a device ordinal, 

50 ``-1`` for host, or ``None`` for no preference. 

51  

52 When set explicitly, the type determines both the kind of 

53 preferred location and the valid values for 

54 ``preferred_location``: 

55  

56 - ``"device"``: prefer a specific GPU. ``preferred_location`` 

57 must be a device ordinal (``>= 0``). 

58 - ``"host"``: prefer host memory (OS-managed NUMA placement). 

59 ``preferred_location`` must be ``None``. 

60 - ``"host_numa"``: prefer a specific host NUMA node. 

61 ``preferred_location`` must be a NUMA node ID (``>= 0``), 

62 or ``None`` to derive the NUMA node from the current CUDA 

63 device's ``host_numa_id`` attribute (requires an active 

64 CUDA context). 

65  

66 (Default to ``None``) 

67 """ 

68 preferred_location: int | None = None (empty)

69 preferred_location_type: ManagedMemoryLocationType | str | None = None (empty)

70  

71  

72cdef class ManagedMemoryResource(_MemPool): 

73 """ 

74 A managed memory resource managing a stream-ordered memory pool. 

75  

76 Managed memory is accessible from both the host and device, with automatic 

77 migration between them as needed. 

78  

79 Parameters 

80 ---------- 

81 options : ManagedMemoryResourceOptions 

82 Memory resource creation options. 

83  

84 If set to `None`, the memory resource uses the driver's current 

85 stream-ordered memory pool. If no memory pool is set as current, 

86 the driver's default memory pool is used. 

87  

88 If not set to `None`, a new memory pool is created, which is owned by 

89 the memory resource. 

90  

91 When using an existing (current or default) memory pool, the returned 

92 managed memory resource does not own the pool (`is_handle_owned` is 

93 `False`), and closing the resource has no effect. 

94  

95 Notes 

96 ----- 

97 IPC (Inter-Process Communication) is not currently supported for managed 

98 memory pools. 

99 """ 

100  

101 @cython.annotation_typing(False) 

102 def __init__(self, options: ManagedMemoryResourceOptions | None = None) -> None: 

103 _MMR_init(self, options) 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

104  

105 def allocate(self, size_t size, *, stream: Stream | GraphBuilder) -> ManagedBuffer: (empty)

106 """Allocate a managed-memory buffer of the requested size. 

107  

108 Parameters 

109 ---------- 

110 size : int 

111 The size of the buffer to allocate, in bytes. 

112 stream : :obj:`~_stream.Stream` 

113 Keyword-only. The stream on which to perform the allocation 

114 asynchronously. Must be passed explicitly; pass 

115 ``device.default_stream`` to use the default stream. 

116  

117 Returns 

118 ------- 

119 ManagedBuffer 

120 A :class:`ManagedBuffer` (a :class:`Buffer` subclass) that 

121 exposes the property-style advice API 

122 (``read_mostly``, ``preferred_location``, ``accessed_by``) 

123 and instance methods (``prefetch``, ``discard``, 

124 ``discard_prefetch``). 

125 """ 

126 MP_check_open(self) 28 ctx1phijklm345tuvw6789!#$cqnefrs

127 assert isinstance(stream, Stream), "Only Stream is supported for managed memory allocations" 28 ctx1phijklm345tuvw6789!#$cqnefrs

128 if self.is_mapped: 28 ctx1phijklm345tuvw6789!#$cqnefrs

129 raise TypeError("Cannot allocate from a mapped IPC-enabled memory resource") 

130 cdef Stream s = Stream_accept(stream) 28 ctx1phijklm345tuvw6789!#$cqnefrs

131 return _MP_allocate(self, size, s, ManagedBuffer) 28 ctx1phijklm345tuvw6789!#$cqnefrs

132  

133 @property 

134 def device_id(self) -> int: 

135 """The preferred device ordinal, or -1 if the preferred location is not a device.""" 

136 if self._pref_loc_type == "device": 1 ctx1e

137 return self._pref_loc_id 1 ctx1e

138 return -1 

139  

140 @property 

141 def preferred_location(self) -> tuple[ManagedMemoryLocationType, int | None] | None: 

142 """The preferred location for managed memory allocations. 

143  

144 Returns ``None`` if no preferred location is set (driver decides), 

145 or a tuple ``(type, id)`` where *type* is one of ``"device"``, 

146 ``"host"``, or ``"host_numa"``, and *id* is the device ordinal, 

147 ``None`` (for ``"host"``), or the NUMA node ID, respectively. 

148 """ 

149 if self._pref_loc_type is None: 4 ctx1ydbg

150 return None 1 ctx1y

151 if self._pref_loc_type == "host": 3 ctx1dbg

152 return (ManagedMemoryLocationType.HOST, None) 1 ctx1g

153 return (ManagedMemoryLocationType(self._pref_loc_type), self._pref_loc_id) 2 ctx1db

154  

155 @property 

156 def is_device_accessible(self) -> bool: 

157 """Return True. This memory resource provides device-accessible buffers.""" 

158 return True 3 ctx1cqn

159  

160 @property 

161 def is_host_accessible(self) -> bool: 

162 """Return True. This memory resource provides host-accessible buffers.""" 

163 return True 3 ctx1cqn

164  

165 @property 

166 def is_managed(self) -> bool: 

167 """Return True. This memory resource provides managed (unified) memory buffers.""" 

168 return True 1 ctx1c

169  

170  

171IF CUDA_CORE_BUILD_MAJOR >= 13: 

172 cdef _resolve_preferred_location(ManagedMemoryResourceOptions opts): (empty)

173 """Resolve preferred location options into driver and stored values. 

174  

175 Returns a 4-tuple: 

176 (CUmemLocationType, loc_id, pref_loc_type_str, pref_loc_id) 

177 """ 

178 cdef object pref_loc = opts.preferred_location if opts is not None else None 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

179 cdef object pref_type = opts.preferred_location_type if opts is not None else None 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

180  

181 check_str_enum(pref_type, ManagedMemoryLocationType, allow_none=True) 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

182  

183 if pref_type is None: 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

184 # Legacy behavior 

185 if pref_loc is None: 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

186 return ( 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

187 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_NONE, 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

188 -1, None, -1, 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

189 ) 

190 if pref_loc == -1: 5 ctx1cdgae

191 return ( 1 ctx1g

192 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1g

193 -1, "host", -1, 

194 ) 

195 if pref_loc < 0: 4 ctx1cdae

196 raise ValueError( 1 ctx1a

197 f"preferred_location must be a device ordinal (>= 0), -1 for " 1 ctx1a

198 f"host, or None for no preference, got {pref_loc}" 1 ctx1a

199 ) 

200 return ( 3 ctx1cde

201 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 3 ctx1cde

202 pref_loc, "device", pref_loc, 3 ctx1cde

203 ) 

204  

205 if pref_type == "device": 5 ctx1odbga

206 if pref_loc is None or pref_loc < 0: 2 ctx1da

207 raise ValueError( 1 ctx1a

208 f"preferred_location must be a device ordinal (>= 0) when " 1 ctx1a

209 f"preferred_location_type is 'device', got {pref_loc!r}" 1 ctx1a

210 ) 

211 return ( 1 ctx1d

212 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 1 ctx1d

213 pref_loc, "device", pref_loc, 1 ctx1d

214 ) 

215  

216 if pref_type == "host": 4 ctx1obga

217 if pref_loc is not None: 2 ctx1ga

218 raise ValueError( 1 ctx1a

219 f"preferred_location must be None when " 1 ctx1a

220 f"preferred_location_type is 'host', got {pref_loc!r}" 1 ctx1a

221 ) 

222 return ( 1 ctx1g

223 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1g

224 -1, "host", -1, 

225 ) 

226  

227 # pref_type == "host_numa" 

228 if pref_loc is None: 3 ctx1oba

229 from .._device import Device 2 ctx1ob

230 dev = Device() 2 ctx1ob

231 numa_id = dev.properties.host_numa_id 2 ctx1ob

232 if numa_id < 0: 2 ctx1ob

233 raise RuntimeError( 1 ctx1o

234 "Cannot determine host NUMA ID for the current CUDA device. " 

235 "The system may not support NUMA, or no CUDA context is " 

236 "active. Set preferred_location to an explicit NUMA node ID " 

237 "or call Device.set_current() first." 

238 ) 

239 return ( 1 ctx1b

240 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

241 numa_id, "host_numa", numa_id, 1 ctx1b

242 ) 

243 if pref_loc < 0: 2 ctx1ba

244 raise ValueError( 1 ctx1a

245 f"preferred_location must be a NUMA node ID (>= 0) or None " 1 ctx1a

246 f"when preferred_location_type is 'host_numa', got {pref_loc}" 1 ctx1a

247 ) 

248 return ( 1 ctx1b

249 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

250 pref_loc, "host_numa", pref_loc, 1 ctx1b

251 ) 

252  

253  

254cdef inline _MMR_init(ManagedMemoryResource self, options): (empty)

255 IF CUDA_CORE_BUILD_MAJOR >= 13: 

256 cdef ManagedMemoryResourceOptions opts = check_or_create_options( 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

257 ManagedMemoryResourceOptions, options, "ManagedMemoryResource options", 

258 keep_none=True 

259 ) 

260 cdef cydriver.CUmemLocationType loc_type 

261 cdef int loc_id 

262  

263 loc_type, loc_id, self._pref_loc_type, self._pref_loc_id = ( 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

264 _resolve_preferred_location(opts) 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

265 ) 

266  

267 if opts is None: 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

268 try: 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

269 MP_init_current_pool( 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

270 self, 

271 loc_type, 

272 loc_id, 

273 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

274 ) 

275 except CUDAError as e: 6 ctx1hijklm

276 if "CUDA_ERROR_NOT_SUPPORTED" in str(e): 6 ctx1hijklm

277 from .._device import Device 6 ctx1hijklm

278 if not Device().properties.concurrent_managed_access: 6 ctx1hijklm

279 raise RuntimeError( 6 ctx1hijklm

280 "The default memory pool on this device does not support " 

281 "managed allocations (concurrent managed access is not " 

282 "available). Use " 

283 "ManagedMemoryResource(options=ManagedMemoryResourceOptions(...)) " 

284 "to create a dedicated managed pool." 

285 ) from e 6 ctx1hijklm

286 raise 

287 else: 

288 MP_init_create_pool( 17 ctx1pcdbgnezfABCDrEFs

289 self, 

290 loc_type, 

291 loc_id, 

292 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

293 False, 17 ctx1pcdbgnezfABCDrEFs

294 0, 

295 ) 

296  

297 _check_concurrent_managed_access() 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

298 ELSE: 

299 raise RuntimeError("ManagedMemoryResource requires CUDA 13.0 or later") 

300  

301  

302cdef bint _concurrent_access_warned = False (empty)

303cdef object _concurrent_access_lock = threading.Lock() (empty)

304  

305  

306cdef inline _check_concurrent_managed_access(): (empty)

307 """Warn once if the platform lacks concurrent managed memory access.""" 

308 global _concurrent_access_warned 

309 if _concurrent_access_warned: 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

310 return 54 ctx1GphijklmHIJKtuLvMwNOPQRSTUcoqydbganezVfWAXBYCZD0r1E2Fs

311  

312 cdef int c_concurrent = 0 1 ctx1f

313 with _concurrent_access_lock: 1 ctx1f

314 if _concurrent_access_warned: 1 ctx1f

315 return 

316  

317 # concurrent_managed_access is a system-level attribute for sm_60 and 

318 # later, so any device will do. 

319 with nogil: 1 ctx1f

320 HANDLE_RETURN(cydriver.cuDeviceGetAttribute( 1 ctx1f

321 &c_concurrent, 

322 cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CONCURRENT_MANAGED_ACCESS, 

323 0)) 

324 if not c_concurrent: 1 ctx1f

325 warnings.warn( 

326 "This platform does not support concurrent managed memory access " 

327 "(Device.properties.concurrent_managed_access is False). Host access to any managed " 

328 "allocation is forbidden while any GPU kernel is in flight, even " 

329 "if the kernel does not touch that allocation. Failing to " 

330 "synchronize before host access will cause a segfault. " 

331 "See: https://docs.nvidia.com/cuda/cuda-c-programming-guide/" 

332 "index.html#gpu-exclusive-access-to-managed-memory", 

333 UserWarning, 

334 stacklevel=3 

335 ) 

336  

337 _concurrent_access_warned = True 1 ctx1f

338  

339  

340def reset_concurrent_access_warning() -> None: (empty)

341 """Reset the concurrent access warning flag for testing purposes.""" 

342 global _concurrent_access_warned 

343 _concurrent_access_warned = False