Coverage for cuda/core/_memory/_virtual_memory_resource.py: 91.50%

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 (empty)

6 

7from dataclasses import dataclass, field (empty)

8from typing import TYPE_CHECKING, Iterable (empty)

9 

10if TYPE_CHECKING: (empty)

11 from cuda.core._stream import Stream 

12 from cuda.core.graph import GraphBuilder 

13 

14from cuda.core._device import Device (empty)

15from cuda.core._memory._buffer import Buffer, MemoryResource (empty)

16from cuda.core._utils.cuda_utils import ( (empty)

17 Transaction, 

18 check_or_create_options, 

19 driver, 

20) 

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

22 _check_driver_error as raise_if_driver_error, 

23) 

24from cuda.core._utils.version import binding_version (empty)

25from cuda.core.typing import ( (empty)

26 DevicePointerType, 

27 VirtualMemoryAccessType, 

28 VirtualMemoryAllocationType, 

29 VirtualMemoryGranularityType, 

30 VirtualMemoryHandleType, 

31 VirtualMemoryLocationType, 

32) 

33 

34__all__ = ["VirtualMemoryResource", "VirtualMemoryResourceOptions"] (empty)

35 

36# Location types whose physical backing lives in host memory. Shared by 

37# VirtualMemoryResource.__init__ and is_host_accessible so the two cannot drift. 

38_HOST_LOCATION_TYPES = frozenset( (empty)

39 { 

40 VirtualMemoryLocationType.HOST, 

41 VirtualMemoryLocationType.HOST_NUMA, 

42 VirtualMemoryLocationType.HOST_NUMA_CURRENT, 

43 } 

44) 

45 

46 

47@dataclass (empty)

48class VirtualMemoryResourceOptions: (empty)

49 """A configuration object for the VirtualMemoryResource 

50 Stores configuration information which tells the resource how to use the CUDA VMM APIs 

51 

52 Attributes 

53 ---------- 

54 allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str 

55 Controls the type of allocation. 

56 location_type: :obj:`~_memory.VirtualMemoryLocationType` | str 

57 Controls the location of the allocation. 

58 handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str 

59 Export handle type for the physical allocation. Use ``"posix_fd"`` on 

60 Linux if you plan to import/export the allocation. Use `None` if you 

61 don't need an exportable handle. 

62 gpu_direct_rdma: bool 

63 Hint that the allocation should be GDR-capable (if supported). 

64 granularity: :obj:`~_memory.VirtualMemoryGranularityType` | str 

65 Controls granularity query and size rounding. 

66 addr_hint: int 

67 A (optional) virtual address hint to try to reserve at. Setting it to 0 lets the CUDA driver decide. 

68 addr_align: int 

69 Alignment for the VA reservation. If `None`, use the queried granularity. 

70 peers: Iterable[int] 

71 Extra device IDs that should be granted access in addition to ``device``. 

72 self_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str 

73 Access flags for the owning device. 

74 peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str 

75 Access flags for peers. 

76 """ 

77 

78 allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED (empty)

79 location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE (empty)

80 handle_type: VirtualMemoryHandleType = VirtualMemoryHandleType.POSIX_FD (empty)

81 granularity: VirtualMemoryGranularityType = VirtualMemoryGranularityType.RECOMMENDED (empty)

82 gpu_direct_rdma: bool = False (empty)

83 addr_hint: int | None = 0 (empty)

84 addr_align: int | None = None (empty)

85 peers: Iterable[int] = field(default_factory=tuple) (empty)

86 self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE (empty)

87 peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE (empty)

88 

89 _a = driver.CUmemAccess_flags (empty)

90 _access_flags = { # noqa: RUF012 (empty)

91 VirtualMemoryAccessType.READ_WRITE: _a.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, 

92 VirtualMemoryAccessType.READ: _a.CU_MEM_ACCESS_FLAGS_PROT_READ, 

93 None: 0, 

94 } 

