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

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) 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

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) 29 ctx1enopqrs345tuvwx6789!#$cjhfklm

125 assert isinstance(stream, Stream), "Only Stream is supported for managed memory allocations" 29 ctx1enopqrs345tuvwx6789!#$cjhfklm

126 if self.is_mapped: 29 ctx1enopqrs345tuvwx6789!#$cjhfklm

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

128 cdef Stream s = Stream_accept(stream) 29 ctx1enopqrs345tuvwx6789!#$cjhfklm

129 return _MP_allocate(self, size, s, ManagedBuffer) 29 ctx1enopqrs345tuvwx6789!#$cjhfklm

130  

131 @property 

132 def device_id(self) -> int: 

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

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

135 return self._pref_loc_id 1 ctx1f

136 return -1 

137  

138 @property 

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

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

141  

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

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

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

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

146 """ 

147 if self._pref_loc_type is None: 4 ctx1ydbg

148 return None 1 ctx1y

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

150 return (ManagedMemoryLocationType.HOST, None) 1 ctx1g

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

152  

153 @property 

154 def is_device_accessible(self) -> bool: 

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

156 return True 3 ctx1cjh

157  

158 @property 

159 def is_host_accessible(self) -> bool: 

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

161 return True 3 ctx1cjh

162  

163 @property 

164 def is_managed(self) -> bool: 

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

166 return True 1 ctx1c

167  

168  

169IF CUDA_CORE_BUILD_MAJOR >= 13: 

170 cdef _resolve_preferred_location(ManagedMemoryResourceOptions opts): 

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

172  

173 Returns a 4-tuple: 

174 (CUmemLocationType, loc_id, pref_loc_type_str, pref_loc_id) 

175 """ 

176 cdef object pref_loc = opts.preferred_location if opts is not None else None 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

177 cdef object pref_type = opts.preferred_location_type if opts is not None else None 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

178  

179 check_str_enum(pref_type, ManagedMemoryLocationType, allow_none=True) 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

180  

181 if pref_type is None: 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

182 # Legacy behavior 

183 if pref_loc is None: 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

184 return ( 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

185 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_NONE, 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

186 -1, None, -1, 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

187 ) 

188 if pref_loc == -1: 5 ctx1cdgaf

189 return ( 1 ctx1g

190 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1g

191 -1, "host", -1, 

192 ) 

193 if pref_loc < 0: 4 ctx1cdaf

194 raise ValueError( 1 ctx1a

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

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

197 ) 

198 return ( 3 ctx1cdf

199 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 3 ctx1cdf

200 pref_loc, "device", pref_loc, 3 ctx1cdf

201 ) 

202  

203 if pref_type == "device": 5 ctx1idbga

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

205 raise ValueError( 1 ctx1a

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

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

208 ) 

209 return ( 1 ctx1d

210 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 1 ctx1d

211 pref_loc, "device", pref_loc, 1 ctx1d

212 ) 

213  

214 if pref_type == "host": 4 ctx1ibga

215 if pref_loc is not None: 2 ctx1ga

216 raise ValueError( 1 ctx1a

217 f"preferred_location must be None when " 1 ctx1a

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

219 ) 

220 return ( 1 ctx1g

221 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1g

222 -1, "host", -1, 

223 ) 

224  

225 # pref_type == "host_numa" 

226 if pref_loc is None: 3 ctx1iba

227 from .._device import Device 2 ctx1ib

228 dev = Device() 2 ctx1ib

229 numa_id = dev.properties.host_numa_id 2 ctx1ib

230 if numa_id < 0: 2 ctx1ib

231 raise RuntimeError( 1 ctx1i

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

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

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

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

236 ) 

237 return ( 1 ctx1b

238 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

239 numa_id, "host_numa", numa_id, 1 ctx1b

240 ) 

241 if pref_loc < 0: 2 ctx1ba

242 raise ValueError( 1 ctx1a

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

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

245 ) 

246 return ( 1 ctx1b

247 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

248 pref_loc, "host_numa", pref_loc, 1 ctx1b

249 ) 

250  

251  

252cdef inline _MMR_init(ManagedMemoryResource self, options): 

253 IF CUDA_CORE_BUILD_MAJOR >= 13: 

254 cdef ManagedMemoryResourceOptions opts = check_or_create_options( 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

255 ManagedMemoryResourceOptions, options, "ManagedMemoryResource options", 

256 keep_none=True 

257 ) 

258 cdef cydriver.CUmemLocationType loc_type 

259 cdef int loc_id 

260  

261 loc_type, loc_id, self._pref_loc_type, self._pref_loc_id = ( 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

262 _resolve_preferred_location(opts) 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

263 ) 

264  

265 if opts is None: 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

266 try: 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

267 MP_init_current_pool( 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

268 self, 

269 loc_type, 

270 loc_id, 

271 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

272 ) 

273 except CUDAError as e: 

274 if "CUDA_ERROR_NOT_SUPPORTED" in str(e): 

275 from .._device import Device 

276 if not Device().properties.concurrent_managed_access: 

277 raise RuntimeError( 

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

279 "managed allocations (concurrent managed access is not " 

280 "available). Use " 

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

282 "to create a dedicated managed pool." 

283 ) from e 

284 raise 

285 else: 

286 MP_init_create_pool( 17 ctx1ecdbghfzkABCDlEFm

287 self, 

288 loc_type, 

289 loc_id, 

290 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

291 False, 17 ctx1ecdbghfzkABCDlEFm

292 0, 

293 ) 

294  

295 _check_concurrent_managed_access() 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

296 ELSE: 

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

298  

299  

300cdef bint _concurrent_access_warned = False 

301cdef object _concurrent_access_lock = threading.Lock() 

302  

303  

304cdef inline _check_concurrent_managed_access(): 

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

306 global _concurrent_access_warned 

307 if _concurrent_access_warned: 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

308 return 55 ctx1enopqrsGHIJtuKvLwMxNOPQRSTUcijydbgahfzVkWAXBYCZD0l1E2Fm

309  

310 cdef int c_concurrent = 0 1 ctx1e

311 with _concurrent_access_lock: 1 ctx1e

312 if _concurrent_access_warned: 1 ctx1e

313 return 

314  

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

316 # later, so any device will do. 

317 with nogil: 1 ctx1e

318 HANDLE_RETURN(cydriver.cuDeviceGetAttribute( 1 ctx1e

319 &c_concurrent, 

320 cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CONCURRENT_MANAGED_ACCESS, 

321 0)) 

322 if not c_concurrent: 1 ctx1e

323 warnings.warn( 

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

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

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

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

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

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

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

331 UserWarning, 

332 stacklevel=3 

333 ) 

334  

335 _concurrent_access_warned = True 1 ctx1e

336  

337  

338def reset_concurrent_access_warning() -> None: 

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

340 global _concurrent_access_warned 

341 _concurrent_access_warned = False