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

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

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 

15  

16from dataclasses import dataclass 

17import threading 

18from typing import TYPE_CHECKING 

19import warnings 

20  

21from cuda.core._memory._managed_buffer import ManagedBuffer 

22from cuda.core.typing import ManagedMemoryLocationType 

23  

24IF CUDA_CORE_BUILD_MAJOR >= 13: 

25 from cuda.core._utils.validators import check_str_enum 

26  

27if TYPE_CHECKING: 

28 from cuda.core.graph import GraphBuilder 

29  

30__all__ = ['ManagedMemoryResource', 'ManagedMemoryResourceOptions'] 

31  

32  

33@dataclass 

34cdef class ManagedMemoryResourceOptions: 

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

36  

37 Attributes 

38 ---------- 

39 preferred_location : int | None, optional 

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

41 meaning depends on ``preferred_location_type``. 

42 (Default to ``None``) 

43  

44 preferred_location_type : ManagedMemoryLocationType | str | None, optional 

45 Controls how ``preferred_location`` is interpreted. 

46  

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

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

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

50  

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

52 preferred location and the valid values for 

53 ``preferred_location``: 

54  

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

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

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

58 ``preferred_location`` must be ``None``. 

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

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

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

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

63 CUDA context). 

64  

65 (Default to ``None``) 

66 """ 

67 preferred_location: int | None = None 

68 preferred_location_type: ManagedMemoryLocationType | str | None = None 

69  

70  

71cdef class ManagedMemoryResource(_MemPool): 

72 """ 

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

74  

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

76 migration between them as needed. 

77  

78 Parameters 

79 ---------- 

80 options : ManagedMemoryResourceOptions 

81 Memory resource creation options. 

82  

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

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

85 the driver's default memory pool is used. 

86  

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

88 the memory resource. 

89  

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

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

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

93  

94 Notes 

95 ----- 

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

97 memory pools. 

98 """ 

99  

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

101 _MMR_init(self, options) 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

102  

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

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

105  

106 Parameters 

107 ---------- 

108 size : int 

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

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

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

112 asynchronously. Must be passed explicitly; pass 

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

114  

115 Returns 

116 ------- 

117 ManagedBuffer 

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

119 exposes the property-style advice API 

120 (``read_mostly``, ``preferred_location``, 

121 ``last_prefetch_location``, ``accessed_by``) and instance methods 

122 (``prefetch``, ``discard``, ``discard_prefetch``). 

123 """ 

124 MP_check_open(self) 33 ctx1erstuvw'()xyzAB*+,-./:hckiflmnopq

125 if not isinstance(stream, Stream): 33 ctx1erstuvw'()xyzAB*+,-./:hckiflmnopq

126 raise TypeError(f"Only Stream is supported for managed memory allocations, got {type(stream).__name__}") 1 ctx1h

127 if self.is_mapped: 32 ctx1erstuvw'()xyzAB*+,-./:ckiflmnopq

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

129 cdef Stream s = Stream_accept(stream) 32 ctx1erstuvw'()xyzAB*+,-./:ckiflmnopq

130 return _MP_allocate(self, size, s, ManagedBuffer) 32 ctx1erstuvw'()xyzAB*+,-./:ckiflmnopq

131  

132 @property 

133 def device_id(self) -> int: 

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

135 if self._pref_loc_type == "device": 1 ctx1f

136 return self._pref_loc_id 1 ctx1f

137 return -1 

138  

139 @property 

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

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

142  

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

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

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

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

147 """ 

148 if self._pref_loc_type is None: 4 ctx1Cdbg

149 return None 1 ctx1C

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

151 return (ManagedMemoryLocationType.HOST, None) 1 ctx1g

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

153  

154 @property 

155 def is_device_accessible(self) -> bool: 

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

157 return True 3 ctx1cki

158  

159 @property 

160 def is_host_accessible(self) -> bool: 

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

162 return True 3 ctx1cki

163  

164 @property 

165 def is_managed(self) -> bool: 

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

167 return True 1 ctx1c

168  

169  

170IF CUDA_CORE_BUILD_MAJOR >= 13: 

171 cdef _resolve_preferred_location(ManagedMemoryResourceOptions opts): 

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

173  

174 Returns a 4-tuple: 

175 (CUmemLocationType, loc_id, pref_loc_type_str, pref_loc_id) 

176 """ 

177 cdef object pref_loc = opts.preferred_location if opts is not None else None 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

178 cdef object pref_type = opts.preferred_location_type if opts is not None else None 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

179  

180 check_str_enum(pref_type, ManagedMemoryLocationType, allow_none=True) 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

181  

182 if pref_type is None: 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

183 # Legacy behavior 

184 if pref_loc is None: 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

185 return ( 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

186 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_NONE, 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

187 -1, None, -1, 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

188 ) 

