Coverage for cuda / core / experimental / _memory / _device_memory_resource.pyx: 73%

202 statements  

« prev     ^ index     » next       coverage.py v7.13.0, created at 2025-12-10 01:19 +0000

1# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from __future__ import annotations 

6  

7from libc.limits cimport ULLONG_MAX 

8from libc.stdint cimport uintptr_t 

9from libc.stdlib cimport malloc, free 

10from libc.string cimport memset 

11  

12from cuda.bindings cimport cydriver 

13from cuda.core.experimental._memory._buffer cimport Buffer, MemoryResource 

14from cuda.core.experimental._memory cimport _ipc 

15from cuda.core.experimental._memory._ipc cimport IPCAllocationHandle, IPCDataForMR 

16from cuda.core.experimental._stream cimport default_stream, Stream_accept, Stream 

17from cuda.core.experimental._utils.cuda_utils cimport ( 

18 check_or_create_options, 

19 HANDLE_RETURN, 

20) 

21  

22from dataclasses import dataclass 

23from typing import Optional, TYPE_CHECKING 

24import platform # no-cython-lint 

25import uuid 

26import weakref 

27  

28from cuda.core.experimental._utils.cuda_utils import driver 

29  

30if TYPE_CHECKING: 

31 from cuda.core.experimental._memory.buffer import DevicePointerT 

32 from .._device import Device 

33  

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

35  

36  

37@dataclass 

38cdef class DeviceMemoryResourceOptions: 

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

40  

41 Attributes 

42 ---------- 

43 ipc_enabled : bool, optional 

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

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

46 processes. (Default to False) 

47  

48 max_size : int, optional 

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

50 (Default to 0) 

51 """ 

52 ipc_enabled : bool = False 

53 max_size : int = 0 

54  

55  

56cdef class DeviceMemoryResourceAttributes: 

57 cdef: 

58 object _mr_weakref 

59  

60 def __init__(self, *args, **kwargs): 

61 raise RuntimeError("DeviceMemoryResourceAttributes cannot be instantiated directly. Please use MemoryResource APIs.") 

62  

63 @classmethod 

64 def _init(cls, mr): 

65 cdef DeviceMemoryResourceAttributes self = DeviceMemoryResourceAttributes.__new__(cls) 

66 self._mr_weakref = mr 

67 return self 

68  

69 def __repr__(self): 

70 return f"{self.__class__.__name__}(%s)" % ", ".join( 

71 f"{attr}={getattr(self, attr)}" for attr in dir(self) 

72 if not attr.startswith("_") 

73 ) 

74  

75 cdef int _getattribute(self, cydriver.CUmemPool_attribute attr_enum, void* value) except?-1: 

76 cdef DeviceMemoryResource mr = <DeviceMemoryResource>(self._mr_weakref()) 

77 if mr is None: 

78 raise RuntimeError("DeviceMemoryResource is expired") 

79 cdef cydriver.CUmemoryPool pool_handle = mr._handle 

80 with nogil: 

81 HANDLE_RETURN(cydriver.cuMemPoolGetAttribute(pool_handle, attr_enum, value)) 

82 return 0 

83  

84 @property 

85 def reuse_follow_event_dependencies(self): 

86 """Allow memory to be reused when there are event dependencies between streams.""" 

87 cdef int value 

88 self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES, &value) 

89 return bool(value) 

90  

91 @property 

92 def reuse_allow_opportunistic(self): 

93 """Allow reuse of completed frees without dependencies.""" 

94 cdef int value 

95 self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC, &value) 

96 return bool(value) 

97  

98 @property 

99 def reuse_allow_internal_dependencies(self): 

100 """Allow insertion of new stream dependencies for memory reuse.""" 

101 cdef int value 

102 self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES, &value) 

103 return bool(value) 

104  

105 @property 

106 def release_threshold(self): 

107 """Amount of reserved memory to hold before OS release.""" 

108 cdef cydriver.cuuint64_t value 

109 self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, &value) 

110 return int(value) 

111  

112 @property 

113 def reserved_mem_current(self): 

114 """Current amount of backing memory allocated.""" 

115 cdef cydriver.cuuint64_t value 

116 self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT, &value) 

117 return int(value) 

118  

119 @property 

120 def reserved_mem_high(self): 

121 """High watermark of backing memory allocated.""" 

122 cdef cydriver.cuuint64_t value 

123 self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH, &value) 

124 return int(value) 

125  

126 @property 

127 def used_mem_current(self): 

128 """Current amount of memory in use.""" 

129 cdef cydriver.cuuint64_t value 

130 self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_CURRENT, &value) 

131 return int(value) 

132  

133 @property 

134 def used_mem_high(self): 

135 """High watermark of memory in use.""" 

136 cdef cydriver.cuuint64_t value 

137 self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_HIGH, &value) 

138 return int(value) 

139  

140  

141cdef class DeviceMemoryResource(MemoryResource): 

142 """ 

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

