Coverage for cuda/core/_memory/_device_memory_resource.pyx: 81.08%

74 statements  

« prev     ^ index     » next       coverage.py v7.15.0, created at 2026-07-03 01:38 +0000

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._memory_pool cimport ( 

9 _MemPool, MP_init_create_pool, MP_raise_release_threshold, 

10) 

11from cuda.core._memory cimport _ipc 

12from cuda.core._memory._ipc cimport IPCAllocationHandle 

13from cuda.core._resource_handles cimport ( 

14 as_cu, 

15 get_device_mempool, 

16 get_last_error, 

17) 

18from cuda.core._utils.cuda_utils cimport ( 

19 check_or_create_options, 

20 HANDLE_RETURN, 

21) 

22from dataclasses import dataclass 

23import multiprocessing 

24import platform # no-cython-lint 

25import uuid 

26  

27from cuda.core._memory._peer_access_utils import PeerAccessibleBySetProxy, replace_peer_accessible_by 

28from cuda.core._utils.cuda_utils import check_multiprocessing_start_method 

29  

30from typing import TYPE_CHECKING 

31  

32if TYPE_CHECKING: 

33 from cuda.core._device import Device 

34  

35__all__ = ['DeviceMemoryResource', 'DeviceMemoryResourceOptions'] 

36  

37  

38@dataclass 

39cdef class DeviceMemoryResourceOptions: 

40 """Customizable :obj:`~_memory.DeviceMemoryResource` options. 

41  

42 Attributes 

43 ---------- 

44 ipc_enabled : bool, optional 

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

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

47 processes. (Default to False) 

48  

49 max_size : int, optional 

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

51 (Default to 0) 

52 """ 

53 ipc_enabled : bool = False 

54 max_size : int = 0 

55  

56  

57cdef class DeviceMemoryResource(_MemPool): 

58 """ 

59 A device memory resource managing a stream-ordered memory pool. 

60  

61 Parameters 

62 ---------- 

63 device_id : Device | int 

64 Device or Device ordinal for which a memory resource is constructed. 

65  

66 options : DeviceMemoryResourceOptions 

67 Memory resource creation options. 

68  

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

70 stream-ordered memory pool for the specified `device_id`. If no memory 

71 pool is set as current, the driver's default memory pool for the device 

72 is used. 

73  

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

75 the memory resource. 

76  

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

78 device memory resource does not own the pool (`is_handle_owned` is 

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

80  

81 Notes 

82 ----- 

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

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

85 option. Sharing an allocation is a two-step procedure that involves 

86 mapping a memory resource and then mapping buffers owned by that resource. 

87 These steps can be accomplished in several ways. 

88  

89 An IPC-enabled memory resource can allocate memory buffers but cannot 

90 receive shared buffers. Mapping an MR to another process creates a "mapped 

91 memory resource" (MMR). An MMR cannot allocate memory buffers and can only 

92 receive shared buffers. MRs and MMRs are both of type 

93 :class:`DeviceMemoryResource` and can be distinguished via 

94 :attr:`DeviceMemoryResource.is_mapped`. 

95  

96 An MR is shared via an allocation handle accessed through the 

97 :attr:`DeviceMemoryResource.allocation_handle` property. The allocation 

98 handle has a platform-specific interpretation; however, memory IPC is 

99 currently only supported for Linux, and in that case allocation handles 

100 are file descriptors. After sending an allocation handle to another 

101 process, it can be used to create an MMR by invoking 

102 :meth:`DeviceMemoryResource.from_allocation_handle`. 

103  

104 Buffers can be shared as serializable descriptors accessed through the 

105 :attr:`Buffer.ipc_descriptor` property. In a receiving process, a shared 

106 buffer is created by invoking :meth:`Buffer.from_ipc_descriptor` with an 

107 MMR and buffer descriptor, where the MMR corresponds to the MR that 

108 created the described buffer. 

109  

110 To help manage the association between memory resources and buffers, a 

111 registry is provided. Every MR has a unique identifier (UUID). MMRs can be 

112 registered by calling :meth:`DeviceMemoryResource.register` with the UUID 

113 of the corresponding MR. Registered MMRs can be looked up via 

114 :meth:`DeviceMemoryResource.from_registry`. When registering MMRs in this 

115 way, the use of buffer descriptors can be avoided. Instead, buffer objects 

116 can themselves be serialized and transferred directly. Serialization embeds 

117 the UUID, which is used to locate the correct MMR during reconstruction. 

118  

119 IPC-enabled memory resources interoperate with the :mod:`multiprocessing` 

120 module to provide a simplified interface. This approach can avoid direct 

121 use of allocation handles, buffer descriptors, MMRs, and the registry. When 

122 using :mod:`multiprocessing` to spawn processes or send objects through 

123 communication channels such as :class:`multiprocessing.Queue`, 

124 :class:`multiprocessing.Pipe`, or :class:`multiprocessing.Connection`, 

125 :class:`Buffer` objects may be sent directly, and in such cases the process 

126 for creating MMRs and mapping buffers will be handled automatically. 

127  

128 For greater efficiency when transferring many buffers, one may also send 

129 MRs and buffers separately. When an MR is sent via :mod:`multiprocessing`, 

130 an MMR is created and registered in the receiving process. Subsequently, 

131 buffers may be serialized and transferred using ordinary :mod:`pickle` 

132 methods. The reconstruction procedure uses the registry to find the 

133 associated MMR. Unpickling a :class:`Buffer` performs an IPC import from 

134 the embedded descriptor; only unpickle buffers received from trusted peers. 

135  

136 Warning 

137 ------- 

138 IPC descriptors and pickled buffers cross a trust boundary between 

139 cooperating same-host processes. A malicious peer can supply crafted 

140 descriptor fields. Use :meth:`Buffer.from_ipc_descriptor` only with 

141 descriptors from trusted peers, and do not unpickle buffers from 

142 untrusted sources. 

143 """ 