95 _h = driver.CUmemAllocationHandleType (empty)

96 _handle_types = { # noqa: RUF012 (empty)

97 None: _h.CU_MEM_HANDLE_TYPE_NONE, 

98 VirtualMemoryHandleType.POSIX_FD: _h.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 

99 VirtualMemoryHandleType.WIN32_KMT: _h.CU_MEM_HANDLE_TYPE_WIN32_KMT, 

100 VirtualMemoryHandleType.FABRIC: _h.CU_MEM_HANDLE_TYPE_FABRIC, 

101 } 

102 _g = driver.CUmemAllocationGranularity_flags (empty)

103 _granularity = { # noqa: RUF012 (empty)

104 VirtualMemoryGranularityType.RECOMMENDED: _g.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, 

105 VirtualMemoryGranularityType.MINIMUM: _g.CU_MEM_ALLOC_GRANULARITY_MINIMUM, 

106 } 

107 _l = driver.CUmemLocationType (empty)

108 _location_type = { # noqa: RUF012 (empty)

109 VirtualMemoryLocationType.DEVICE: _l.CU_MEM_LOCATION_TYPE_DEVICE, 

110 VirtualMemoryLocationType.HOST: _l.CU_MEM_LOCATION_TYPE_HOST, 

111 VirtualMemoryLocationType.HOST_NUMA: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA, 

112 VirtualMemoryLocationType.HOST_NUMA_CURRENT: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, 

113 } 

114 _t = driver.CUmemAllocationType (empty)

115 # CUDA 13+ exposes MANAGED in CUmemAllocationType; older 12.x does not 

116 _allocation_type = {VirtualMemoryAllocationType.PINNED: _t.CU_MEM_ALLOCATION_TYPE_PINNED} # noqa: RUF012 (empty)

117 if binding_version() >= (13, 0, 0): (empty)

118 _allocation_type[VirtualMemoryAllocationType.MANAGED] = _t.CU_MEM_ALLOCATION_TYPE_MANAGED (empty)

119 

120 @staticmethod (empty)

121 def _access_to_flags(spec: VirtualMemoryAccessType | None) -> int: (empty)

122 flags = VirtualMemoryResourceOptions._access_flags.get(spec) 6 ctx1cdafbk

123 if flags is None: 6 ctx1cdafbk

124 raise ValueError(f"Unknown access spec: {spec!r}") 1 ctx1k

125 return flags # type: ignore[no-any-return] 5 ctx1cdafb

126 

127 @staticmethod (empty)

128 def _allocation_type_to_driver(spec: VirtualMemoryAllocationType) -> int: (empty)

129 alloc_type = VirtualMemoryResourceOptions._allocation_type.get(spec) 5 ctx1cdabl

130 if alloc_type is None: 5 ctx1cdabl

131 raise ValueError(f"Unsupported allocation_type: {spec!r}") 1 ctx1l

132 return alloc_type # type: ignore[no-any-return] 4 ctx1cdab

133 

134 @staticmethod (empty)

135 def _location_type_to_driver(spec: VirtualMemoryLocationType) -> int: (empty)

136 loc_type = VirtualMemoryResourceOptions._location_type.get(spec) 5 ctx1cdabm

137 if loc_type is None: 5 ctx1cdabm

138 raise ValueError(f"Unsupported location_type: {spec!r}") 1 ctx1m

139 return loc_type # type: ignore[no-any-return] 4 ctx1cdab

140 

141 @staticmethod (empty)

142 def _handle_type_to_driver(spec: VirtualMemoryHandleType | None) -> int: (empty)

143 if spec == "win32": 6 ctx1cdaboj

144 raise NotImplementedError("win32 is currently not supported, please reach out to the CUDA Python team") 1 ctx1o

145 handle_type = VirtualMemoryResourceOptions._handle_types.get(spec) 5 ctx1cdabj

146 if handle_type is None: 5 ctx1cdabj