144  

145 Parameters 

146 ---------- 

147 device_id : Device | int 

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

149  

150 options : DeviceMemoryResourceOptions 

151 Memory resource creation options. 

152  

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

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

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

156 is used. 

157  

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

159 the memory resource. 

160  

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

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

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

164  

165 Notes 

166 ----- 

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

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

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

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

171 These steps can be accomplished in several ways. 

172  

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

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

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

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

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

178 :attr:`DeviceMemoryResource.is_mapped`. 

179  

180 An MR is shared via an allocation handle obtained by calling 

181 :meth:`DeviceMemoryResource.get_allocation_handle`. The allocation handle 

182 has a platform-specific interpretation; however, memory IPC is currently 

183 only supported for Linux, and in that case allocation handles are file 

184 descriptors. After sending an allocation handle to another process, it can 

185 be used to create an MMR by invoking 

186 :meth:`DeviceMemoryResource.from_allocation_handle`. 

187  

188 Buffers can be shared as serializable descriptors obtained by calling 

189 :meth:`Buffer.get_ipc_descriptor`. In a receiving process, a shared buffer is 

190 created by invoking :meth:`Buffer.from_ipc_descriptor` with an MMR and 

191 buffer descriptor, where the MMR corresponds to the MR that created the 

192 described buffer. 

193  

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

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

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

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

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

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

200 can themselves be serialized and transferred directly. Serialization embeds 

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

202  

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

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

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

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

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

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

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

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

211  

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

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

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

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

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

217 associated MMR. 

218 """ 

219  

220 def __cinit__(self): 

221 self._dev_id = cydriver.CU_DEVICE_INVALID 

222 self._handle = NULL 

223 self._mempool_owned = False 

224 self._ipc_data = None 

225 self._attributes = None 

226 self._peer_accessible_by = () 

227  

228 def __init__(self, device_id: Device | int, options=None): 

229 from .._device import Device 

230 cdef int dev_id = Device(device_id).device_id 

231 opts = check_or_create_options( 

232 DeviceMemoryResourceOptions, options, "DeviceMemoryResource options", 

233 keep_none=True 

234 ) 

235  

236 if opts is None: 

237 DMR_init_current(self, dev_id) 

238 else: 

239 DMR_init_create(self, dev_id, opts) 

240  

241 def __dealloc__(self): 

242 DMR_close(self) 

243  

244 def close(self): 

245 """ 

246 Close the device memory resource and destroy the associated memory pool 

247 if owned. 

248 """ 

249 DMR_close(self) 

250  

251 def __reduce__(self): 

252 return DeviceMemoryResource.from_registry, (self.uuid,) 

253  

254 @staticmethod 

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

256 """ 

257 Obtain a registered mapped memory resource. 

258  

259 Raises 

260 ------ 

261 RuntimeError 

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

263 """ 

264 return _ipc.DMR_from_registry(uuid) 

265  

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

267 """ 

268 Register a mapped memory resource. 

269  

270 Returns 

271 ------- 

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

273 with the given key, it is returned. 

274 """ 

275 return _ipc.DMR_register(self, uuid) 

276  

277 @classmethod 

278 def from_allocation_handle( 

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

280 ) -> DeviceMemoryResource: 

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

282  

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

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

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

286  

287 Parameters 

288 ---------- 

289 device_id : int | Device 

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

291 resource is created. 

292  

293 alloc_handle : int | IPCAllocationHandle 

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

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

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

297  

298 Returns 

299 ------- 

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

301 """ 

