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

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._memory_pool cimport _MemPool, MP_init_create_pool, MP_init_current_pool 

9from cuda.core._memory cimport _ipc 

10from cuda.core._memory._ipc cimport IPCAllocationHandle 

11from cuda.core._utils.cuda_utils cimport ( 

12 check_or_create_options, 

13 HANDLE_RETURN, 

14) 

15  

16from dataclasses import dataclass (empty)

17import multiprocessing (empty)

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

19import uuid (empty)

20  

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

22  

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

24  

25  

26@dataclass (empty)

27cdef class PinnedMemoryResourceOptions: 

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

29  

30 Attributes 

31 ---------- 

32 ipc_enabled : bool, optional 

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

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

35 processes. (Default to False) 

36  

37 max_size : int, optional 

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

39 (Default to 0) 

40  

41 numa_id : int or None, optional 

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

43 the behavior depends on ``ipc_enabled``: 

44  

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

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

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

48 context. 

49  

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

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

52 """ 

53 ipc_enabled : bool = False (empty)

54 max_size : int = 0 

55 numa_id : int | None = None (empty)

56  

57  

58cdef class PinnedMemoryResource(_MemPool): 

59 """ 

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

61  

62 Parameters 

63 ---------- 

64 options : PinnedMemoryResourceOptions 

65 Memory resource creation options. 

66  

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

68 stream-ordered memory pool. If no memory 

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

70 is used. 

71  

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

73 the memory resource. 

74  

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

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

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

78  

79 Notes 

80 ----- 

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

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

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

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

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

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

87  

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

89 """ 

90  

91 def __init__(self, options=None): 

92 _PMR_init(self, options) 55 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN1

93  

94 def __reduce__(self): (empty)

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

96  

97 @staticmethod (empty)

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

99 """ 

100 Obtain a registered mapped memory resource. 

101  

102 Raises 

103 ------ 

104 RuntimeError 

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

106 """ 

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

108  

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

110 """ 

111 Register a mapped memory resource. 

112  

113 Returns 

114 ------- 

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

116 with the given key, it is returned. 

117 """ 

118 return <PinnedMemoryResource>(_ipc.MP_register(self, uuid)) 

119  

120 @classmethod (empty)

121 def from_allocation_handle( 

122 cls, alloc_handle: int | IPCAllocationHandle 

123 ) -> PinnedMemoryResource: 

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

125  

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

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

128  

129 Parameters 

130 ---------- 

131 alloc_handle : int | IPCAllocationHandle 

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

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

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

135  

136 Returns 

137 ------- 

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

139 """ 

140 # cuMemPoolImportFromShareableHandle requires CUDA to be initialized, but in 

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

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

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

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

145 HANDLE_RETURN(cydriver.cuInit(0)) 

146  

147 cdef PinnedMemoryResource mr = <PinnedMemoryResource>( 

148 _ipc.MP_from_allocation_handle(cls, alloc_handle)) 

149 return mr 

150  

151 @property 

152 def allocation_handle(self) -> IPCAllocationHandle: 

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

154  

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

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

157 """ 

158 if not self.is_ipc_enabled: 26 ctx1564789:;!#=$%?'()@*+,-./ab

159 raise RuntimeError("Memory resource is not IPC-enabled") 2 ctx14b

160 return self._ipc_data._alloc_handle 25 ctx1564789:;!#=$%?'()@*+,-./a

161  

162 @property 

163 def device_id(self) -> int: 

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

165 return -1 5 ctx120Mab

166  

167 @property 

168 def numa_id(self) -> int: 

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

170 return self._numa_id 5 ctx1abOcN

171  

172 @property 

173 def is_device_accessible(self) -> bool: 

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

175 return True 3 ctx10Ma

176  

177 @property 

178 def is_host_accessible(self) -> bool: 

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

180 return True 3 ctx10Ma

181  

182  

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

184 from .._device import Device 55 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN1

185  

186 cdef PinnedMemoryResourceOptions opts = check_or_create_options( 55 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN1

187 PinnedMemoryResourceOptions, options, "PinnedMemoryResource options", 

188 keep_none=True 

189 ) 

190 cdef bint ipc_enabled = False 55 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN1

191 cdef size_t max_size = 0 55 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN1

192 cdef cydriver.CUmemLocationType loc_type 

193 cdef int numa_id = -1 55 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN1

194  

195 if opts is not None: 55 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN1

196 ipc_enabled = opts.ipc_enabled 52 ctx1defghijklmnopqrstuvwxyzABCDPEQFRGSHTIUJVKWLXYMabOcN1

197 if ipc_enabled and not _ipc.is_supported(): 52 ctx1defghijklmnopqrstuvwxyzABCDPEQFRGSHTIUJVKWLXYMabOcN1

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

199 max_size = opts.max_size 52 ctx1defghijklmnopqrstuvwxyzABCDPEQFRGSHTIUJVKWLXYMabOcN1

200  

201 if opts.numa_id is not None: 52 ctx1defghijklmnopqrstuvwxyzABCDPEQFRGSHTIUJVKWLXYMabOcN1

202 numa_id = opts.numa_id 2 ctx1N1

203 if numa_id < 0: 2 ctx1N1

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

205 elif ipc_enabled: 50 ctx1defghijklmnopqrstuvwxyzABCDPEQFRGSHTIUJVKWLXYMabOc

206 dev = Device() 37 ctx1defghijklmnopqrstuvwxyzABCDEFGHIJKLac

207 numa_id = dev.properties.host_numa_id 37 ctx1defghijklmnopqrstuvwxyzABCDEFGHIJKLac

208 if numa_id < 0: 37 ctx1defghijklmnopqrstuvwxyzABCDEFGHIJKLac

209 raise RuntimeError( 

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

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

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

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

214  

215 if numa_id >= 0: 54 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN

216 loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA 38 ctx1defghijklmnopqrstuvwxyzABCDEFGHIJKLacN

217 else: 

218 loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST 16 ctx12PQRSTUVWXY30MbO

219  

220 self._numa_id = numa_id 54 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN

221  

222 if opts is None: 54 ctx1defghijklmnopqrstuvwxyzABCD2PEQFRGSHTIUJVKWLXY30MabOcN

223 MP_init_current_pool( 3 ctx1230

224 self, 

225 loc_type, 

226 numa_id, 

227 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

228 ) 

229 else: 

230 MP_init_create_pool( 51 ctx1defghijklmnopqrstuvwxyzABCDPEQFRGSHTIUJVKWLXYMabOcN

231 self, 

232 loc_type, 

233 numa_id, 

234 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

235 ipc_enabled, 

236 max_size, 51 ctx1defghijklmnopqrstuvwxyzABCDPEQFRGSHTIUJVKWLXYMabOcN

237 ) 

238  

239  

240def _deep_reduce_pinned_memory_resource(mr): (empty)

241 check_multiprocessing_start_method() 19 ctx1564789!#$%'()*+,-./

242 alloc_handle = mr.allocation_handle 19 ctx1564789!#$%'()*+,-./

243 return mr.from_allocation_handle, (alloc_handle,) 19 ctx1564789!#$%'()*+,-./

244  

245  

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