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

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 

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  

24import cython 

25from dataclasses import dataclass (empty)

26import multiprocessing (empty)

27import platform # no-cython-lint (empty)

28import uuid (empty)

29  

30from cuda.core._utils.cuda_utils import check_multiprocessing_start_method (empty)

31  

32from typing import TYPE_CHECKING (empty)

33  

34if TYPE_CHECKING: (empty)

35 from cuda.core.graph import GraphBuilder (empty)

36  

37__all__ = ['PinnedMemoryResource', 'PinnedMemoryResourceOptions'] (empty)

38  

39  

40@dataclass (empty)

41cdef class PinnedMemoryResourceOptions: 

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

43  

44 Attributes 

45 ---------- 

46 ipc_enabled : bool, optional 

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

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

49 processes. (Default to False) 

50  

51 max_size : int, optional 

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

53 (Default to 0) 

54  

55 numa_id : int or None, optional 

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

57 the behavior depends on ``ipc_enabled``: 

58  

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

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

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

62 context. 

63  

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

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

66 """ 

67 ipc_enabled : bool = False (empty)

68 max_size : int = 0 

69 numa_id : int | None = None (empty)

70  

71  

72cdef class PinnedMemoryResource(_MemPool): 

73 """ 

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

75  

76 Parameters 

77 ---------- 

78 options : PinnedMemoryResourceOptions 

79 Memory resource creation options. 

80  

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

82 stream-ordered memory pool. If no memory 

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

84 is used. 

85  

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

87 the memory resource. 

88  

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

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

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

92  

93 Notes 

94 ----- 

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

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

97 You can query these capabilities through 

98 ``Device.properties.host_memory_pools_supported`` and 

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

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

101 :class:`LegacyPinnedMemoryResource`. 

102  

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

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

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

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

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

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

109  

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

111 """ 

112  

113 @cython.annotation_typing(False) 

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

115 _PMR_init(self, options) 61 ctx1kfmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY8

116  

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

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

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

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

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

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

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

124 cdef bint supported = ( 

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

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

127 else device.properties.host_memory_pools_supported 9 ctx1fQghiXjeb

128 ) 

129  

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

131 raise RuntimeError( 

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

133 "host memory pool for PinnedMemoryResource. Use " 

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

135 ) 

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

137  

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

139 MP_check_open(self) 

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

141  

142 @staticmethod (empty)

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

144 """ 

145 Obtain a registered mapped memory resource. 

146  

147 Raises 

148 ------ 

149 RuntimeError 

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

151 """ 

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

153  

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

155 """ 

156 Register a mapped memory resource. 

157  

158 Returns 

159 ------- 

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

161 with the given key, it is returned. 

162 """ 

163 return <PinnedMemoryResource>(_ipc.MP_register(self, uuid)) 1 ctx1@

164  

165 @classmethod (empty)

166 def from_allocation_handle( 

167 cls, alloc_handle: int | IPCAllocationHandle 

168 ) -> PinnedMemoryResource: 

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

170  

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

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

173  

174 Parameters 

175 ---------- 

176 alloc_handle : int | IPCAllocationHandle 

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

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

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

180  

181 Returns 

182 ------- 

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

184 """ 

185 # cuMemPoolImportFromShareableHandle requires CUDA to be initialized, but in 

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

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

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

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

190 HANDLE_RETURN(cydriver.cuInit(0)) 1 ctx1@

191  

192 cdef PinnedMemoryResource mr = <PinnedMemoryResource>( 

193 _ipc.MP_from_allocation_handle(cls, alloc_handle)) 1 ctx1@

194 return mr 

195  

196 @property 

197 def allocation_handle(self) -> IPCAllocationHandle: 

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

199  

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

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

202 """ 

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

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

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

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

207  

208 @property 

209 def device_id(self) -> int: 

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

211 return -1 5 ctx1Qjeab

212  

213 @property 

214 def numa_id(self) -> int: 

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

216 return self._numa_id 5 ctx1abZlY

217  

218 @property 

219 def is_device_accessible(self) -> bool: 

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

221 return True 3 ctx1jea

222  

223 @property 

224 def is_host_accessible(self) -> bool: 

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

226 return True 3 ctx1jea

227  

228  

229cdef inline _PMR_init(PinnedMemoryResource self, options): (empty)

230 from .._device import Device 60 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY8

231  

232 cdef PinnedMemoryResourceOptions opts = check_or_create_options( 60 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY8

233 PinnedMemoryResourceOptions, options, "PinnedMemoryResource options", 

234 keep_none=True 

235 ) 

236 cdef bint ipc_enabled = False 60 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY8

237 cdef size_t max_size = 0 60 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY8

238 cdef cydriver.CUmemLocationType loc_type 

239 cdef int numa_id = -1 60 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY8

240  

241 if opts is not None: 60 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY8

242 ipc_enabled = opts.ipc_enabled 57 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZlY8

243 if ipc_enabled and not _ipc.is_supported(): 57 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZlY8

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

245 max_size = opts.max_size 57 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZlY8

246  

247 if opts.numa_id is not None: 57 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZlY8

248 numa_id = opts.numa_id 2 ctx1Y8

249 if numa_id < 0: 2 ctx1Y8

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

251 elif ipc_enabled: 55 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZl

252 dev = Device() 40 ctx1mnopqrstuvwxyzABCDEFGHIJKLMNOPRcSTUVdWal

253 numa_id = dev.properties.host_numa_id 40 ctx1mnopqrstuvwxyzABCDEFGHIJKLMNOPRcSTUVdWal

254 if numa_id < 0: 40 ctx1mnopqrstuvwxyzABCDEFGHIJKLMNOPRcSTUVdWal

255 raise RuntimeError( 

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

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

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

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

260  

261 if numa_id >= 0: 59 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY

262 loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA 41 ctx1mnopqrstuvwxyzABCDEFGHIJKLMNOPRcSTUVdWalY

263 else: 

264 loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST 18 ctx1fQ0g1234h56iXj7ebZ

265  

266 self._numa_id = numa_id 59 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY

267  

268 if opts is None: 59 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOPQ0Rgc1S2T3U4Vhd5W6iXj7eabZlY

269 MP_init_current_pool( 3 ctx1QXj

270 self, 

271 loc_type, 

272 numa_id, 

273 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

274 ) 

275 else: 

276 MP_init_create_pool( 56 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZlY

277 self, 

278 loc_type, 

279 numa_id, 

280 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

281 ipc_enabled, 

282 max_size, 56 ctx1fmnopqrstuvwxyzABCDEFGHIJKLMNOP0Rgc1S2T3U4Vhd5W6i7eabZlY

283 ) 

284  

285  

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

287 check_multiprocessing_start_method() 20 ctx1[]9!^#$%_'`({)*+,-./

288 alloc_handle = mr.allocation_handle 20 ctx1[]9!^#$%_'`({)*+,-./

289 return mr.from_allocation_handle, (alloc_handle,) 20 ctx1[]9!^#$%_'`({)*+,-./

290  

291  

292multiprocessing.reduction.register(PinnedMemoryResource, _deep_reduce_pinned_memory_resource) (empty)