302 return _ipc.DMR_from_allocation_handle(cls, device_id, alloc_handle) 

303  

304 def get_allocation_handle(self) -> IPCAllocationHandle: 

305 """Export the memory pool handle to be shared (requires IPC). 

306  

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

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

309  

310 Returns 

311 ------- 

312 The shareable handle for the memory pool. 

313 """ 

314 if not self.is_ipc_enabled: 

315 raise RuntimeError("Memory resource is not IPC-enabled") 

316 return self._ipc_data._alloc_handle 

317  

318 def allocate(self, size_t size, stream: Stream | GraphBuilder | None = None) -> Buffer: 

319 """Allocate a buffer of the requested size. 

320  

321 Parameters 

322 ---------- 

323 size : int 

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

325 stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`, optional 

326 The stream on which to perform the allocation asynchronously. 

327 If None, an internal stream is used. 

328  

329 Returns 

330 ------- 

331 Buffer 

332 The allocated buffer object, which is accessible on the device that this memory 

333 resource was created for. 

334 """ 

335 if self.is_mapped: 

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

337 stream = Stream_accept(stream) if stream is not None else default_stream() 

338 return DMR_allocate(self, size, <Stream> stream) 

339  

340 def deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream | GraphBuilder | None = None): 

341 """Deallocate a buffer previously allocated by this resource. 

342  

343 Parameters 

344 ---------- 

345 ptr : :obj:`~_memory.DevicePointerT` 

346 The pointer or handle to the buffer to deallocate. 

347 size : int 

348 The size of the buffer to deallocate, in bytes. 

349 stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`, optional 

350 The stream on which to perform the deallocation asynchronously. 

351 If the buffer is deallocated without an explicit stream, the allocation stream 

352 is used. 

353 """ 

354 stream = Stream_accept(stream) if stream is not None else default_stream() 

355 DMR_deallocate(self, <uintptr_t>ptr, size, <Stream> stream) 

356  

357 @property 

358 def attributes(self) -> DeviceMemoryResourceAttributes: 

359 """Memory pool attributes.""" 

360 if self._attributes is None: 

361 ref = weakref.ref(self) 

362 self._attributes = DeviceMemoryResourceAttributes._init(ref) 

363 return self._attributes 

364  

365 @property 

366 def device_id(self) -> int: 

367 """The associated device ordinal.""" 

368 return self._dev_id 

369  

370 @property 

371 def handle(self) -> driver.CUmemoryPool: 

372 """Handle to the underlying memory pool.""" 

373 return driver.CUmemoryPool(<uintptr_t>(self._handle)) 

374  

375 @property 

376 def is_device_accessible(self) -> bool: 

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

378 return True 

379  

380 @property 

381 def is_handle_owned(self) -> bool: 

382 """Whether the memory resource handle is owned. If False, ``close`` has no effect.""" 

383 return self._mempool_owned 

384  

385 @property 

386 def is_host_accessible(self) -> bool: 

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

388 return False 

389  

390 @property 

391 def is_ipc_enabled(self) -> bool: 

392 """Whether this memory resource has IPC enabled.""" 

393 return self._ipc_data is not None 

394  

395 @property 

396 def is_mapped(self) -> bool: 

397 """ 

398 Whether this is a mapping of an IPC-enabled memory resource from 

399 another process. If True, allocation is not permitted. 

400 """ 

401 return self._ipc_data is not None and self._ipc_data._is_mapped 

402  

403 @property 

404 def uuid(self) -> Optional[uuid.UUID]: 

405 """ 

406 A universally unique identifier for this memory resource. Meaningful 

407 only for IPC-enabled memory resources. 

408 """ 

409 return getattr(self._ipc_data, 'uuid', None) 

410  

411 @property 

412 def peer_accessible_by(self): 

413 """ 

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

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

416 from this memory pool. 

417  

418 Returns a tuple of sorted device IDs that currently have peer access to 

419 allocations from this memory pool. 

420  

421 When setting, accepts a sequence of Device objects or device IDs. 

422 Setting to an empty sequence revokes all peer access. 

423  

424 Examples 

425 -------- 

426 >>> dmr = DeviceMemoryResource(0) 

427 >>> dmr.peer_accessible_by = [1] # Grant access to device 1 

428 >>> assert dmr.peer_accessible_by == (1,) 

429 >>> dmr.peer_accessible_by = [] # Revoke access 

430 """ 