144  

145 def __cinit__(self, *args, **kwargs) -> None: 

146 self._dev_id = cydriver.CU_DEVICE_INVALID 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

147  

148 def __init__( 

149 self, 

150 device_id: Device | int, 

151 options: DeviceMemoryResourceOptions | dict[str, object] | None = None 

152 ) -> None: 

153 _DMR_init(self, device_id, options) 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

154  

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

156 return DeviceMemoryResource.from_registry, (self.uuid,) 1adbc

157  

158 @staticmethod 

159 def from_registry(uuid: uuid.UUID) -> DeviceMemoryResource: # no-cython-lint 

160 """ 

161 Obtain a registered mapped memory resource. 

162  

163 Raises 

164 ------ 

165 RuntimeError 

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

167 """ 

168 return <DeviceMemoryResource>(_ipc.MP_from_registry(uuid)) 

169  

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

171 """ 

172 Register a mapped memory resource. 

173  

174 Returns 

175 ------- 

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

177 with the given key, it is returned. 

178 """ 

179 return <DeviceMemoryResource>(_ipc.MP_register(self, uuid)) 

180  

181 @classmethod 

182 def from_allocation_handle( 

183 cls, device_id: Device | int, alloc_handle: int | IPCAllocationHandle 

184 ) -> DeviceMemoryResource: 

185 """Create a device memory resource from an allocation handle. 

186  

187 Construct a new `DeviceMemoryResource` instance that imports a memory 

188 pool from a shareable handle. The memory pool is marked as owned, and 

189 the resource is associated with the specified `device_id`. 

190  

191 Parameters 

192 ---------- 

193 device_id : int | Device 

194 The ID of the device or a Device object for which the memory 

195 resource is created. 

196  

197 alloc_handle : int | IPCAllocationHandle 

198 The shareable handle of the device memory resource to import. If an 

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

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

201  

202 Returns 

203 ------- 

204 A new device memory resource instance with the imported handle. 

205 """ 

206 cdef DeviceMemoryResource mr = <DeviceMemoryResource>( 

207 _ipc.MP_from_allocation_handle(cls, alloc_handle)) 

208 from .._device import Device 

209 mr._dev_id = Device(device_id).device_id 

210 return mr 

211  

212 @property 

213 def allocation_handle(self) -> IPCAllocationHandle: 

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

215  

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

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

218 """ 

219 if not self.is_ipc_enabled: 2wca d ucvce xcycMcNczcAcOcBcCcPcDcEcFcQcGcHcIcf g JcKcLch i j k b c u l v m n

220 raise RuntimeError("Memory resource is not IPC-enabled") 2a d ucvcu

221 return self._ipc_data._alloc_handle 2wca ucvce xcycMcNczcAcOcBcCcPcDcEcFcQcGcHcIcf g JcKcLch i j k b c l v m n

222  

223 @property 

224 def device_id(self) -> int: 

225 """The associated device ordinal.""" 

226 return self._dev_id 2U V W X Y Z 0 1 wca d ucvce xcyczcAcBcCcDcEcFcGcHcIcf g JcKcLch i j k b c ` { r s q } l m n w Rcx y Scz A B C TcUcD VcWcE XcF G YcH I ZcJ 0cK L M N 1cO 2c3cP Q R S T

227  

228 @property 

229 def peer_accessible_by(self) -> PeerAccessibleBySetProxy: 

230 """ 

231 Get or set the devices that can access allocations from this memory 

232 pool. Access can be modified at any time and affects all allocations 

233 from this memory pool. 

234  

235 Returns a set-like proxy of :obj:`~_device.Device` objects that manages 

236 peer access. Inputs are accepted as either :obj:`~_device.Device` 

237 objects or device-ordinal :class:`int` values. 

238  

239 Examples 

240 -------- 

241 >>> dmr = DeviceMemoryResource(0) 

242 >>> dmr.peer_accessible_by = {1} # grant access to device 1 

243 >>> assert 1 in dmr.peer_accessible_by 

244 >>> dmr.peer_accessible_by.add(2) # update access to include device 2 

245 >>> dmr.peer_accessible_by = [] # revoke peer access 

