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

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 

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  

16from dataclasses import dataclass (empty)

17import threading (empty)

18from typing import TYPE_CHECKING (empty)

19import warnings (empty)

20  

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

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

23  

24IF CUDA_CORE_BUILD_MAJOR >= 13: 

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

26  

27if TYPE_CHECKING: (empty)

28 from cuda.core.graph import GraphBuilder 

29  

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

31  

32  

33@dataclass (empty)

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 (empty)

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

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 | dict[str, object] | None = None) -> None: 

101 _MMR_init(self, options) 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

102  

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

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``, ``accessed_by``) 

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

122 ``discard_prefetch``). 

123 """ 

124 assert isinstance(stream, Stream), "Only Stream is supported for managed memory allocations" 21 ctx1VWXnopqYZ01234deigjkl

125 if self.is_mapped: 21 ctx1VWXnopqYZ01234deigjkl

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

127 cdef Stream s = Stream_accept(stream) 21 ctx1VWXnopqYZ01234deigjkl

128 return _MP_allocate(self, size, s, ManagedBuffer) 21 ctx1VWXnopqYZ01234deigjkl

129  

130 @property 

131 def device_id(self) -> int: 

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

133 if self._pref_loc_type == "device": 2 ctx1mg

134 return self._pref_loc_id 1 ctx1g

135 return -1 

136  

137 @property 

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

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

140  

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

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

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

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

145 """ 

146 if self._pref_loc_type is None: 4 ctx1rcbf

147 return None 1 ctx1r

148 if self._pref_loc_type == "host": 3 ctx1cbf

149 return (ManagedMemoryLocationType.HOST, None) 1 ctx1f

150 return (ManagedMemoryLocationType(self._pref_loc_type), self._pref_loc_id) 2 ctx1cb

151  

152 @property 

153 def is_device_accessible(self) -> bool: 

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

155 return True 3 ctx1dei

156  

157 @property 

158 def is_host_accessible(self) -> bool: 

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

160 return True 3 ctx1dei

161  

162 @property 

163 def is_managed(self) -> bool: 

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

165 return True 1 ctx1d

166  

167  

168IF CUDA_CORE_BUILD_MAJOR >= 13: 

169 cdef _resolve_preferred_location(ManagedMemoryResourceOptions opts): (empty)

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

171  

172 Returns a 4-tuple: 

173 (CUmemLocationType, loc_id, pref_loc_type_str, pref_loc_id) 

174 """ 

175 cdef object pref_loc = opts.preferred_location if opts is not None else None 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

176 cdef object pref_type = opts.preferred_location_type if opts is not None else None 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

177  

178 check_str_enum(pref_type, ManagedMemoryLocationType, allow_none=True) 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

179  

180 if pref_type is None: 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

181 # Legacy behavior 

182 if pref_loc is None: 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

183 return ( 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

184 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_NONE, 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

185 -1, None, -1, 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

186 ) 

187 if pref_loc == -1: 5 ctx1dcfag

188 return ( 1 ctx1f

189 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1f

190 -1, "host", -1, 

191 ) 

192 if pref_loc < 0: 4 ctx1dcag

193 raise ValueError( 1 ctx1a

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

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

196 ) 

197 return ( 3 ctx1dcg

198 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 3 ctx1dcg

199 pref_loc, "device", pref_loc, 3 ctx1dcg

200 ) 

201  

202 if pref_type == "device": 5 ctx1hcbfa

203 if pref_loc is None or pref_loc < 0: 2 ctx1ca

204 raise ValueError( 1 ctx1a

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

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

207 ) 

208 return ( 1 ctx1c

209 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 1 ctx1c

210 pref_loc, "device", pref_loc, 1 ctx1c

211 ) 

212  

213 if pref_type == "host": 4 ctx1hbfa

214 if pref_loc is not None: 2 ctx1fa

215 raise ValueError( 1 ctx1a

216 f"preferred_location must be None when " 1 ctx1a

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

218 ) 

219 return ( 1 ctx1f

220 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1f

221 -1, "host", -1, 

222 ) 

223  

224 # pref_type == "host_numa" 

225 if pref_loc is None: 3 ctx1hba

226 from .._device import Device 2 ctx1hb

227 dev = Device() 2 ctx1hb

228 numa_id = dev.properties.host_numa_id 2 ctx1hb

229 if numa_id < 0: 2 ctx1hb

230 raise RuntimeError( 1 ctx1h

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

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

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

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

235 ) 

236 return ( 1 ctx1b

237 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

238 numa_id, "host_numa", numa_id, 1 ctx1b

239 ) 

240 if pref_loc < 0: 2 ctx1ba

241 raise ValueError( 1 ctx1a

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

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

244 ) 

245 return ( 1 ctx1b

246 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

247 pref_loc, "host_numa", pref_loc, 1 ctx1b

248 ) 

249  

250  

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

252 IF CUDA_CORE_BUILD_MAJOR >= 13: 

253 cdef ManagedMemoryResourceOptions opts = check_or_create_options( 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

254 ManagedMemoryResourceOptions, options, "ManagedMemoryResource options", 

255 keep_none=True 

256 ) 

257 cdef cydriver.CUmemLocationType loc_type 

258 cdef int loc_id 

259  

260 loc_type, loc_id, self._pref_loc_type, self._pref_loc_id = ( 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

261 _resolve_preferred_location(opts) 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

262 ) 

263  

264 if opts is None: 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

265 try: 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

266 MP_init_current_pool( 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

267 self, 

268 loc_type, 

269 loc_id, 

270 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

271 ) 

272 except CUDAError as e: 

273 if "CUDA_ERROR_NOT_SUPPORTED" in str(e): 

274 from .._device import Device 

275 if not Device().properties.concurrent_managed_access: 

276 raise RuntimeError( 

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

278 "managed allocations (concurrent managed access is not " 

279 "available). Use " 

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

281 "to create a dedicated managed pool." 

282 ) from e 

283 raise 

284 else: 

285 MP_init_create_pool( 16 ctx1dcbfigsjtuvwkxyl

286 self, 

287 loc_type, 

288 loc_id, 

289 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

290 False, 16 ctx1dcbfigsjtuvwkxyl

291 0, 

292 ) 

293  

294 _check_concurrent_managed_access() 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

295 ELSE: 

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

297  

298  

299cdef bint _concurrent_access_warned = False (empty)

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

301  

302  

303cdef inline _check_concurrent_managed_access(): (empty)

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

305 global _concurrent_access_warned 

306 if _concurrent_access_warned: 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

307 return 46 ctx1zABCnoDpEqFGHIJKLMdhercbfaigsNjOtPuQvRwSkTxUyl

308  

309 cdef int c_concurrent = 0 1 ctx1e

310 with _concurrent_access_lock: 1 ctx1e

311 if _concurrent_access_warned: 1 ctx1e

312 return 

313  

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

315 # later, so any device will do. 

316 with nogil: 1 ctx1e

317 HANDLE_RETURN(cydriver.cuDeviceGetAttribute( 1 ctx1e

318 &c_concurrent, 

319 cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CONCURRENT_MANAGED_ACCESS, 

320 0)) 

321 if not c_concurrent: 1 ctx1e

322 warnings.warn( 

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

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

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

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

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

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

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

330 UserWarning, 

331 stacklevel=3 

332 ) 

333  

334 _concurrent_access_warned = True 1 ctx1e

335  

336  

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

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

339 global _concurrent_access_warned 

340 _concurrent_access_warned = False