Coverage for cuda/core/_memory/_pinned_memory_resource.pyx: 62.03%

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 

8from cuda.core._memory._buffer cimport Buffer 

9from cuda.core._memory._memory_pool cimport ( 

10 _MemPool, 

11 _MP_allocate, 

12 MP_check_open, 

13 MP_init_create_pool, 

14 MP_init_current_pool, 

15) 

16from cuda.core._memory cimport _ipc 

17from cuda.core._memory._ipc cimport IPCAllocationHandle 

18from cuda.core._stream cimport Stream, Stream_accept 

19from cuda.core._utils.cuda_utils cimport ( 

20 check_or_create_options, 

21 HANDLE_RETURN, 

22) 

23  

24from dataclasses import dataclass 

25import multiprocessing 

26import platform # no-cython-lint 

27import uuid 

28  

29from cuda.core._utils.cuda_utils import check_multiprocessing_start_method 

30  

31from typing import TYPE_CHECKING 

32  

33if TYPE_CHECKING: 

34 from cuda.core.graph import GraphBuilder 

35  

36__all__ = ['PinnedMemoryResource', 'PinnedMemoryResourceOptions'] 

37  

38  

39@dataclass 

40cdef class PinnedMemoryResourceOptions: 

41 """Customizable :obj:`~_memory.PinnedMemoryResource` options. 

42  

43 Attributes 

44 ---------- 

45 ipc_enabled : bool, optional 

46 Specifies whether to create an IPC-enabled memory pool. When set to 

47 True, the memory pool and its allocations can be shared with other 

48 processes. (Default to False) 

49  

50 max_size : int, optional 

51 Maximum pool size. When set to 0, defaults to a system-dependent value. 

52 (Default to 0) 

53  

54 numa_id : int or None, optional 

55 Host NUMA node ID for pool placement. When set to None (the default), 

56 the behavior depends on ``ipc_enabled``: 

57  

58 - ``ipc_enabled=False``: OS-managed placement (location type HOST). 

59 - ``ipc_enabled=True``: automatically derived from the current CUDA 

60 device's ``host_numa_id`` attribute, requiring an active CUDA 

61 context. 

62  

63 When set to a non-negative integer, that NUMA node is used explicitly 

64 regardless of ``ipc_enabled`` (location type HOST_NUMA). 

65 """ 

66 ipc_enabled : bool = False 

67 max_size : int = 0 

68 numa_id : int | None = None 

69  

70  

71cdef class PinnedMemoryResource(_MemPool): 

72 """ 

73 A host-pinned memory resource managing a stream-ordered memory pool. 

74  

75 Parameters 

76 ---------- 

77 options : PinnedMemoryResourceOptions 

78 Memory resource creation options. 

79  

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

81 stream-ordered memory pool. If no memory 

82 pool is set as current, the driver's default memory pool 

83 is used. 

84  

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

86 the memory resource. 

87  

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

89 host-pinned memory resource does not own the pool (`is_handle_owned` is 

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

91  

92 Notes 

93 ----- 

94 The device associated with ``stream`` must support host memory pools. If 

95 ``numa_id`` is set or derived for IPC, it must support host NUMA memory pools. 

96 You can query these capabilities through 

97 ``Device.properties.host_memory_pools_supported`` and 

98 ``Device.properties.host_numa_memory_pools_supported``. If the required pool 

99 is unsupported and stream-ordered allocation is not needed, use 

100 :class:`LegacyPinnedMemoryResource`. 

101  

102 To create an IPC-Enabled memory resource (MR) that is capable of sharing 

103 allocations between processes, specify ``ipc_enabled=True`` in the initializer 

104 option. When IPC is enabled and ``numa_id`` is not specified, the NUMA node 

105 is automatically derived from the current CUDA device's ``host_numa_id`` 

106 attribute, which requires an active CUDA context. If ``numa_id`` is 

107 explicitly set, that value is used regardless of ``ipc_enabled``. 

108  

109 See :class:`DeviceMemoryResource` for more details on IPC usage patterns. 

110 """ 

111  

112 def __init__(self, options: PinnedMemoryResourceOptions | None = None) -> None: 

113 _PMR_init(self, options) 61 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY8

114  

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

116 """Allocate a host-pinned buffer asynchronously on the supplied stream.""" 