147 raise ValueError(f"Unsupported handle_type: {spec!r}") 1 ctx1j

148 return handle_type # type: ignore[no-any-return] 4 ctx1cdab

149 

150 @staticmethod (empty)

151 def _granularity_to_driver(spec: VirtualMemoryGranularityType) -> int: (empty)

152 granularity = VirtualMemoryResourceOptions._granularity.get(spec) 5 ctx1cdabn

153 if granularity is None: 5 ctx1cdabn

154 raise ValueError(f"Unsupported granularity: {spec!r}") 1 ctx1n

155 return granularity # type: ignore[no-any-return] 4 ctx1cdab

156 

157 

158class VirtualMemoryResource(MemoryResource): (empty)

159 """Create a device memory resource that uses the CUDA VMM APIs to allocate memory. 

160 

161 Parameters 

162 ---------- 

163 device_id : Device | int 

164 Device for which a memory resource is constructed. 

165 

166 config : VirtualMemoryResourceOptions, optional 

167 A configuration object for the VirtualMemoryResource 

168 

169 

170 Warning 

171 ------- 

172 This is a low-level API that is provided only for convenience. Make sure you fully understand 

173 how CUDA Virtual Memory Management works before using this. Other MemoryResource subclasses 

174 in cuda.core should already meet the common needs. 

175 """ 

176 

177 def __init__(self, device_id: Device | int, config: VirtualMemoryResourceOptions | None = None) -> None: (empty)

178 self.device: Device | None = Device(device_id) 8 ctx1cdafbghi

179 self.config: VirtualMemoryResourceOptions = check_or_create_options( # type: ignore[assignment] 8 ctx1cdafbghi

180 VirtualMemoryResourceOptions, config, "VirtualMemoryResource options", keep_none=False 

181 ) 

182 if self.config.location_type in _HOST_LOCATION_TYPES: 8 ctx1cdafbghi

183 self.device = None 3 ctx1ghi

184 

185 if not self.device and self.config.location_type == "device": 8 ctx1cdafbghi

186 raise RuntimeError("VirtualMemoryResource requires a device for device memory allocations") 

187 

188 if self.device and not self.device.properties.virtual_memory_management_supported: 8 ctx1cdafbghi

189 raise RuntimeError("VirtualMemoryResource requires CUDA VMM API support") 

190 

191 # Validate RDMA support if requested 

192 if ( 

193 self.config.gpu_direct_rdma 

194 and self.device is not None 

195 and not self.device.properties.gpu_direct_rdma_supported 

196 ): 

197 raise RuntimeError("GPU Direct RDMA is not supported on this device") 

198 

199 @staticmethod (empty)

200 def _align_up(size: int, gran: int) -> int: (empty)

201 """ 

202 Align a size up to the nearest multiple of a granularity. 

203 """ 

204 return (size + gran - 1) & ~(gran - 1) 4 ctx1cdab

205 

206 def modify_allocation( (empty)

207 self, buf: Buffer, new_size: int, config: VirtualMemoryResourceOptions | None = None 

208 ) -> Buffer: 

209 """ 

210 Grow an existing allocation using CUDA VMM, with a configurable policy. 

211 

212 This implements true growing allocations that preserve the base pointer 

213 by extending the virtual address range and mapping additional physical memory. 

214 

215 This function uses transactional allocation: if any step fails, the original buffer is not modified and 

216 all steps the function took are rolled back so a new allocation is not created. 

217 

218 Parameters 

219 ---------- 

220 buf : Buffer 

221 The existing buffer to grow 

222 new_size : int 

223 The new total size for the allocation 

224 config : VirtualMemoryResourceOptions, optional 

225 Configuration for the new physical memory chunks. If None, uses current config. 

226 

227 Returns 

228 ------- 

229 Buffer 

230 The same buffer with updated size and properties, preserving the original pointer 

231 """ 

232 if not isinstance(buf, Buffer): 2 ctx1ab

233 raise TypeError(f"buf must be a Buffer, got {type(buf).__name__}") 

