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

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 

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

11from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

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

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

14  

15from dataclasses import dataclass (empty)

16import threading (empty)

17import warnings (empty)

18  

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

20  

21  

22@dataclass (empty)

23cdef class ManagedMemoryResourceOptions: 

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

25  

26 Attributes 

27 ---------- 

28 preferred_location : int | None, optional 

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

30 meaning depends on ``preferred_location_type``. 

31 (Default to ``None``) 

32  

33 preferred_location_type : ``"device"`` | ``"host"`` | ``"host_numa"`` | None, optional 

34 Controls how ``preferred_location`` is interpreted. 

35  

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

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

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

39  

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

41 preferred location and the valid values for 

42 ``preferred_location``: 

43  

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

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

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

47 ``preferred_location`` must be ``None``. 

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

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

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

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

52 CUDA context). 

53  

54 (Default to ``None``) 

55 """ 

56 preferred_location: int | None = None (empty)

57 preferred_location_type: str | None = None (empty)

58  

59  

60cdef class ManagedMemoryResource(_MemPool): 

61 """ 

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

63  

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

65 migration between them as needed. 

66  

67 Parameters 

68 ---------- 

69 options : ManagedMemoryResourceOptions 

70 Memory resource creation options. 

71  

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

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

74 the driver's default memory pool is used. 

75  

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

77 the memory resource. 

78  

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

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

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

82  

83 Notes 

84 ----- 

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

86 memory pools. 

87 """ 

88  

89 def __init__(self, options=None): 

90 _MMR_init(self, options) 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

91  

92 @property 

93 def device_id(self) -> int: 

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

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

96 return self._pref_loc_id 1 ctx1f

97 return -1 

98  

99 @property 

100 def preferred_location(self) -> tuple | None: 

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

102  

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

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

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

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

107 """ 

108 if self._pref_loc_type is None: 4 ctx1jcbd

109 return None 1 ctx1j

110 if self._pref_loc_type == "host": 3 ctx1cbd

111 return ("host", None) 1 ctx1d

112 return (self._pref_loc_type, self._pref_loc_id) 2 ctx1cb

113  

114 @property 

115 def is_device_accessible(self) -> bool: 

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

117 return True 3 ctx1eih

118  

119 @property 

120 def is_host_accessible(self) -> bool: 

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

122 return True 3 ctx1eih

123  

124 @property 

125 def is_managed(self) -> bool: 

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

127 return True 1 ctx1e

128  

129  

130IF CUDA_CORE_BUILD_MAJOR >= 13: 

131 cdef tuple _VALID_LOCATION_TYPES = ("device", "host", "host_numa") (empty)

132  

133  

134 cdef _resolve_preferred_location(ManagedMemoryResourceOptions opts): (empty)

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

136  

137 Returns a 4-tuple: 

138 (CUmemLocationType, loc_id, pref_loc_type_str, pref_loc_id) 