431 return self._peer_accessible_by 

432  

433 @peer_accessible_by.setter 

434 def peer_accessible_by(self, devices): 

435 """Set which devices can access this memory pool.""" 

436 from .._device import Device 

437  

438 # Convert all devices to device IDs 

439 cdef set[int] target_ids = {Device(dev).device_id for dev in devices} 

440 target_ids.discard(self._dev_id) # exclude this device from peer access list 

441 this_dev = Device(self._dev_id) 

442 cdef list bad = [dev for dev in target_ids if not this_dev.can_access_peer(dev)] 

443 if bad: 

444 raise ValueError(f"Device {self._dev_id} cannot access peer(s): {', '.join(map(str, bad))}") 

445 cdef set[int] cur_ids = set(self._peer_accessible_by) 

446 cdef set[int] to_add = target_ids - cur_ids 

447 cdef set[int] to_rm = cur_ids - target_ids 

448 cdef size_t count = len(to_add) + len(to_rm) # transaction size 

449 cdef cydriver.CUmemAccessDesc* access_desc = NULL 

450 cdef size_t i = 0 

451  

452 if count > 0: 

453 access_desc = <cydriver.CUmemAccessDesc*>malloc(count * sizeof(cydriver.CUmemAccessDesc)) 

454 if access_desc == NULL: 

455 raise MemoryError("Failed to allocate memory for access descriptors") 

456  

457 try: 

458 for dev_id in to_add: 

459 access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE 

460 access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

461 access_desc[i].location.id = dev_id 

462 i += 1 

463  

464 for dev_id in to_rm: 

465 access_desc[i].flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_NONE 

466 access_desc[i].location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

467 access_desc[i].location.id = dev_id 

468 i += 1 

469  

470 with nogil: 

471 HANDLE_RETURN(cydriver.cuMemPoolSetAccess(self._handle, access_desc, count)) 

472 finally: 

473 if access_desc != NULL: 

474 free(access_desc) 

475  

476 self._peer_accessible_by = tuple(target_ids) 

477  

478  

479# DeviceMemoryResource Implementation 

480# ----------------------------------- 

481  

482cdef void DMR_init_current(DeviceMemoryResource self, int dev_id): 

483 # Get the current memory pool. 

484 cdef cydriver.cuuint64_t current_threshold 

485 cdef cydriver.cuuint64_t max_threshold = ULLONG_MAX 

486  

487 self._dev_id = dev_id 

488 self._mempool_owned = False 

489  

490 with nogil: 

491 HANDLE_RETURN(cydriver.cuDeviceGetMemPool(&(self._handle), dev_id)) 

492  

493 # Set a higher release threshold to improve performance when there are 

494 # no active allocations. By default, the release threshold is 0, which 

495 # means memory is immediately released back to the OS when there are no 

496 # active suballocations, causing performance issues. 

497 HANDLE_RETURN( 

498 cydriver.cuMemPoolGetAttribute( 

499 self._handle, 

500 cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, 

501 &current_threshold 

502 ) 

503 ) 

504  

505 # If threshold is 0 (default), set it to maximum to retain memory in the pool. 

506 if current_threshold == 0: 

507 HANDLE_RETURN(cydriver.cuMemPoolSetAttribute( 

508 self._handle, 

509 cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, 

510 &max_threshold 

511 )) 

512  

513  

514cdef void DMR_init_create( 

515 DeviceMemoryResource self, int dev_id, DeviceMemoryResourceOptions opts 

516): 

517 # Create a new memory pool. 

518 cdef cydriver.CUmemPoolProps properties 

519  

520 if opts.ipc_enabled and not _ipc.is_supported(): 

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

522  

523 memset(&properties, 0, sizeof(cydriver.CUmemPoolProps)) 

524 properties.allocType = cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED 

525 properties.handleTypes = _ipc.IPC_HANDLE_TYPE if opts.ipc_enabled else cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE 