234 if buf.is_closed: 2 ctx1ab

235 raise RuntimeError("Buffer has been closed") 

236 if config is not None: 2 ctx1ab

237 self.config = config 1 ctx1b

238 

239 # Build allocation properties for new chunks 

240 prop = driver.CUmemAllocationProp() 2 ctx1ab

241 prop.type = VirtualMemoryResourceOptions._allocation_type_to_driver(self.config.allocation_type) 2 ctx1ab

242 prop.location.type = VirtualMemoryResourceOptions._location_type_to_driver(self.config.location_type) 2 ctx1ab

243 # Caller must not invoke modify_allocation on a host-located resource; 

244 # we rely on the dataclass invariant that self.device is non-None for 

245 # device-located resources (it's only None when location is host). 

246 assert self.device is not None, "modify_allocation requires a device-located resource" 2 ctx1ab

247 prop.location.id = self.device.device_id 2 ctx1ab

248 prop.allocFlags.gpuDirectRDMACapable = 1 if self.config.gpu_direct_rdma else 0 2 ctx1ab

249 prop.requestedHandleTypes = VirtualMemoryResourceOptions._handle_type_to_driver(self.config.handle_type) 2 ctx1ab

250 prop.win32HandleMetaData = 0 2 ctx1ab

251 

252 # Query granularity 

253 gran_flag = VirtualMemoryResourceOptions._granularity_to_driver(self.config.granularity) 2 ctx1ab

254 res, gran = driver.cuMemGetAllocationGranularity(prop, gran_flag) 2 ctx1ab

255 raise_if_driver_error(res) 2 ctx1ab

256 

257 # Calculate sizes 

258 additional_size = new_size - buf.size 2 ctx1ab

259 if additional_size <= 0: 2 ctx1ab

260 # Same size: only update access policy if needed; avoid zero-sized driver calls 

261 descs = self._build_access_descriptors(prop) 2 ctx1ab

262 if descs: 2 ctx1ab

263 (res,) = driver.cuMemSetAccess(int(buf.handle), buf.size, descs, len(descs)) 2 ctx1ab

264 raise_if_driver_error(res) 2 ctx1ab

265 return buf 2 ctx1ab

266 

267 aligned_additional_size = VirtualMemoryResource._align_up(additional_size, gran) 1 ctx1a

268 total_aligned_size = VirtualMemoryResource._align_up(new_size, gran) 1 ctx1a

269 aligned_prev_size = total_aligned_size - aligned_additional_size 1 ctx1a

270 addr_align = self.config.addr_align or gran 1 ctx1a

271 

272 # Try to extend the existing VA range first 

273 res, new_ptr = driver.cuMemAddressReserve( 1 ctx1a

274 aligned_additional_size, 

275 addr_align, 

276 int(buf.handle) + aligned_prev_size, # fixedAddr hint - aligned end of current range 

277 0, 

278 ) 

279 

280 if res != driver.CUresult.CUDA_SUCCESS or new_ptr != (int(buf.handle) + aligned_prev_size): 1 ctx1a

281 # Check for specific errors that are not recoverable with the slow path 

282 if res in ( 1 ctx1a

283 driver.CUresult.CUDA_ERROR_INVALID_VALUE, 

284 driver.CUresult.CUDA_ERROR_NOT_PERMITTED, 

285 driver.CUresult.CUDA_ERROR_NOT_INITIALIZED, 

286 driver.CUresult.CUDA_ERROR_NOT_SUPPORTED, 

287 ): 

288 raise_if_driver_error(res) 

289 (res2,) = driver.cuMemAddressFree(new_ptr, aligned_additional_size) 1 ctx1a

290 raise_if_driver_error(res2) 1 ctx1a

291 # Fallback: couldn't extend contiguously, need full remapping 

292 return self._grow_allocation_slow_path( 1 ctx1a

293 buf, new_size, prop, aligned_additional_size, total_aligned_size, addr_align 

294 ) 