117 MP_check_open(self) 31 ctx1f9!#$%'=(?)@*;+,-./:QgchdiXjeab

118 if self.is_mapped: 31 ctx1f9!#$%'=(?)@*;+,-./:QgchdiXjeab

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

120 cdef Stream s = Stream_accept(stream) 31 ctx1f9!#$%'=(?)@*;+,-./:QgchdiXjeab

121 device = s.device 31 ctx1f9!#$%'=(?)@*;+,-./:QgchdiXjeab

122 cdef bint supported = ( 

123 device.properties.host_numa_memory_pools_supported 22 ctx19!#$%'=(?)@*;+,-./:cda

124 if self._numa_id >= 0 31 ctx1f9!#$%'=(?)@*;+,-./:QgchdiXjeab

125 else device.properties.host_memory_pools_supported 9 ctx1fQghiXjeb

126 ) 

127  

128 if not supported: 31 ctx1f9!#$%'=(?)@*;+,-./:QgchdiXjeab

129 raise RuntimeError( 

130 f"CUDA device {device.device_id} does not support the requested " 

131 "host memory pool for PinnedMemoryResource. Use " 

132 "LegacyPinnedMemoryResource if memory-pool features are not required." 

133 ) 

134 return _MP_allocate(self, size, s) 31 ctx1f9!#$%'=(?)@*;+,-./:QgchdiXjeab

135  

136 def __reduce__(self) -> tuple[object, ...]: 

137 MP_check_open(self) 

138 return PinnedMemoryResource.from_registry, (self.uuid,) 

139  

140 @staticmethod 

141 def from_registry(uuid: uuid.UUID) -> PinnedMemoryResource: # no-cython-lint 

142 """ 

143 Obtain a registered mapped memory resource. 

144  

145 Raises 

146 ------ 

147 RuntimeError 

148 If no mapped memory resource is found in the registry. 

149 """ 

150 return <PinnedMemoryResource>(_ipc.MP_from_registry(uuid)) 

151  

152 def register(self, uuid: uuid.UUID) -> PinnedMemoryResource: # no-cython-lint 

153 """ 

154 Register a mapped memory resource. 

155  

156 Returns 

157 ------- 

158 The registered mapped memory resource. If one was previously registered 

159 with the given key, it is returned. 

160 """ 