139 """ 

140 cdef object pref_loc = opts.preferred_location if opts is not None else None 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

141 cdef object pref_type = opts.preferred_location_type if opts is not None else None 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

142  

143 if pref_type is not None and pref_type not in _VALID_LOCATION_TYPES: 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

144 raise ValueError( 1 ctx1a

145 f"preferred_location_type must be one of {_VALID_LOCATION_TYPES!r} " 1 ctx1a

146 f"or None, got {pref_type!r}" 1 ctx1a

147 ) 

148  

149 if pref_type is None: 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

150 # Legacy behavior 

151 if pref_loc is None: 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

152 return ( 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

153 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_NONE, 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

154 -1, None, -1, 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

155 ) 

156 if pref_loc == -1: 5 ctx1ecdaf

157 return ( 1 ctx1d

158 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1d

159 -1, "host", -1, 

160 ) 

161 if pref_loc < 0: 4 ctx1ecaf

162 raise ValueError( 1 ctx1a

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

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

165 ) 

166 return ( 3 ctx1ecf

167 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 3 ctx1ecf

168 pref_loc, "device", pref_loc, 3 ctx1ecf

169 ) 

170  

171 if pref_type == "device": 5 ctx1gcbda

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

173 raise ValueError( 1 ctx1a

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

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

176 ) 

177 return ( 1 ctx1c

178 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 1 ctx1c

179 pref_loc, "device", pref_loc, 1 ctx1c

180 ) 

181  

182 if pref_type == "host": 4 ctx1gbda

183 if pref_loc is not None: 2 ctx1da

184 raise ValueError( 1 ctx1a

185 f"preferred_location must be None when " 1 ctx1a

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

187 ) 

188 return ( 1 ctx1d

189 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 1 ctx1d

190 -1, "host", -1, 

191 ) 

192  

193 # pref_type == "host_numa" 

194 if pref_loc is None: 3 ctx1gba

195 from .._device import Device 2 ctx1gb

196 dev = Device() 2 ctx1gb

197 numa_id = dev.properties.host_numa_id 2 ctx1gb

198 if numa_id < 0: 2 ctx1gb

199 raise RuntimeError( 1 ctx1g

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

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

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

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

204 ) 

205 return ( 1 ctx1b

206 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

207 numa_id, "host_numa", numa_id, 1 ctx1b

208 ) 

209 if pref_loc < 0: 2 ctx1ba

210 raise ValueError( 1 ctx1a

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

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

213 ) 

214 return ( 1 ctx1b

215 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, 1 ctx1b

216 pref_loc, "host_numa", pref_loc, 1 ctx1b

217 ) 

218  

219  

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

221 IF CUDA_CORE_BUILD_MAJOR >= 13: 

222 cdef ManagedMemoryResourceOptions opts = check_or_create_options( 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

223 ManagedMemoryResourceOptions, options, "ManagedMemoryResource options", 

224 keep_none=True 

225 ) 

226 cdef cydriver.CUmemLocationType loc_type 

227 cdef int loc_id 

228  

229 loc_type, loc_id, self._pref_loc_type, self._pref_loc_id = ( 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

230 _resolve_preferred_location(opts) 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

231 ) 

232  

233 if opts is None: 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

234 try: 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

235 MP_init_current_pool( 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

236 self, 

237 loc_type, 

238 loc_id, 

239 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

240 ) 

241 except CUDAError as e: 

242 if "CUDA_ERROR_NOT_SUPPORTED" in str(e): 

243 from .._device import Device 

244 if not Device().properties.concurrent_managed_access: 

245 raise RuntimeError( 

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

247 "managed allocations (concurrent managed access is not " 

248 "available). Use " 

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

250 "to create a dedicated managed pool." 

251 ) from e 

252 raise 

253 else: 

254 MP_init_create_pool( 16 ctx1ecbdhfklmnopqrst

255 self, 

256 loc_type, 

257 loc_id, 

258 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED, 

259 False, 16 ctx1ecbdhfklmnopqrst

260 0, 

261 ) 

262  

263 _check_concurrent_managed_access() 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

264 ELSE: 

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

266  

267  

268cdef bint _concurrent_access_warned = False (empty)

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

270  

271  

272cdef inline _check_concurrent_managed_access(): (empty)

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

274 global _concurrent_access_warned 

275 if _concurrent_access_warned: 29 ctx1uegijcbdahfkvlwmxnyozpAqBrCst

276 return 28 ctx1uegijcbdhfkvlwmxnyozpAqBrCst

277  

278 cdef int c_concurrent = 0 1 ctx1a

279 with _concurrent_access_lock: 1 ctx1a

280 if _concurrent_access_warned: 1 ctx1a

281 return 

282  

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

284 # later, so any device will do. 

285 with nogil: 1 ctx1a

286 HANDLE_RETURN(cydriver.cuDeviceGetAttribute( 1 ctx1a

287 &c_concurrent, 

288 cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CONCURRENT_MANAGED_ACCESS, 

289 0)) 

290 if not c_concurrent: 1 ctx1a

291 warnings.warn( 

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

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

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

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

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

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

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

299 UserWarning, 

300 stacklevel=3 

301 ) 

302  

303 _concurrent_access_warned = True 1 ctx1a

304  

305  

306def reset_concurrent_access_warning(): (empty)

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

308 global _concurrent_access_warned 

309 _concurrent_access_warned = False