295 else: 

296 # Success! We can extend the VA range contiguously 

297 return self._grow_allocation_fast_path(buf, new_size, prop, aligned_additional_size, new_ptr) 

298 

299 def _grow_allocation_fast_path( (empty)

300 self, buf: Buffer, new_size: int, prop: driver.CUmemAllocationProp, aligned_additional_size: int, new_ptr: int 

301 ) -> Buffer: 

302 """ 

303 Fast path for growing a virtual memory allocation when the new region can be 

304 reserved contiguously after the existing buffer. 

305 

306 This function creates and maps new physical memory for the additional size, 

307 sets access permissions, and updates the buffer size in place (the pointer 

308 remains unchanged). 

309 

310 Args: 

311 buf (Buffer): 

312 The buffer to grow. 

313 

314 new_size (int): 

315 The new total size in bytes. 

316 

317 prop (driver.CUmemAllocationProp): 

318 Allocation properties for the new memory. 

319 

320 aligned_additional_size (int): 

321 The size of the new region to allocate, aligned to granularity. 

322 

323 new_ptr (int): 

324 The address of the newly reserved contiguous VA region (should 

325 be at the end of the current buffer). 

326 

327 Returns: 

328 Buffer: The same buffer object with its size updated to `new_size`. 

329 """ 

330 with Transaction() as trans: 1 ctx1f

331 # Create new physical memory for the additional size 

332 trans.append( 1 ctx1f

333 lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0]) 

334 ) 

335 res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) 1 ctx1f

336 raise_if_driver_error(res) 1 ctx1f

337 # Register undo for creation 

338 trans.append(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) 1 ctx1f

339 

340 # Map the new physical memory to the extended VA range 

341 (res,) = driver.cuMemMap(new_ptr, aligned_additional_size, 0, new_handle, 0) 1 ctx1f

342 raise_if_driver_error(res) 1 ctx1f

343 # Register undo for mapping 

344 trans.append( 1 ctx1f

345 lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0]) 

346 ) 

347 

348 # Set access permissions for the new portion 

349 descs = self._build_access_descriptors(prop) 1 ctx1f

350 if descs: 1 ctx1f

351 (res,) = driver.cuMemSetAccess(new_ptr, aligned_additional_size, descs, len(descs)) 1 ctx1f

352 raise_if_driver_error(res) 1 ctx1f

353 

354 # All succeeded, cancel undo actions 

355 trans.commit() 1 ctx1f

356 

357 # Update the buffer size (pointer stays the same). `Buffer.size` has 

358 # no public setter, so this reaches into the private attribute. 

359 buf._size = new_size 1 ctx1f

360 return buf 1 ctx1f

361 

362 def _grow_allocation_slow_path( (empty)

363 self, 

364 buf: Buffer, 

365 new_size: int, 

366 prop: driver.CUmemAllocationProp, 

367 aligned_additional_size: int, 

368 total_aligned_size: int, 

369 addr_align: int, 

370 ) -> Buffer: 

371 """ 

372 Slow path for growing a virtual memory allocation when the new region cannot be 

373 reserved contiguously after the existing buffer. 

374 

375 This function reserves a new, larger virtual address (VA) range, remaps the old 

376 physical memory to the beginning of the new VA range, creates and maps new physical 

377 memory for the additional size, sets access permissions, and updates the buffer's 

378 pointer and size. 

379 

380 Args: 

381 buf (Buffer): The buffer to grow. 

382 new_size (int): The new total size in bytes. 

383 prop (driver.CUmemAllocationProp): Allocation properties for the new memory. 

384 aligned_additional_size (int): The size of the new region to allocate, aligned to granularity. 

385 total_aligned_size (int): The total new size to reserve, aligned to granularity. 

386 addr_align (int): The required address alignment for the new VA range. 

387 

388 Returns: 

389 Buffer: The buffer object updated with the new pointer and size. 

390 """ 

