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) 72 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2,

114  

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

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

117  

118 Parameters 

119 ---------- 

120 size : int 

121 The size of the buffer to allocate, in bytes. 

122 stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder` 

123 Keyword-only. The stream on which to perform the allocation 

124 asynchronously. Must be passed explicitly; pass 

125 ``device.default_stream`` to use the default stream. 

126  

127 Returns 

128 ------- 

129 Buffer 

130 The allocated buffer object, which is accessible from the host. 

131  

132 Raises 

133 ------ 

134 RuntimeError 

135 If the stream's device does not support the requested host memory 

136 pool. Use :class:`LegacyPinnedMemoryResource` when stream-ordered 

137 allocation is not required. 

138 """ 

139 MP_check_open(self) 34 ctx2f - . / : ; = ~ ? ab@ bb[ } ] ^ _ ` { | U g c h d i j k l 1 m e a b

140 if self.is_mapped: 34 ctx2f - . / : ; = ~ ? ab@ bb[ } ] ^ _ ` { | U g c h d i j k l 1 m e a b

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

142 cdef Stream s = Stream_accept(stream) 34 ctx2f - . / : ; = ~ ? ab@ bb[ } ] ^ _ ` { | U g c h d i j k l 1 m e a b

143 device = s.device 34 ctx2f - . / : ; = ~ ? ab@ bb[ } ] ^ _ ` { | U g c h d i j k l 1 m e a b

144 cdef bint supported = ( 

145 device.properties.host_numa_memory_pools_supported 22 ctx2- . / : ; = ~ ? ab@ bb[ } ] ^ _ ` { | c d a

146 if self._numa_id >= 0 34 ctx2f - . / : ; = ~ ? ab@ bb[ } ] ^ _ ` { | U g c h d i j k l 1 m e a b

147 else device.properties.host_memory_pools_supported 12 ctx1fUghijkl1meb

148 ) 

149  

150 if not supported: 34 ctx2f - . / : ; = ~ ? ab@ bb[ } ] ^ _ ` { | U g c h d i j k l 1 m e a b

151 raise RuntimeError( 

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

153 "host memory pool for PinnedMemoryResource. Use " 

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

155 ) 

156 return _MP_allocate(self, size, s) 34 ctx2f - . / : ; = ~ ? ab@ bb[ } ] ^ _ ` { | U g c h d i j k l 1 m e a b

157  

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

159 MP_check_open(self) 

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

161  

162 @staticmethod 

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

164 """ 

165 Obtain a registered mapped memory resource. 

166  

167 Raises 

168 ------ 

169 RuntimeError 

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

171 """ 

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

173  

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

175 """ 

176 Register a mapped memory resource. 

177  

178 Returns 

179 ------- 

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

181 with the given key, it is returned. 

182  

183 Raises 

184 ------ 

185 RuntimeError 

186 If the memory resource is not IPC-enabled. 

187 ValueError 

188 If the memory resource already has a different UUID. 

189 """ 

190 return <PinnedMemoryResource>(_ipc.MP_register(self, uuid)) 2 ctx2cbnb

191  

192 @classmethod 

193 def from_allocation_handle( 

194 cls, alloc_handle: int | IPCAllocationHandle 

195 ) -> PinnedMemoryResource: 

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

197  

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

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

200  

201 Parameters 

202 ---------- 

203 alloc_handle : int | IPCAllocationHandle 

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

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

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

207  

208 Returns 

209 ------- 

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

211 """ 

212 # cuMemPoolImportFromShareableHandle requires CUDA to be initialized, but in 

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

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

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

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

217 HANDLE_RETURN(cydriver.cuInit(0)) 1 ctx2cb

218  

219 cdef PinnedMemoryResource mr = <PinnedMemoryResource>( 

220 _ipc.MP_from_allocation_handle(cls, alloc_handle)) 1 ctx2cb

221 return mr 

222  

223 @property 

224 def allocation_handle(self) -> IPCAllocationHandle: 

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

226  

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

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

229 """ 

230 MP_check_open(self) 29 ctx2dbeb- . fbcb/ : ; jbkb= gblb? hbmb@ ib[ } ] ^ _ ` { | a b

231 if not self.is_ipc_enabled: 29 ctx2dbeb- . fbcb/ : ; jbkb= gblb? hbmb@ ib[ } ] ^ _ ` { | a b

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

233 return self._ipc_data._alloc_handle 28 ctx2dbeb- . fbcb/ : ; jbkb= gblb? hbmb@ ib[ } ] ^ _ ` { | a

234  

235 @property 

236 def device_id(self) -> int: 

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

238 return -1 5 ctx1Umeab

239  

240 @property 

241 def numa_id(self) -> int: 

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

243 return self._numa_id 5 ctx1ab3n2

244  

245 @property 

246 def is_device_accessible(self) -> bool: 

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

248 return True 3 ctx1mea

249  

250 @property 

251 def is_host_accessible(self) -> bool: 

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

253 return True 3 ctx1mea

254  

255  

256cdef inline _PMR_init(PinnedMemoryResource self, options): 

257 from .._device import Device 72 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2,

258  

259 cdef PinnedMemoryResourceOptions opts = check_or_create_options( 72 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2,

260 PinnedMemoryResourceOptions, options, "PinnedMemoryResource options", 

261 keep_none=True 

262 ) 

263 cdef bint ipc_enabled = False 72 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2,

264 cdef size_t max_size = 0 72 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2,

265 cdef cydriver.CUmemLocationType loc_type 

266 cdef int numa_id = -1 72 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2,

267  

268 if opts is not None: 72 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2,

269 ipc_enabled = opts.ipc_enabled 69 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRST4Vgc5W6X7Y8Zhd90!i#$%'()*jkl+eab3n2,

270 if ipc_enabled and not _ipc.is_supported(): 69 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRST4Vgc5W6X7Y8Zhd90!i#$%'()*jkl+eab3n2,

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

272 max_size = opts.max_size 69 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRST4Vgc5W6X7Y8Zhd90!i#$%'()*jkl+eab3n2,

273  

274 if opts.numa_id is not None: 69 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRST4Vgc5W6X7Y8Zhd90!i#$%'()*jkl+eab3n2,

275 numa_id = opts.numa_id 2 ctx12,

276 if numa_id < 0: 2 ctx12,

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

278 elif ipc_enabled: 67 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRST4Vgc5W6X7Y8Zhd90!i#$%'()*jkl+eab3n

279 dev = Device() 42 ctx1opqrstuvwxyzABCDEFGHIJKLMNOPQRSTVcWXYZd0an

280 numa_id = dev.properties.host_numa_id 42 ctx1opqrstuvwxyzABCDEFGHIJKLMNOPQRSTVcWXYZd0an

281 if numa_id < 0: 42 ctx1opqrstuvwxyzABCDEFGHIJKLMNOPQRSTVcWXYZd0an

282 raise RuntimeError( 

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

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

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

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

287  

288 if numa_id >= 0: 71 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2

289 loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA 43 ctx1opqrstuvwxyzABCDEFGHIJKLMNOPQRSTVcWXYZd0an2

290 else: 

291 loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST 28 ctx1fU4g5678h9!i#$%'()*jkl1m+eb3

292  

293 self._numa_id = numa_id 71 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2

294  

295 if opts is None: 71 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU4Vgc5W6X7Y8Zhd90!i#$%'()*jkl1m+eab3n2

296 MP_init_current_pool( 3 ctx1U1m

297 self, 

298 loc_type, 

299 numa_id, 

300 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

301 ) 

302 else: 

303 MP_init_create_pool( 68 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRST4Vgc5W6X7Y8Zhd90!i#$%'()*jkl+eab3n2

304 self, 

305 loc_type, 

306 numa_id, 

307 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

308 ipc_enabled, 

309 max_size, 68 ctx1fopqrstuvwxyzABCDEFGHIJKLMNOPQRST4Vgc5W6X7Y8Zhd90!i#$%'()*jkl+eab3n2

310 ) 

311  

312  

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

314 check_multiprocessing_start_method() 21 ctx2dbeb- . fb/ : ; = gb? hb@ ib[ ] ^ _ ` { |

315 alloc_handle = mr.allocation_handle 21 ctx2dbeb- . fb/ : ; = gb? hb@ ib[ ] ^ _ ` { |

316 return mr.from_allocation_handle, (alloc_handle,) 21 ctx2dbeb- . fb/ : ; = gb? hb@ ib[ ] ^ _ ` { |

317  

318  

319multiprocessing.reduction.register(PinnedMemoryResource, _deep_reduce_pinned_memory_resource)