189 if pref_loc == -1: 6 ctx1hcdgaf

190 return ( 1 ctx1g

191 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1g

192 -1, "host", -1, 

193 ) 

194 if pref_loc < 0: 5 ctx1hcdaf

195 raise ValueError( 1 ctx1a

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

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

198 ) 

199 return ( 4 ctx1hcdf

200 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 4 ctx1hcdf

201 pref_loc, "device", pref_loc, 4 ctx1hcdf

202 ) 

203  

204 if pref_type == "device": 5 ctx1jdbga

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

206 raise ValueError( 1 ctx1a

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

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

209 ) 

210 return ( 1 ctx1d

211 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 1 ctx1d

212 pref_loc, "device", pref_loc, 1 ctx1d

213 ) 

214  

215 if pref_type == "host": 4 ctx1jbga

216 if pref_loc is not None: 2 ctx1ga

217 raise ValueError( 1 ctx1a

218 f"preferred_location must be None when " 1 ctx1a

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

220 ) 

221 return ( 1 ctx1g

222 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1g

223 -1, "host", -1, 

224 ) 

225  

226 # pref_type == "host_numa" 

227 if pref_loc is None: 3 ctx1jba

228 from .._device import Device 2 ctx1jb

229 dev = Device() 2 ctx1jb

230 numa_id = dev.properties.host_numa_id 2 ctx1jb

231 if numa_id < 0: 2 ctx1jb

232 raise RuntimeError( 1 ctx1j

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

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

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

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

237 ) 

238 return ( 1 ctx1b

239 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

240 numa_id, "host_numa", numa_id, 1 ctx1b

241 ) 

242 if pref_loc < 0: 2 ctx1ba

243 raise ValueError( 1 ctx1a

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

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

246 ) 

247 return ( 1 ctx1b

248 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

249 pref_loc, "host_numa", pref_loc, 1 ctx1b

250 ) 

251  

252  

253cdef inline _MMR_init(ManagedMemoryResource self, options): 

254 IF CUDA_CORE_BUILD_MAJOR >= 13: 

255 cdef ManagedMemoryResourceOptions opts = check_or_create_options( 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

256 ManagedMemoryResourceOptions, options, "ManagedMemoryResource options", 

257 keep_none=True 

258 ) 

259 cdef cydriver.CUmemLocationType loc_type 

260 cdef int loc_id 

261  

262 loc_type, loc_id, self._pref_loc_type, self._pref_loc_id = ( 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

263 _resolve_preferred_location(opts) 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

264 ) 

265  

266 if opts is None: 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

267 try: 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

268 MP_init_current_pool( 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

269 self, 

270 loc_type, 

271 loc_id, 

272 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

273 ) 

274 except CUDAError as e: 

275 if "CUDA_ERROR_NOT_SUPPORTED" in str(e): 

276 from .._device import Device 

277 if not Device().properties.concurrent_managed_access: 

278 raise RuntimeError( 

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

280 "managed allocations (concurrent managed access is not " 

281 "available). Use " 

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

283 "to create a dedicated managed pool." 

284 ) from e 

285 raise 

286 else: 

287 MP_init_create_pool( 28 ctx1ehcdbgifDlEFGHmIJnKLMNOPQopq

288 self, 

289 loc_type, 

290 loc_id, 

291 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

292 False, 28 ctx1ehcdbgifDlEFGHmIJnKLMNOPQopq

293 0, 

294 ) 

295  

296 _check_concurrent_managed_access() 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

297 ELSE: 

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

299  

300  

301cdef bint _concurrent_access_warned = False 

302cdef object _concurrent_access_lock = threading.Lock() 

303  

304  

305cdef inline _check_concurrent_managed_access(): 

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

307 global _concurrent_access_warned 

308 if _concurrent_access_warned: 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

309 return 66 ctx1erstuvwRSTUxyVzWAXBYZ012345hcjkCdbgaifD6l7E8F9G!H#m$I%JnKLMNOPQopq

310  

311 cdef int c_concurrent = 0 1 ctx1e

312 with _concurrent_access_lock: 1 ctx1e

313 if _concurrent_access_warned: 1 ctx1e

314 return 

315  

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

317 # later, so any device will do. 

318 with nogil: 1 ctx1e

319 HANDLE_RETURN(cydriver.cuDeviceGetAttribute( 1 ctx1e

320 &c_concurrent, 

321 cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CONCURRENT_MANAGED_ACCESS, 

322 0)) 

323 if not c_concurrent: 1 ctx1e

324 warnings.warn( 

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

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

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

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

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

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

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

332 UserWarning, 

333 stacklevel=3 

334 ) 

335  

336 _concurrent_access_warned = True 1 ctx1e

337  

338  

339def reset_concurrent_access_warning() -> None: 

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

341 global _concurrent_access_warned 

342 _concurrent_access_warned = False