391 with Transaction() as trans: 1 ctx1a

392 # Reserve a completely new, larger VA range 

393 res, new_ptr = driver.cuMemAddressReserve(total_aligned_size, addr_align, 0, 0) 1 ctx1a

394 raise_if_driver_error(res) 1 ctx1a

395 # Register undo for VA reservation 

396 trans.append( 1 ctx1a

397 lambda np=new_ptr, s=total_aligned_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0]) 

398 ) 

399 

400 # Get the old allocation handle for remapping 

401 result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle) 1 ctx1a

402 raise_if_driver_error(result) 1 ctx1a

403 # Register undo for old_handle 

404 trans.append(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) 1 ctx1a

405 

406 # Unmap the old VA range (aligned previous size) 

407 aligned_prev_size = total_aligned_size - aligned_additional_size 1 ctx1a

408 (result,) = driver.cuMemUnmap(int(buf.handle), aligned_prev_size) 1 ctx1a

409 raise_if_driver_error(result) 1 ctx1a

410 

411 def _remap_old() -> None: 1 ctx1a

412 # Try to remap the old physical memory back to the original VA range 

413 try: 

414 (res,) = driver.cuMemMap(int(buf.handle), aligned_prev_size, 0, old_handle, 0) 

415 raise_if_driver_error(res) 

416 except Exception: # noqa: S110 

417 # TODO: consider logging this exception 

418 pass 

419 

420 trans.append(_remap_old) 1 ctx1a

421 

422 # Remap the old physical memory to the new VA range (aligned previous size) 

423 (res,) = driver.cuMemMap(int(new_ptr), aligned_prev_size, 0, old_handle, 0) 1 ctx1a

424 raise_if_driver_error(res) 1 ctx1a

425 

426 # Register undo for mapping 

427 trans.append(lambda np=new_ptr, s=aligned_prev_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0])) 1 ctx1a

428 

429 # Create new physical memory for the additional size 

430 res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) 1 ctx1a

431 raise_if_driver_error(res) 1 ctx1a

432 

433 # Register undo for new physical memory 

434 trans.append(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) 1 ctx1a

435 

436 # Map the new physical memory to the extended portion (aligned offset) 

437 (res,) = driver.cuMemMap(int(new_ptr) + aligned_prev_size, aligned_additional_size, 0, new_handle, 0) 1 ctx1a

438 raise_if_driver_error(res) 1 ctx1a

439 

440 # Register undo for mapping 

441 trans.append( 1 ctx1a

442 lambda base=int(new_ptr), offs=aligned_prev_size, s=aligned_additional_size: raise_if_driver_error( 

443 driver.cuMemUnmap(base + offs, s)[0] 

444 ) 

445 ) 

446 

447 # Set access permissions for the entire new range 

448 descs = self._build_access_descriptors(prop) 1 ctx1a

449 if descs: 1 ctx1a

450 (res,) = driver.cuMemSetAccess(new_ptr, total_aligned_size, descs, len(descs)) 1 ctx1a

451 raise_if_driver_error(res) 1 ctx1a

452 

453 # All succeeded, cancel undo actions 

454 trans.commit() 1 ctx1a

455 

456 # Free the old VA range (aligned previous size) 

457 (res2,) = driver.cuMemAddressFree(int(buf.handle), aligned_prev_size) 1 ctx1a

458 raise_if_driver_error(res2) 1 ctx1a

459 

460 # Invalidate the old buffer so its destructor won't try to free again 

461 buf._clear() 1 ctx1a

462 

463 # Return a new Buffer for the new mapping 

464 return Buffer.from_handle(ptr=new_ptr, size=new_size, mr=self) 1 ctx1a

465 

466 def _build_access_descriptors(self, prop: driver.CUmemAllocationProp) -> list[driver.CUmemAccessDesc]: (empty)

467 """ 

468 Build access descriptors for memory access permissions. 

469 

470 Returns 

471 ------- 

472 list 

473 List of CUmemAccessDesc objects for setting memory access 

474 """ 