161 return <PinnedMemoryResource>(_ipc.MP_register(self, uuid)) 1 ctx1[

162  

163 @classmethod 

164 def from_allocation_handle( 

165 cls, alloc_handle: int | IPCAllocationHandle 

166 ) -> PinnedMemoryResource: 

167 """Create a host-pinned memory resource from an allocation handle. 

168  

169 Construct a new `PinnedMemoryResource` instance that imports a memory 

170 pool from a shareable handle. The memory pool is marked as owned. 

171  

172 Parameters 

173 ---------- 

174 alloc_handle : int | IPCAllocationHandle 

175 The shareable handle of the host-pinned memory resource to import. If an 

176 integer is supplied, it must represent a valid platform-specific 

177 handle. It is the caller's responsibility to close that handle. 

178  

179 Returns 

180 ------- 

181 A new host-pinned memory resource instance with the imported handle. 

182 """ 

183 # cuMemPoolImportFromShareableHandle requires CUDA to be initialized, but in 

184 # a child process CUDA may not be initialized yet. For DeviceMemoryResource, 

185 # this is not a concern because most likely when retrieving the device_id the 

186 # user would have already initialized CUDA. But since PinnedMemoryResource is 

187 # not device-specific it is unlikelt the case. 

188 HANDLE_RETURN(cydriver.cuInit(0)) 1 ctx1[

189  

190 cdef PinnedMemoryResource mr = <PinnedMemoryResource>( 

191 _ipc.MP_from_allocation_handle(cls, alloc_handle)) 1 ctx1[

192 return mr 

193  

194 @property 

195 def allocation_handle(self) -> IPCAllocationHandle: 

196 """Shareable handle for this memory pool (requires IPC). 

197  

198 The handle can be used to share the memory pool with other processes. 

199 The handle is cached in this `MemoryResource` and owned by it. 

200 """ 

201 MP_check_open(self) 29 ctx2] ^ 9 ! _ [ # $ % } ~ ' ` ab( { bb) | * ; + , - . / : a b

202 if not self.is_ipc_enabled: 29 ctx2] ^ 9 ! _ [ # $ % } ~ ' ` ab( { bb) | * ; + , - . / : a b

203 raise RuntimeError("Memory resource is not IPC-enabled") 1 ctx1b

204 return self._ipc_data._alloc_handle 28 ctx2] ^ 9 ! _ [ # $ % } ~ ' ` ab( { bb) | * ; + , - . / : a

205  

206 @property 

207 def device_id(self) -> int: 

208 """Return -1. Pinned memory is host memory and is not associated with a specific device.""" 

209 return -1 5 ctx1Qjeab

210  

211 @property 

212 def numa_id(self) -> int: 

213 """The host NUMA node ID used for pool placement, or -1 for OS-managed placement.""" 

214 return self._numa_id 5 ctx1abZkY

215  

216 @property 

217 def is_device_accessible(self) -> bool: 

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

219 return True 3 ctx1jea

220  

221 @property 

222 def is_host_accessible(self) -> bool: 

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

224 return True 3 ctx1jea

225  

226  

227cdef inline _PMR_init(PinnedMemoryResource self, options): 

228 from .._device import Device 61 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY8

229  

230 cdef PinnedMemoryResourceOptions opts = check_or_create_options( 61 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY8

231 PinnedMemoryResourceOptions, options, "PinnedMemoryResource options", 

232 keep_none=True 

233 ) 

234 cdef bint ipc_enabled = False 61 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY8

235 cdef size_t max_size = 0 61 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY8

236 cdef cydriver.CUmemLocationType loc_type 

237 cdef int numa_id = -1 61 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY8

238  

239 if opts is not None: 61 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY8

240 ipc_enabled = opts.ipc_enabled 58 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZkY8

241 if ipc_enabled and not _ipc.is_supported(): 58 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZkY8

242 raise RuntimeError(f"IPC is not available on {platform.system()}") 

243 max_size = opts.max_size 58 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZkY8

244  

245 if opts.numa_id is not None: 58 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZkY8

246 numa_id = opts.numa_id 2 ctx1Y8

247 if numa_id < 0: 2 ctx1Y8

248 raise ValueError(f"numa_id must be >= 0, got {numa_id}") 1 ctx18

249 elif ipc_enabled: 56 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZk

250 dev = Device() 41 ctx1lmnopqrstuvwxyzABCDEFGHIJKLMNOPRcSTUVdWak

251 numa_id = dev.properties.host_numa_id 41 ctx1lmnopqrstuvwxyzABCDEFGHIJKLMNOPRcSTUVdWak

252 if numa_id < 0: 41 ctx1lmnopqrstuvwxyzABCDEFGHIJKLMNOPRcSTUVdWak

253 raise RuntimeError( 

254 "Cannot determine host NUMA ID for IPC-enabled pinned " 

255 "memory pool. The system may not support NUMA, or no " 

256 "CUDA context is active. Set numa_id explicitly or " 

257 "call Device.set_current() first.") 

258  

259 if numa_id >= 0: 60 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY

260 loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA 42 ctx1lmnopqrstuvwxyzABCDEFGHIJKLMNOPRcSTUVdWakY

261 else: 

262 loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST 18 ctx1fQ0g1234h56iXj7ebZ

263  

264 self._numa_id = numa_id 60 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY

265  

266 if opts is None: 60 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZkY

267 MP_init_current_pool( 3 ctx1QXj

268 self, 

269 loc_type, 

270 numa_id, 

271 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

272 ) 

273 else: 

274 MP_init_create_pool( 57 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZkY

275 self, 

276 loc_type, 

277 numa_id, 

278 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

279 ipc_enabled, 

280 max_size, 57 ctx1flmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZkY

281 ) 

282  

283  

284def _deep_reduce_pinned_memory_resource(mr: object) -> tuple[object, ...]: 

285 check_multiprocessing_start_method() 21 ctx1]^9!_#$%'`({)|*+,-./:

286 alloc_handle = mr.allocation_handle 21 ctx1]^9!_#$%'`({)|*+,-./:

287 return mr.from_allocation_handle, (alloc_handle,) 21 ctx1]^9!_#$%'`({)|*+,-./:

288  

289  

290multiprocessing.reduction.register(PinnedMemoryResource, _deep_reduce_pinned_memory_resource)