526 properties.location.id = dev_id 

527 properties.location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

528 properties.maxSize = opts.max_size 

529 properties.win32SecurityAttributes = NULL 

530 properties.usage = 0 

531  

532 self._dev_id = dev_id 

533 self._mempool_owned = True 

534  

535 with nogil: 

536 HANDLE_RETURN(cydriver.cuMemPoolCreate(&(self._handle), &properties)) 

537 # TODO: should we also set the threshold here? 

538  

539 if opts.ipc_enabled: 

540 alloc_handle = _ipc.DMR_export_mempool(self) 

541 self._ipc_data = IPCDataForMR(alloc_handle, False) 

542  

543  

544# Raise an exception if the given stream is capturing. 

545# A result of CU_STREAM_CAPTURE_STATUS_INVALIDATED is considered an error. 

546cdef inline int check_not_capturing(cydriver.CUstream s) except?-1 nogil: 

547 cdef cydriver.CUstreamCaptureStatus capturing 

548 HANDLE_RETURN(cydriver.cuStreamIsCapturing(s, &capturing)) 

549 if capturing != cydriver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_NONE: 

550 raise RuntimeError("DeviceMemoryResource cannot perform memory operations on " 

551 "a capturing stream (consider using GraphMemoryResource).") 

552  

553  

554cdef inline Buffer DMR_allocate(DeviceMemoryResource self, size_t size, Stream stream): 

555 cdef cydriver.CUstream s = stream._handle 

556 cdef cydriver.CUdeviceptr devptr 

557 with nogil: 

558 check_not_capturing(s) 

559 HANDLE_RETURN(cydriver.cuMemAllocFromPoolAsync(&devptr, size, self._handle, s)) 

560 cdef Buffer buf = Buffer.__new__(Buffer) 

561 buf._ptr = <uintptr_t>(devptr) 

562 buf._ptr_obj = None 

563 buf._size = size 

564 buf._memory_resource = self 

565 buf._alloc_stream = stream 

566 return buf 

567  

568  

569cdef inline void DMR_deallocate( 

570 DeviceMemoryResource self, uintptr_t ptr, size_t size, Stream stream 

571) noexcept: 

572 cdef cydriver.CUstream s = stream._handle 

573 cdef cydriver.CUdeviceptr devptr = <cydriver.CUdeviceptr>ptr 

574 cdef cydriver.CUresult r 

575 with nogil: 

576 r = cydriver.cuMemFreeAsync(devptr, s) 

577 if r != cydriver.CUDA_ERROR_INVALID_CONTEXT: 

578 HANDLE_RETURN(r) 

579  

580  

581cdef inline DMR_close(DeviceMemoryResource self): 

582 if self._handle == NULL: 

583 return 

584  

585 # This works around nvbug 5698116. When a memory pool handle is recycled 

586 # the new handle inherits the peer access state of the previous handle. 

587 if self._peer_accessible_by: 

588 self.peer_accessible_by = [] 

589  

590 try: 

591 if self._mempool_owned: 

592 with nogil: 

593 HANDLE_RETURN(cydriver.cuMemPoolDestroy(self._handle)) 

594 finally: 

595 self._dev_id = cydriver.CU_DEVICE_INVALID 

596 self._handle = NULL 

597 self._attributes = None 

598 self._mempool_owned = False 

599 self._ipc_data = None 

600 self._peer_accessible_by = () 

601  

602  

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

604cpdef DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id): 

605 """ 

606 Probes peer access from the given device using cuMemPoolGetAccess. 

607  

608 Parameters 

609 ---------- 

610 device_id : int or Device 

611 The device to query access for. 

612  

613 Returns 

614 ------- 

615 str 

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

617 """ 

618 from .._device import Device 

619  

620 cdef int dev_id = Device(device_id).device_id 

621 cdef cydriver.CUmemAccess_flags flags 

622 cdef cydriver.CUmemLocation location 

623  

624 location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

625 location.id = dev_id 

626  

627 with nogil: 

628 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, dmr._handle, &location)) 

629  

630 if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE: 

631 return "rw" 

632 elif flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READ: 

633 return "r" 

634 else: 

635 return ""