475 descs = [] 5 ctx1cdafb

476 

477 # Owner access 

478 owner_flags = VirtualMemoryResourceOptions._access_to_flags(self.config.self_access) 5 ctx1cdafb

479 if owner_flags: 5 ctx1cdafb

480 d = driver.CUmemAccessDesc() 5 ctx1cdafb

481 d.location.type = prop.location.type 5 ctx1cdafb

482 d.location.id = prop.location.id 5 ctx1cdafb

483 d.flags = owner_flags 5 ctx1cdafb

484 descs.append(d) 5 ctx1cdafb

485 

486 # Peer device access 

487 peer_flags = VirtualMemoryResourceOptions._access_to_flags(self.config.peer_access) 5 ctx1cdafb

488 if peer_flags: 5 ctx1cdafb

489 for peer_dev in self.config.peers: 5 ctx1cdafb

490 d = driver.CUmemAccessDesc() 

491 d.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

492 d.location.id = int(peer_dev) 

493 d.flags = peer_flags 

494 descs.append(d) 

495 

496 return descs 5 ctx1cdafb

497 

498 def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> Buffer: (empty)

499 """ 

500 Allocate a buffer of the given size using CUDA virtual memory. 

501 

502 Parameters 

503 ---------- 

504 size : int 

505 The size in bytes of the buffer to allocate. 

506 stream : Stream, optional 

507 Keyword-only. Unused because virtual memory operations are 

508 synchronous. 

509 

510 Returns 

511 ------- 

512 Buffer 

513 A Buffer object representing the allocated virtual memory. 

514 

515 Raises 

516 ------ 

517 CUDAError 

518 If any CUDA driver API call fails during allocation. 

519 

520 Notes 

521 ----- 

522 This method uses transactional allocation: if any step fails, all resources 

523 allocated so far are automatically cleaned up. The allocation is performed 

524 with the configured granularity, access permissions, and peer access as 

525 specified in the resource's configuration. 

526 """ 

527 if stream is not None: 4 ctx1cdab

528 from cuda.core._stream import Stream_accept 

529 

530 Stream_accept(stream) 

531 

532 config = self.config 4 ctx1cdab

533 # ---- Build allocation properties ---- 

534 prop = driver.CUmemAllocationProp() 4 ctx1cdab

535 prop.type = VirtualMemoryResourceOptions._allocation_type_to_driver(config.allocation_type) 4 ctx1cdab

536 prop.location.type = VirtualMemoryResourceOptions._location_type_to_driver(config.location_type) 4 ctx1cdab

537 prop.location.id = self.device.device_id if self.device is not None else -1 4 ctx1cdab

538 prop.allocFlags.gpuDirectRDMACapable = 1 if config.gpu_direct_rdma else 0 4 ctx1cdab

539 prop.requestedHandleTypes = VirtualMemoryResourceOptions._handle_type_to_driver(config.handle_type) 4 ctx1cdab

540 prop.win32HandleMetaData = 0 4 ctx1cdab

541 

542 # ---- Query and apply granularity ---- 

543 # Choose min vs recommended granularity per config 

544 gran_flag = VirtualMemoryResourceOptions._granularity_to_driver(config.granularity) 4 ctx1cdab

545 res, gran = driver.cuMemGetAllocationGranularity(prop, gran_flag) 4 ctx1cdab

546 raise_if_driver_error(res) 4 ctx1cdab

547 

548 aligned_size = VirtualMemoryResource._align_up(size, gran) 4 ctx1cdab

549 addr_align = config.addr_align or gran 4 ctx1cdab

550 

551 # ---- Transactional allocation ---- 

552 with Transaction() as trans: 4 ctx1cdab

553 # ---- Create physical memory ---- 

554 res, handle = driver.cuMemCreate(aligned_size, prop, 0) 4 ctx1cdab

555 raise_if_driver_error(res) 4 ctx1cdab