246 """ 

247 return PeerAccessibleBySetProxy(self) 1|t

248  

249 @peer_accessible_by.setter 

250 def peer_accessible_by(self, devices) -> None: 

251 replace_peer_accessible_by(self, devices) 1t

252  

253 @property 

254 def is_device_accessible(self) -> bool: 

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

256 return True 2r s q w Rcx y Scz A B C TcUcD VcWcE XcF G YcH I ZcJ 0cK L M N 1cO 2c3cP Q R S T

257  

258 @property 

259 def is_host_accessible(self) -> bool: 

260 """Return False. This memory resource does not provide host-accessible buffers.""" 

261 return False 1rsq

262  

263  

264cdef inline _DMR_init(DeviceMemoryResource self, device_id, options): 

265 from .._device import Device 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

266 cdef int dev_id = Device(device_id).device_id 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

267 cdef DeviceMemoryResourceOptions opts = check_or_create_options( 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

268 DeviceMemoryResourceOptions, options, "DeviceMemoryResource options", 

269 keep_none=True 

270 ) 

271 cdef bint ipc_enabled = False 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

272 cdef size_t max_size = 0 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

273  

274 self._dev_id = dev_id 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

275  

276 if opts is not None: 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

277 ipc_enabled = opts.ipc_enabled 22 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c q t ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu l v m n rbsb

278 if ipc_enabled and not _ipc.is_supported(): 22 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c q t ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu l v m n rbsb

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

280 max_size = opts.max_size 22 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c q t ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu l v m n rbsb

281  

282 if opts is None: 2tbubvbwbxbU V W X Y Z 0 1 2 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c ` { ybzbAbBbCbr s | q o t } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu DbEbl v m n FbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^brbsb_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

283 self._h_pool = get_device_mempool(dev_id) 2tbubvbwbxbU V W X Y Z 0 1 ` { ybzbAbBbCbr s | o } DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

284 if not self._h_pool: 2tbubvbwbxbU V W X Y Z 0 1 ` { ybzbAbBbCbr s | o } DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

285 HANDLE_RETURN(get_last_error()) 

286 raise RuntimeError( 

287 f"Failed to initialize DeviceMemoryResource for device {dev_id}: " 

288 "cuda-core returned an empty memory pool handle without recording a CUDA error. " 

289 "This is an internal cuda-core error; please report it with your CUDA driver, " 

290 "CUDA Toolkit, and cuda-python versions." 

291 ) 

292 self._mempool_owned = False 2tbubvbwbxbU V W X Y Z 0 1 ` { ybzbAbBbCbr s | o } DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

293 MP_raise_release_threshold(self) 2tbubvbwbxbU V W X Y Z 0 1 ` { ybzbAbBbCbr s | o } DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcw x y z A B C D E F G H I J K L M N O scP Q tcR S T

294 else: 

295 MP_init_create_pool( 22 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c q t ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu l v m n rbsb

296 self, 

297 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 

298 dev_id, 

299 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

300 ipc_enabled, 

301 max_size, 22 a 3 d 4 5 e 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ f g ] ^ _ h i j k b c q t ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbu l v m n rbsb

302 ) 

303  

304  

305# Note: this is referenced in instructions to debug nvbug 5698116. 

306cpdef str DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id): 

307 """ 

308 Probes peer access from the given device using cuMemPoolGetAccess. 

309  

310 Parameters 

311 ---------- 

312 device_id : int or Device 

313 The device to query access for. 

314  

315 Returns 

316 ------- 

317 str 

318 Access permissions: "rw" for read-write, "r" for read-only, "" for no access. 

319 """ 

320 from .._device import Device 1o

321  

322 cdef int dev_id = Device(device_id).device_id 1o

323 cdef cydriver.CUmemAccess_flags flags 

324 cdef cydriver.CUmemLocation location 

325  

326 location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 1o

327 location.id = dev_id 1o

328  

329 with nogil: 1o

330 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(dmr._h_pool), &location)) 1o

331  

332 if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE: 1o

333 return "rw" 1o

334 elif flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READ: 

335 return "r" 

336 else: 

337 return "" 

338  

339  

340def _deep_reduce_device_memory_resource(mr) -> tuple[object, ...]: 

341 check_multiprocessing_start_method() 2wca d ucvce xcyczcAcBcCcDcEcFcGcHcIcf g JcKcLch i j k b c l m n

342 from .._device import Device 2wca d ucvce xcyczcAcBcCcDcEcFcGcHcIcf g JcKcLch i j k b c l m n

343 device = Device(mr.device_id) 2wca d ucvce xcyczcAcBcCcDcEcFcGcHcIcf g JcKcLch i j k b c l m n

344 alloc_handle = mr.allocation_handle 2wca d ucvce xcyczcAcBcCcDcEcFcGcHcIcf g JcKcLch i j k b c l m n

345 return mr.from_allocation_handle, (device, alloc_handle) 2wca ucvce xcyczcAcBcCcDcEcFcGcHcIcf g JcKcLch i j k b c l m n

346  

347  

348multiprocessing.reduction.register(DeviceMemoryResource, _deep_reduce_device_memory_resource)