556 # Register undo for physical memory 

557 trans.append(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) 4 ctx1cdab

558 

559 # ---- Reserve VA space ---- 

560 # Potentially, use a separate size for the VA reservation from the physical allocation size 

561 res, ptr = driver.cuMemAddressReserve(aligned_size, addr_align, config.addr_hint, 0) 4 ctx1cdab

562 raise_if_driver_error(res) 4 ctx1cdab

563 # Register undo for VA reservation 

564 trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemAddressFree(p, s)[0])) 4 ctx1cdab

565 

566 # ---- Map physical memory into VA ---- 

567 (res,) = driver.cuMemMap(ptr, aligned_size, 0, handle, 0) 4 ctx1cdab

568 trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0])) 4 ctx1cdab

569 raise_if_driver_error(res) 4 ctx1cdab

570 

571 # ---- Set access for owner + peers ---- 

572 descs = self._build_access_descriptors(prop) 4 ctx1cdab

573 if descs: 4 ctx1cdab

574 (res,) = driver.cuMemSetAccess(ptr, aligned_size, descs, len(descs)) 4 ctx1cdab

575 raise_if_driver_error(res) 4 ctx1cdab

576 

577 trans.commit() 4 ctx1cdab

578 

579 # Done — return a Buffer that tracks this VA range 

580 buf = Buffer.from_handle(ptr=ptr, size=aligned_size, mr=self) 4 ctx1cdab

581 return buf 4 ctx1cdab

582 

583 def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | GraphBuilder | None = None) -> None: (empty)

584 """ 

585 Deallocate memory on the device using CUDA VMM APIs. 

586 

587 Parameters 

588 ---------- 

589 ptr : DevicePointerType 

590 The pointer to the memory to deallocate. 

591 size : int 

592 The size in bytes of the memory to deallocate. 

593 stream : Stream, optional 

594 Keyword-only. Unused because virtual memory operations are 

595 synchronous. 

596 """ 

597 ptr = 0 if ptr is None else int(ptr) 4 ctx1cdab

598 

599 if stream is not None: 4 ctx1cdab

600 from cuda.core._stream import Stream_accept 4 ctx1cdab

601 

602 Stream_accept(stream) 4 ctx1cdab

603 result, handle = driver.cuMemRetainAllocationHandle(ptr) 4 ctx1cdab

604 raise_if_driver_error(result) 4 ctx1cdab

605 (result,) = driver.cuMemUnmap(ptr, size) 4 ctx1cdab

606 raise_if_driver_error(result) 4 ctx1cdab

607 (result,) = driver.cuMemAddressFree(ptr, size) 4 ctx1cdab

608 raise_if_driver_error(result) 4 ctx1cdab

609 (result,) = driver.cuMemRelease(handle) 4 ctx1cdab

610 raise_if_driver_error(result) 4 ctx1cdab

611 

612 @property (empty)

613 def is_device_accessible(self) -> bool: (empty)

614 """ 

615 Indicates whether the allocated memory is accessible from the device. 

616 """ 

617 return self.config.location_type == "device" 

618 

619 @property (empty)

620 def is_host_accessible(self) -> bool: (empty)

621 """ 

622 Indicates whether the allocated memory is accessible from the host. 

623 """ 

624 return self.config.location_type in _HOST_LOCATION_TYPES 3 ctx1ghi

625 

626 @property (empty)

627 def device_id(self) -> int: (empty)

628 """ 

629 Get the device ID associated with this memory resource. 

630 

631 Returns: 

632 int: CUDA device ID. -1 if the memory resource allocates host memory 

633 """ 

634 return self.device.device_id if self.device is not None else -1 3 ctx1cdb

635 

636 def __repr__(self) -> str: (empty)

637 """ 

638 Return a string representation of the VirtualMemoryResource. 

639 

640 Returns: 

641 str: A string describing the object 

642 """ 

643 return f"<VirtualMemoryResource device={self.device}>"