From 2b4c46f8ab5c314d254bdc842228b87ba3bc7a49 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Fri, 18 Sep 2026 12:40:37 -0700 Subject: [PATCH 1/6] fix(cuda.core): move VirtualMemoryResource onto the _rt handle layer Each physical allocation, address reservation and mapping now lives in a std::shared_ptr handle whose deleter knows the exact driver call to undo it. A buffer owns a range of mappings through its device pointer handle, so everything a buffer maps is released when the last buffer that maps it closes, and a failed multi-step operation unwinds by letting its local handles die. The module moves from Python to Cython. The design is in cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md. Behavior changes: - modify_allocation returns a new VirtualMemoryBuffer and leaves the input open; the two alias the same physical memory, which is freed when the last of them closes. The pointer is preserved when the driver grants the adjacent address range. - Buffer.size after a grow is the aligned total. - config= applies to the chunk the call adds and is not stored on the resource. - Buffers from allocate() free themselves on close and do not call deallocate(), which now serves pointers wrapped with Buffer.from_handle. - A buffer records the stream passed to allocate(); the last close of an aliased range synchronizes every recorded stream before it unmaps. An explicit close on a capturing stream raises. - location_type="host" requires handle_type=None. allocate(0) returns an empty buffer without a driver call. Fixes #2887 Fixes #2907 Fixes #2908 Fixes #2909 Fixes #2886 Fixes #2345 Addresses #2388 item 2 and the size-0, misaligned-probe and host handle-type parts of #2910. Part of #2906. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/DESIGN.md | 5 + cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 189 +++++ cuda_core/cuda/core/_cpp/rt/api.hpp | 64 ++ cuda_core/cuda/core/_cpp/rt/driver_api.cpp | 11 + cuda_core/cuda/core/_cpp/rt/driver_api.hpp | 11 + cuda_core/cuda/core/_cpp/rt/internal.hpp | 3 + cuda_core/cuda/core/_cpp/rt/memory.cpp | 26 + cuda_core/cuda/core/_cpp/rt/py.hpp | 12 + cuda_core/cuda/core/_cpp/rt/types.hpp | 40 + .../cuda/core/_cpp/rt/virtual_memory.cpp | 316 ++++++++ cuda_core/cuda/core/_cpp/rt/vmm_range.hpp | 38 + cuda_core/cuda/core/_memory/_buffer.pxd | 5 +- cuda_core/cuda/core/_memory/_buffer.pyi | 2 - .../core/_memory/_virtual_memory_resource.py | 636 ---------------- .../core/_memory/_virtual_memory_resource.pyi | 250 +++++++ .../core/_memory/_virtual_memory_resource.pyx | 701 ++++++++++++++++++ cuda_core/cuda/core/_rt.pxd | 48 ++ cuda_core/cuda/core/_rt.pyi | 8 + cuda_core/cuda/core/_rt.pyx | 54 ++ cuda_core/cuda/core/_utils/cuda_utils.pyi | 32 - cuda_core/cuda/core/_utils/cuda_utils.pyx | 67 -- cuda_core/docs/source/api.rst | 1 + cuda_core/docs/source/release/1.3.0-notes.rst | 38 + .../graph/test_graph_definition_lifetime.py | 46 +- cuda_core/tests/test_memory.py | 567 ++++++++++---- 25 files changed, 2279 insertions(+), 891 deletions(-) create mode 100644 cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md create mode 100644 cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp create mode 100644 cuda_core/cuda/core/_cpp/rt/vmm_range.hpp delete mode 100644 cuda_core/cuda/core/_memory/_virtual_memory_resource.py create mode 100644 cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi create mode 100644 cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx diff --git a/cuda_core/cuda/core/_cpp/rt/DESIGN.md b/cuda_core/cuda/core/_cpp/rt/DESIGN.md index 0374d5a08be..9d167b708a5 100644 --- a/cuda_core/cuda/core/_cpp/rt/DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/DESIGN.md @@ -63,6 +63,11 @@ Internally, handles use **shared pointer aliasing**: the actual managed object i "box" containing the resource, its dependencies, and any state needed for destruction. The public handle points only to the raw resource field, keeping the API minimal. +The virtual memory resource adds `MemAllocationHandle`, `VaReservationHandle` and +`VaMappingHandle`. Their values are `TaggedHandle` wrappers, because +`CUmemGenericAllocationHandle` and `CUdeviceptr` are both `unsigned long long` and the +accessor overloads must stay distinct. See [VMM_DESIGN.md](VMM_DESIGN.md). + ### Why shared_ptr? - **Automatic reference counting**: Resources are released when the last reference diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md new file mode 100644 index 00000000000..d7d1ee3b107 --- /dev/null +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -0,0 +1,189 @@ +# VirtualMemoryResource on the handle layer + +This document describes how `VirtualMemoryResource` uses the `_rt` handle layer that the +pool-backed resources also use. It complements [DESIGN.md](DESIGN.md), which describes the layer +itself. + +## Summary + +Each physical allocation, address reservation, and mapping has its own `std::shared_ptr` handle +with a deleter that knows the exact driver call to undo it. A buffer owns a list of mappings. +Teardown order follows from what holds what, and a failed multi-step operation unwinds by letting +its local handles die. The module is written in Cython, because the handles are `cdef` types. + +## Driver behavior the design relies on + +- **Reservation.** `cuMemAddressReserve(size, align, hint)` returns `(ptr, size)`. + [`cuMemAddressFree`](https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html) + succeeds only for the exact `(ptr, size)` pair of one reservation. Reservations never overlap; + growing a buffer in place yields two adjacent reservations, each of which is freed separately. + A hint must be a multiple of `max(align, 2 MiB)`; `align == 0` means the default 2 MiB. +- **Physical allocation.** `cuMemCreate` returns a handle with one reference; `cuMemRelease` drops + one; `cuMemRetainAllocationHandle` adds one. Mappings are counted separately. The memory is + freed when references are zero and no mapping remains. Releasing while mapped is legal. +- **Mapping.** `cuMemMap(ptr, size, 0, handle)` maps the whole allocation at `ptr`; offset must be + 0 and size must equal the allocation's size. `cuMemSetAccess(ptr, size, descs, count)` grants + access per mapped range and rejects `count == 0`. `cuMemUnmap(ptr, size)` may cover several + whole mappings, never part of one. One allocation may be mapped at several addresses at once; + each mapping has its own access state. +- **Coherence of aliases.** Two addresses that map one allocation reach the same physical pages. + Writes through one are visible through the other in stream order and at kernel boundaries; the + driver adds no synchronization of its own. +- **Context and synchronization.** No VMM entry point needs a current context. `cuMemUnmap` does + not synchronize. `cuStreamSynchronize` is rejected on a capturing stream. + +So a mapping depends on exactly one reservation and one allocation, mappings never depend on +other mappings, and reservations and allocations are independent of each other. A buffer is a +list of mappings. + +## Handles + +Three new `std::shared_ptr` aliases into boxes, following the conventions in `types.hpp`. +`CUmemGenericAllocationHandle` and `CUdeviceptr` are both `unsigned long long`, so the values are +wrapped in `TaggedHandle` to keep the overload sets distinct. + +| Handle | Box | Deleter | Depends on | +|---|---|---|---| +| `MemAllocationHandle` | `{handle, size, access descriptors}` | `pw_cuMemRelease(handle)` | nothing | +| `VaReservationHandle` | `{ptr, size}` | `pw_cuMemAddressFree(ptr, size)` | nothing | +| `VaMappingHandle` | `{ptr, size, h_alloc, h_reservation}` | `pw_cuMemUnmap(ptr, size)`, then the members release | allocation, reservation | + +The allocation box carries the access descriptors it was created with. A mapping applies its +allocation's descriptors (and skips the call when there are none), so a chunk keeps its access +wherever it is mapped. + +``` +MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocationProp& prop, + const CUmemAccessDesc* descs, size_t count); +VaReservationHandle create_va_reservation_handle(size_t size, size_t align, CUdeviceptr hint); +VaMappingHandle create_va_mapping_handle(CUdeviceptr ptr, const MemAllocationHandle& h_alloc, + const VaReservationHandle& h_res); +size_t mem_allocation_size(...) noexcept; size_t va_reservation_size(...) noexcept; +``` + +Factories return the handle and put the status in thread-local `err`, as the other factories do; +an empty input handle sets `err` too, so an empty result always carries a status. +`create_va_mapping_handle` checks that the range lies inside the reservation, maps, and applies +access; if access fails it unmaps and returns empty. The eight VMM entry points, plus +`cuStreamSynchronize` and `cuStreamGetCaptureInfo`, join the `driver_api` pointer table. + +### The range and the device pointer + +``` +struct VmmRange { // one per buffer base address + std::vector mappings; // ascending, contiguous; sum of sizes = range total + std::mutex mu; // guards `streams`; nothing under it takes the GIL + std::vector streams; // every stream a dying owner forwarded, deduplicated +}; +using VmmRangeHandle = std::shared_ptr; // deleter: sync each stream, then destroy +DevicePtrHandle deviceptr_create_vmm(CUdeviceptr base, VmmRangeHandle range); +VmmRangeHandle vmm_range(const DevicePtrHandle& h); // empty for a non-VMM or closed handle +``` + +There is exactly one ownership chain: `Buffer._h_ptr` -> `DevicePtrBox` (holds the range) -> +`VmmRange` -> mappings -> reservations and allocations. The Cython buffer keeps no other +reference; grow operations call `Buffer_check_open` and then `vmm_range(buf._h_ptr)`. +`Buffer.close()` stays `_h_ptr.reset()`. A buffer's `size` is always a prefix of its range. + +The `DevicePtrHandle` must own the memory because graph memcpy nodes retain `buf._h_ptr` as an +opaque owner; a non-owning handle would let a launched graph outlive its buffer. Any number of +owners may therefore exist at once: aliases from a grow, and graph attachments. + +Deleters: + +- `DevicePtrBox`, VMM flavor: release the GIL; append this box's recorded `DeallocationStream` to + the range under `mu` (skipping an empty stream and duplicates); release `mu`; drop the range + reference. It never blocks. +- `VmmRange`: release the GIL; unless the interpreter is finalizing, for each forwarded stream + check the capture status and skip the stream with one report when a sync would disturb a + capture (the stream is capturing, or it is the legacy stream while a blocking stream in its + context is capturing), otherwise synchronize it under its bound context. Then, whether or not the syncs succeeded, + destroy the mappings. Each mapping unmaps; the reservations free and the allocations release as + their last references go. Every forwarded stream is synchronized because two aliases may have + recorded different streams; synchronizing only the last one to die would unmap under work + queued on the other. This is the first blocking deleter in the layer, and it may run inside + the deferred-cleanup drain on the main thread, with the GIL released. + +Allocations are shared by two ranges after a grow that moves the buffer. Shared ownership is +what makes that safe: the allocation is released exactly once, when its last mapping goes. + +## The resource + +- `VirtualMemoryResourceOptions` describes the allocations. `__init__` rejects `location_type="host"` with a + handle type other than `None`, which the driver rejects, and keeps the RDMA and VMM-support + checks. The resource reports `is_ipc_enabled = False`, which `Buffer.ipc_descriptor` reads. +- `cdef class VirtualMemoryBuffer(Buffer)` carries no extra state. It is created with + `Buffer_from_deviceptr_handle(h_ptr, size, self, cls=VirtualMemoryBuffer)` and documented in + `api.rst` like `ManagedBuffer`. It overrides `close(stream=None)` to reject a capturing stream, + since VMM deallocation is synchronous and cannot be captured. `allocate(0)` returns one with no + mapping. +- `allocate(size, *, stream=None)`: + 1. `size == 0` returns an empty buffer without a driver call, like the other resources. + 2. Build `CUmemAllocationProp` and the access descriptors from the options; query the + granularity; align the size. + 3. Create the allocation, the reservation, and the mapping as locals; on any empty handle, + `HANDLE_RETURN(get_last_error())`. The locals unwind everything. + 4. Build the range and the device pointer handle. + 5. Record the deallocation stream: the caller's stream if it is a real stream, otherwise the + legacy default token bound to the device's primary context, as `_SynchronousMemoryResource` + does. `allocate()` therefore never needs a current context. Host-located resources record no + stream and close without a sync. + 6. Return a `VirtualMemoryBuffer` whose `size` is the aligned size. +- `modify_allocation(buf, new_size, config=None)`: + - `Buffer_check_open(buf)`; `range = vmm_range(buf._h_ptr)`; an empty range means the buffer + did not come from this resource: `TypeError`. `cfg = config or self.config` governs the new + chunk only and is not stored on the resource. Let `req = align_up(new_size)` and `total` be + the range total. + - `req <= buf.size`: return `buf`. The buffer already covers the request. + - `buf.size < req <= total`: return a new `VirtualMemoryBuffer` over the same range with size + `req`; no driver call. This serves a shorter alias asking for what the range already maps. + - `req > total`, in place: probe `cuMemAddressReserve(req - total, align=0, hint=base+total)`. + If the driver grants the hint, create the new allocation with `cfg`'s descriptors and its + mapping as locals; `mappings.reserve(n+1)`; create a second `DevicePtrHandle` on the same + range, copying the input's recorded deallocation stream; build the new buffer; `push_back` + the mapping as the last, non-throwing step. If the driver grants another address, drop the + reservation and move the buffer instead. If the probe fails, drain the status with + `get_last_error()` and move the buffer. + - `req > total`, moved: reserve `req` with `addr_align`; map every mapping in the range (shared + allocation handles, their own descriptors) at `base_new + offset`; create and map the new + allocation; build the new range (copying the input's recorded stream), handle, and buffer. + - Both paths return a new buffer and leave the input open. See "Why `modify_allocation` + returns a new buffer" below. + - `modify_allocation` is not thread-safe with respect to two buffers that share a range; that + synchronization is the caller's responsibility, as elsewhere in cuda.core. +- `deallocate(ptr, size, *, stream=None)` stays for the `MemoryResource` contract. It serves + pointers wrapped with `Buffer.from_handle(ptr, size, mr=self)`: synchronize `stream` if given, + `cuMemUnmap`, `cuMemAddressFree`. It handles one reservation, and the caller must have released + its own `cuMemCreate` reference. It is not called for buffers from `allocate()`, whose ranges + free themselves. A subclass override of `deallocate()` therefore does not run for them. + +## Why `modify_allocation` returns a new buffer + +The input buffer stays open and aliases the result. The chunks the input already mapped are +shared by both buffers and are freed when the last of the two closes; the chunk the grow added +belongs to the result. When the grow happens in place, the two buffers share one range at one +base, and the shorter one pins the whole range until it closes. Callers who are done with the +input close it. + +The alternative, growing the input object in place so every holder sees the new size and +pointer, was rejected: + +- In-place update reaches only holders of the Python object. Holders of the handle, such as + graph memcpy nodes, DLPack capsules, and IPC descriptors, would keep the old mapping alive but + see a different address than the buffer reports. +- An address change should be visible. When the buffer moves, an object that quietly changes + address turns cached `int(buf.handle)` values into dangling pointers. +- `Buffer.__hash__` and `__eq__` include the size, so in-place growth changes the hash of a live + object. +- Leaving the input open costs the caller one line and gives them a valid, shorter alias, which + no in-place scheme can offer. + +## Failure handling + +- Rollback is RAII: locals die in reverse order, deleters run the `pw_*` wrappers, and failures + become `CUDAWarning`. +- Every deleter releases the GIL first; the range deleter holds no C++ lock while it synchronizes + or reports. A failed sync (lost context, capturing stream) is reported and the unmap proceeds; + the driver needs no context for it, so nothing leaks. +- Empty handles always carry a status; the in-place probe drains its status before falling back. +- Factories that allocate are declared `except+` in `_rt.pxd`; deleters only destroy vectors. diff --git a/cuda_core/cuda/core/_cpp/rt/api.hpp b/cuda_core/cuda/core/_cpp/rt/api.hpp index c52e0851cd1..e65617949a4 100644 --- a/cuda_core/cuda/core/_cpp/rt/api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/api.hpp @@ -240,6 +240,70 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept; CUresult set_deallocation_stream( const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept; +// ============================================================================ +// Virtual memory management (VMM_DESIGN.md) +// +// A VirtualMemoryResource buffer is a range of mappings. Each mapping holds +// one physical allocation and one address reservation; the mapping deleter +// unmaps, then the allocation is released and the reservation freed as their +// last references go. A buffer's DevicePtrHandle owns the range. +// ============================================================================ + +// Create a physical allocation via cuMemCreate. The access descriptors are +// applied to every mapping of this allocation. When the last reference is +// released, cuMemRelease is called; the memory is freed once no mapping +// remains. Returns empty handle on error (caller must check). +MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocationProp& prop, + const CUmemAccessDesc* descs, size_t count); + +// Size of the allocation; the only size cuMemMap accepts for it. +size_t mem_allocation_size(const MemAllocationHandle& h) noexcept; + +// Reserve an address range via cuMemAddressReserve. Pass alignment 0 for the +// driver default. When the last reference is released, cuMemAddressFree is +// called with the exact reserved pair. Returns empty handle on error. +VaReservationHandle create_va_reservation_handle(size_t size, size_t alignment, CUdeviceptr hint); + +// Size of the reservation. +size_t va_reservation_size(const VaReservationHandle& h) noexcept; + +// Map the whole allocation at ptr inside the reservation via cuMemMap and +// apply the allocation's access descriptors. The mapping structurally depends +// on both handles. When the last reference is released, cuMemUnmap is called +// first. Returns empty handle on error, including a range outside the +// reservation; a failed cuMemSetAccess unmaps before returning. +VaMappingHandle create_va_mapping_handle(CUdeviceptr ptr, const MemAllocationHandle& h_alloc, + const VaReservationHandle& h_res); + +// Mapping accessors. +size_t va_mapping_size(const VaMappingHandle& h) noexcept; +MemAllocationHandle va_mapping_allocation(const VaMappingHandle& h) noexcept; + +// Create an empty range keyed by its base address. When the last reference is +// released, every deallocation stream its owners recorded is synchronized +// (capturing streams are skipped and reported), then the mappings are +// destroyed. May throw std::bad_alloc. +VmmRangeHandle create_vmm_range(CUdeviceptr base); + +// Recover the range of a VMM device pointer handle; empty for any other +// handle, including a closed one. +VmmRangeHandle vmm_range(const DevicePtrHandle& h); + +// Range accessors and mutators. Mutation is not synchronized: two buffers +// that share a range must not be grown concurrently (caller's responsibility). +size_t vmm_range_count(const VmmRangeHandle& range) noexcept; +VaMappingHandle vmm_range_mapping(const VmmRangeHandle& range, size_t index) noexcept; +size_t vmm_range_total(const VmmRangeHandle& range) noexcept; +void vmm_range_reserve(const VmmRangeHandle& range, size_t count); // may throw +void vmm_range_append(const VmmRangeHandle& range, const VaMappingHandle& mapping); // may throw + +// Create a device pointer handle that owns a range. The box records no +// deallocation stream; set one with set_deallocation_stream. When the last +// reference is released, the recorded stream is forwarded to the range and +// the range reference dropped; the range deleter does the synchronization and +// the unmapping. Returns empty handle for an empty range. +DevicePtrHandle deviceptr_create_vmm(CUdeviceptr base, const VmmRangeHandle& range); + // ============================================================================ // Library handle functions // ============================================================================ diff --git a/cuda_core/cuda/core/_cpp/rt/driver_api.cpp b/cuda_core/cuda/core/_cpp/rt/driver_api.cpp index 860a29a3554..3a98387dd36 100644 --- a/cuda_core/cuda/core/_cpp/rt/driver_api.cpp +++ b/cuda_core/cuda/core/_cpp/rt/driver_api.cpp @@ -60,6 +60,17 @@ decltype(&cuMemFreeHost) p_cuMemFreeHost = nullptr; decltype(&cuMemPoolImportPointer) p_cuMemPoolImportPointer = nullptr; +// Virtual memory management +decltype(&cuMemCreate) p_cuMemCreate = nullptr; +decltype(&cuMemRelease) p_cuMemRelease = nullptr; +decltype(&cuMemAddressReserve) p_cuMemAddressReserve = nullptr; +decltype(&cuMemAddressFree) p_cuMemAddressFree = nullptr; +decltype(&cuMemMap) p_cuMemMap = nullptr; +decltype(&cuMemUnmap) p_cuMemUnmap = nullptr; +decltype(&cuMemSetAccess) p_cuMemSetAccess = nullptr; +decltype(&cuStreamSynchronize) p_cuStreamSynchronize = nullptr; +decltype(&cuStreamGetCaptureInfo) p_cuStreamGetCaptureInfo = nullptr; + decltype(&cuLibraryLoadFromFile) p_cuLibraryLoadFromFile = nullptr; decltype(&cuLibraryLoadData) p_cuLibraryLoadData = nullptr; decltype(&cuLibraryUnload) p_cuLibraryUnload = nullptr; diff --git a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp index 87ed3bd7906..eec3a7fbd32 100644 --- a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp @@ -63,6 +63,17 @@ extern decltype(&cuMemFreeHost) p_cuMemFreeHost; extern decltype(&cuMemPoolImportPointer) p_cuMemPoolImportPointer; +// Virtual memory management (VMM_DESIGN.md) +extern decltype(&cuMemCreate) p_cuMemCreate; +extern decltype(&cuMemRelease) p_cuMemRelease; +extern decltype(&cuMemAddressReserve) p_cuMemAddressReserve; +extern decltype(&cuMemAddressFree) p_cuMemAddressFree; +extern decltype(&cuMemMap) p_cuMemMap; +extern decltype(&cuMemUnmap) p_cuMemUnmap; +extern decltype(&cuMemSetAccess) p_cuMemSetAccess; +extern decltype(&cuStreamSynchronize) p_cuStreamSynchronize; +extern decltype(&cuStreamGetCaptureInfo) p_cuStreamGetCaptureInfo; + // Library extern decltype(&cuLibraryLoadFromFile) p_cuLibraryLoadFromFile; extern decltype(&cuLibraryLoadData) p_cuLibraryLoadData; diff --git a/cuda_core/cuda/core/_cpp/rt/internal.hpp b/cuda_core/cuda/core/_cpp/rt/internal.hpp index 46b5a77f379..b8f17e6087d 100644 --- a/cuda_core/cuda/core/_cpp/rt/internal.hpp +++ b/cuda_core/cuda/core/_cpp/rt/internal.hpp @@ -123,6 +123,9 @@ const WarnOnFailure pw_cuSurfObjectDestroy{"cuSurfObjectD const WarnOnFailure pw_cuGreenCtxDestroy{"cuGreenCtxDestroy"}; const WarnOnFailure pw_cuMemPoolDestroy{"cuMemPoolDestroy"}; const WarnOnFailure pw_cuMemFreeHost{"cuMemFreeHost"}; +const WarnOnFailure pw_cuMemRelease{"cuMemRelease"}; +const WarnOnFailure pw_cuMemUnmap{"cuMemUnmap"}; +const WarnOnFailure pw_cuMemAddressFree{"cuMemAddressFree"}; const WarnOnFailure pw_cuGraphDestroy{"cuGraphDestroy"}; const WarnOnFailure pw_cuGraphExecDestroy{"cuGraphExecDestroy"}; const WarnOnFailure pw_cuGraphicsUnregisterResource{"cuGraphicsUnregisterResource"}; diff --git a/cuda_core/cuda/core/_cpp/rt/memory.cpp b/cuda_core/cuda/core/_cpp/rt/memory.cpp index 68a22c835bc..eb1f36666ef 100644 --- a/cuda_core/cuda/core/_cpp/rt/memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/memory.cpp @@ -8,6 +8,7 @@ #include "driver_api.hpp" #include "error.hpp" #include "internal.hpp" +#include "vmm_range.hpp" #include #include #include @@ -300,6 +301,31 @@ DevicePtrHandle deviceptr_create_mapped_graphics( return DevicePtrHandle(box, &box->resource); } +// ============================================================================ +// Virtual memory ranges (VMM_DESIGN.md) +// ============================================================================ + +DevicePtrHandle deviceptr_create_vmm(CUdeviceptr base, const VmmRangeHandle& range) { + if (!range) { + err = CUDA_ERROR_INVALID_VALUE; + return {}; + } + auto box = std::shared_ptr( + new DevicePtrBox{base, DeallocationStream{}}, + // Init-capture: a plain copy of the `const&` parameter would be const. + [range = range](DevicePtrBox* b) mutable { + GILReleaseGuard gil; + // Hand the recorded stream to the range, then drop the range: if + // this was the last owner, the range deleter synchronizes every + // recorded stream and unmaps. The box never blocks itself. + vmm_range_forward_stream(*range, b->deallocation); + delete b; + range.reset(); + } + ); + return DevicePtrHandle(box, &box->resource); +} + // ============================================================================ // MemoryResource-owned Device Pointer Handles // ============================================================================ diff --git a/cuda_core/cuda/core/_cpp/rt/py.hpp b/cuda_core/cuda/core/_cpp/rt/py.hpp index b606398317e..55212b5c42b 100644 --- a/cuda_core/cuda/core/_cpp/rt/py.hpp +++ b/cuda_core/cuda/core/_cpp/rt/py.hpp @@ -205,6 +205,18 @@ inline PyObject* as_py(const SurfObjectHandle& h) noexcept { return detail::make_py("cuda.bindings.driver", "CUsurfObject", as_intptr(h)); } +inline PyObject* as_py(const MemAllocationHandle& h) noexcept { + return detail::make_py("cuda.bindings.driver", "CUmemGenericAllocationHandle", as_intptr(h)); +} + +inline PyObject* as_py(const VaReservationHandle& h) noexcept { + return detail::make_py("cuda.bindings.driver", "CUdeviceptr", as_intptr(h)); +} + +inline PyObject* as_py(const VaMappingHandle& h) noexcept { + return detail::make_py("cuda.bindings.driver", "CUdeviceptr", as_intptr(h)); +} + // ============================================================================ // Python-coupled API: the prototypes that take or return PyObject* // ============================================================================ diff --git a/cuda_core/cuda/core/_cpp/rt/types.hpp b/cuda_core/cuda/core/_cpp/rt/types.hpp index 59389f78fca..af89a75e23f 100644 --- a/cuda_core/cuda/core/_cpp/rt/types.hpp +++ b/cuda_core/cuda/core/_cpp/rt/types.hpp @@ -43,6 +43,13 @@ using NvJitLinkValue = TaggedHandle; using TexObjectValue = TaggedHandle; using SurfObjectValue = TaggedHandle; +// Virtual memory management (VMM_DESIGN.md). CUmemGenericAllocationHandle is +// also `unsigned long long`, and the reservation and mapping handles carry a +// CUdeviceptr each, so all three are tagged for the same reason. +using MemAllocationValue = TaggedHandle; +using VaReservationValue = TaggedHandle; +using VaMappingValue = TaggedHandle; + // ============================================================================ // Handle type aliases - expose only the raw CUDA resource // ============================================================================ @@ -68,6 +75,15 @@ using MipmappedArrayHandle = std::shared_ptr; using TexObjectHandle = std::shared_ptr; using SurfObjectHandle = std::shared_ptr; +// Virtual memory management: a physical allocation (cuMemCreate), an address +// reservation (cuMemAddressReserve), one mapping of a whole allocation into a +// reservation (cuMemMap), and the range of mappings a buffer owns. +using MemAllocationHandle = std::shared_ptr; +using VaReservationHandle = std::shared_ptr; +using VaMappingHandle = std::shared_ptr; +struct VmmRange; +using VmmRangeHandle = std::shared_ptr; + using DevicePtrHandle = std::shared_ptr; // Type-erased shared owner of an attached resource. Typed handles such as @@ -201,6 +217,18 @@ inline CUsurfObject as_cu(const SurfObjectHandle& h) noexcept { return h ? h->raw : 0; } +inline CUmemGenericAllocationHandle as_cu(const MemAllocationHandle& h) noexcept { + return h ? h->raw : 0; +} + +inline CUdeviceptr as_cu(const VaReservationHandle& h) noexcept { + return h ? h->raw : 0; +} + +inline CUdeviceptr as_cu(const VaMappingHandle& h) noexcept { + return h ? h->raw : 0; +} + // as_intptr() - extract handle as intptr_t for Python interop // Using signed intptr_t per C standard convention and issue #1342 inline std::intptr_t as_intptr(const ContextHandle& h) noexcept { @@ -291,4 +319,16 @@ inline std::intptr_t as_intptr(const SurfObjectHandle& h) noexcept { return static_cast(as_cu(h)); } +inline std::intptr_t as_intptr(const MemAllocationHandle& h) noexcept { + return static_cast(as_cu(h)); +} + +inline std::intptr_t as_intptr(const VaReservationHandle& h) noexcept { + return static_cast(as_cu(h)); +} + +inline std::intptr_t as_intptr(const VaMappingHandle& h) noexcept { + return static_cast(as_cu(h)); +} + } // namespace cuda_core::rt diff --git a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp new file mode 100644 index 00000000000..76be333cabd --- /dev/null +++ b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp @@ -0,0 +1,316 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +// Virtual memory management handles: physical allocations, address +// reservations, mappings, and the range of mappings a buffer owns. +// See VMM_DESIGN.md for the ownership model. + +#include "py.hpp" +#include "api.hpp" +#include "context_scope.hpp" +#include "driver_api.hpp" +#include "error.hpp" +#include "internal.hpp" +#include "vmm_range.hpp" +#include +#include +#include +#include +#include + +namespace cuda_core::rt { + +using namespace detail; + +// ============================================================================ +// Boxes +// ============================================================================ + +namespace { + +struct MemAllocationBox { + MemAllocationValue resource; // from cuMemCreate + size_t size; // the only size cuMemMap accepts for it + std::vector access; // applied to every mapping of this allocation +}; + +struct VaReservationBox { + VaReservationValue resource; // base address from cuMemAddressReserve + size_t size; // exact reserved size; cuMemAddressFree needs both +}; + +struct VaMappingBox { + VaMappingValue resource; // mapped address + size_t size; // == the allocation's size + MemAllocationHandle h_alloc; // released after the unmap + VaReservationHandle h_reservation; // freed after the unmap +}; + +// Recover a box from the aliased handle; the resource is the first member. +template +const Box* box_of(const Handle& h) noexcept { + return reinterpret_cast( + reinterpret_cast(h.get()) - offsetof(Box, resource)); +} + +} // namespace + +// ============================================================================ +// Physical allocations +// ============================================================================ + +MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocationProp& prop, + const CUmemAccessDesc* descs, size_t count) { + // Copy the descriptors before the driver call so a failed copy leaves + // nothing to undo. + std::vector access(descs, descs + count); + GILReleaseGuard gil; + CUmemGenericAllocationHandle handle = 0; + if (CUDA_SUCCESS != (err = p_cuMemCreate(&handle, size, &prop, 0))) { + return {}; + } + auto box = std::shared_ptr( + new MemAllocationBox{{handle}, size, std::move(access)}, + [](const MemAllocationBox* b) { + GILReleaseGuard gil; + pw_cuMemRelease(b->resource.raw); + delete b; + } + ); + return MemAllocationHandle(box, &box->resource); +} + +size_t mem_allocation_size(const MemAllocationHandle& h) noexcept { + return h ? box_of(h)->size : 0; +} + +// ============================================================================ +// Address reservations +// ============================================================================ + +VaReservationHandle create_va_reservation_handle(size_t size, size_t alignment, CUdeviceptr hint) { + GILReleaseGuard gil; + CUdeviceptr ptr = 0; + if (CUDA_SUCCESS != (err = p_cuMemAddressReserve(&ptr, size, alignment, hint, 0))) { + return {}; + } + auto box = std::shared_ptr( + new VaReservationBox{{ptr}, size}, + [](const VaReservationBox* b) { + GILReleaseGuard gil; + pw_cuMemAddressFree(b->resource.raw, b->size); + delete b; + } + ); + return VaReservationHandle(box, &box->resource); +} + +size_t va_reservation_size(const VaReservationHandle& h) noexcept { + return h ? box_of(h)->size : 0; +} + +// ============================================================================ +// Mappings +// ============================================================================ + +VaMappingHandle create_va_mapping_handle(CUdeviceptr ptr, const MemAllocationHandle& h_alloc, + const VaReservationHandle& h_res) { + if (!h_alloc || !h_res) { + err = CUDA_ERROR_INVALID_VALUE; + return {}; + } + const MemAllocationBox* alloc = box_of(h_alloc); + const VaReservationBox* res = box_of(h_res); + const CUdeviceptr base = res->resource.raw; + if (ptr < base || ptr - base > res->size || alloc->size > res->size - (ptr - base)) { + err = CUDA_ERROR_INVALID_VALUE; + return {}; + } + + GILReleaseGuard gil; + if (CUDA_SUCCESS != (err = p_cuMemMap(ptr, alloc->size, 0, alloc->resource.raw, 0))) { + return {}; + } + // cuMemSetAccess rejects an empty descriptor list; a mapping with no + // descriptors is mapped but not accessible, which is what the caller asked for. + if (!alloc->access.empty()) { + const CUresult status = p_cuMemSetAccess(ptr, alloc->size, alloc->access.data(), alloc->access.size()); + if (status != CUDA_SUCCESS) { + pw_cuMemUnmap(ptr, alloc->size); + err = status; + return {}; + } + } + auto box = std::shared_ptr( + new VaMappingBox{{ptr}, alloc->size, h_alloc, h_res}, + [](const VaMappingBox* b) { + GILReleaseGuard gil; + pw_cuMemUnmap(b->resource.raw, b->size); + delete b; // then the allocation and the reservation release + } + ); + return VaMappingHandle(box, &box->resource); +} + +size_t va_mapping_size(const VaMappingHandle& h) noexcept { + return h ? box_of(h)->size : 0; +} + +MemAllocationHandle va_mapping_allocation(const VaMappingHandle& h) noexcept { + return h ? box_of(h)->h_alloc : MemAllocationHandle{}; +} + +// ============================================================================ +// Ranges +// ============================================================================ + +// Base address -> live range, so a DevicePtrHandle can be recognized as a VMM +// buffer and its range recovered. Two buffers that share a base share the +// range. The range deleter removes the entry before it frees the reservation, +// so the address cannot be re-reserved while the key is present. +static HandleRegistry vmm_range_registry; + +namespace detail { +void vmm_range_forward_stream(VmmRange& range, const DeallocationStream& stream) noexcept { + if (!stream.h_stream) { + return; + } + const CUstream s = as_cu(stream.h_stream); + const CUcontext ctx = as_cu(get_stream_context(stream.h_stream)); + std::lock_guard lock(range.mu); + for (const DeallocationStream& recorded : range.streams) { + if (as_cu(recorded.h_stream) == s && as_cu(get_stream_context(recorded.h_stream)) == ctx) { + return; + } + } + try { + range.streams.push_back(stream); + } catch (...) { + range.stream_dropped = true; // reported by the range deleter, outside the lock + } +} +} // namespace detail + +// True when synchronizing `stream` would disturb a graph capture: the stream +// is capturing, or it is the legacy stream while a blocking stream in its +// context is capturing (the query reports that as +// CUDA_ERROR_STREAM_CAPTURE_IMPLICIT). cuStreamSynchronize would invalidate +// such a capture; cuStreamGetCaptureInfo does not. +static bool sync_would_disturb_capture(CUstream stream) noexcept { + CUstreamCaptureStatus status = CU_STREAM_CAPTURE_STATUS_NONE; +#if CUDA_VERSION >= 13000 + const CUresult result = p_cuStreamGetCaptureInfo(stream, &status, nullptr, nullptr, nullptr, nullptr, nullptr); +#else + const CUresult result = p_cuStreamGetCaptureInfo(stream, &status, nullptr, nullptr, nullptr, nullptr); +#endif + if (result == CUDA_ERROR_STREAM_CAPTURE_IMPLICIT) { + return true; + } + return result == CUDA_SUCCESS && status == CU_STREAM_CAPTURE_STATUS_ACTIVE; +} + +// Synchronize a recorded deallocation stream with its bound context current, +// then restore the caller's context. Modeled on cleanup_in_context, with a +// skip message that fits this use: when the sync cannot run, nothing leaks, +// because the range is unmapped regardless. Sets `capture_skipped` instead of +// synchronizing when the sync would disturb a capture. +static void sync_recorded_stream(const DeallocationStream& ds, bool& capture_skipped) noexcept { + const CUstream s = as_cu(ds.h_stream); + CUcontext previous = nullptr; + int changed = 0; + const char* detail = nullptr; + CUresult status = enter_context(deallocation_context(ds), &previous, &changed); + if (status != CUDA_SUCCESS) { + detail = "skipped (context activation failed); the range was unmapped without synchronizing it"; + } else if (sync_would_disturb_capture(s)) { + capture_skipped = true; + } else { + status = p_cuStreamSynchronize(s); + } + const CUresult restore = exit_context(previous, changed, CUDA_SUCCESS); + if (restore != CUDA_SUCCESS) { + // Nothing is raised here, so the detail exit_context recorded has no + // exception to attach to; drop it. + clear_last_error_detail(); + } + if (status != CUDA_SUCCESS || restore != CUDA_SUCCESS) { + char operation_name[160]; + format_operation(operation_name, sizeof(operation_name), "cuStreamSynchronize", handle_bits(s)); + if (status != CUDA_SUCCESS) { + report_cuda_error(operation_name, status, detail); + } + if (restore != CUDA_SUCCESS) { + report_cuda_error(operation_name, restore, "failed while restoring the caller's context"); + } + } +} + +VmmRangeHandle create_vmm_range(CUdeviceptr base) { + auto range = VmmRangeHandle( + new VmmRange(base), + [](VmmRange* r) { + GILReleaseGuard gil; + vmm_range_registry.unregister_handle(r->base); + if (!py_is_finalizing()) { + bool capture_skipped = false; + for (const DeallocationStream& ds : r->streams) { + sync_recorded_stream(ds, capture_skipped); + } + if (capture_skipped) { + report_message( + "a VirtualMemoryResource buffer was released while its deallocation stream " + "is capturing, or is the legacy stream while another stream in its context is " + "capturing; the range was unmapped without synchronizing that stream"); + } + if (r->stream_dropped) { + report_message( + "a VirtualMemoryResource buffer could not record a deallocation stream " + "(out of memory); the range was unmapped without synchronizing it"); + } + } + delete r; // mappings unmap; reservations free and allocations release + } + ); + vmm_range_registry.register_handle(base, range); + return range; +} + +VmmRangeHandle vmm_range(const DevicePtrHandle& h) { + return h ? vmm_range_registry.lookup(as_cu(h)) : VmmRangeHandle{}; +} + +size_t vmm_range_count(const VmmRangeHandle& range) noexcept { + return range ? range->mappings.size() : 0; +} + +VaMappingHandle vmm_range_mapping(const VmmRangeHandle& range, size_t index) noexcept { + if (!range || index >= range->mappings.size()) { + return {}; + } + return range->mappings[index]; +} + +size_t vmm_range_total(const VmmRangeHandle& range) noexcept { + size_t total = 0; + if (range) { + for (const VaMappingHandle& m : range->mappings) { + total += va_mapping_size(m); + } + } + return total; +} + +void vmm_range_reserve(const VmmRangeHandle& range, size_t count) { + if (range) { + range->mappings.reserve(count); + } +} + +void vmm_range_append(const VmmRangeHandle& range, const VaMappingHandle& mapping) { + if (range && mapping) { + range->mappings.push_back(mapping); + } +} + +} // namespace cuda_core::rt diff --git a/cuda_core/cuda/core/_cpp/rt/vmm_range.hpp b/cuda_core/cuda/core/_cpp/rt/vmm_range.hpp new file mode 100644 index 00000000000..8c699d4397f --- /dev/null +++ b/cuda_core/cuda/core/_cpp/rt/vmm_range.hpp @@ -0,0 +1,38 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "types.hpp" +#include "internal.hpp" +#include +#include + +namespace cuda_core::rt { + +// The mappings a VirtualMemoryResource buffer owns, in ascending address +// order, plus every deallocation stream an owner of the range recorded. +// Owned through VmmRangeHandle by the DevicePtrBox of each buffer that maps +// the range (a grow in place creates a second owner at the same base). The +// range deleter synchronizes the recorded streams and then destroys the +// mappings, which unmap, free the reservations and release the allocations +// as their last references go. See VMM_DESIGN.md. +struct VmmRange { + explicit VmmRange(CUdeviceptr base_) noexcept : base(base_) {} + + CUdeviceptr base; // registry key + std::vector mappings; // contiguous; sum of sizes = range total + std::mutex mu; // guards `streams` and `stream_dropped` + std::vector streams; // forwarded by dying owners, deduplicated + bool stream_dropped = false; // a forward failed for lack of memory +}; + +namespace detail { +// Record the stream an owner of `range` used for deallocation, so the range +// deleter synchronizes it before it unmaps. Takes range.mu; nothing under the +// lock acquires the GIL. Implemented in virtual_memory.cpp. +void vmm_range_forward_stream(VmmRange& range, const DeallocationStream& stream) noexcept; +} // namespace detail + +} // namespace cuda_core::rt diff --git a/cuda_core/cuda/core/_memory/_buffer.pxd b/cuda_core/cuda/core/_memory/_buffer.pxd index e929210601f..e74a4833947 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pxd +++ b/cuda_core/cuda/core/_memory/_buffer.pxd @@ -24,10 +24,7 @@ cdef class Buffer: _MemAttrs _mem_attrs std_atomic[cpp_bool] _mem_attrs_inited object __weakref__ - cdef public: - # Python code in _memory/_virtual_memory_resource.py needs to update - # this value, though it is technically private. - size_t _size + size_t _size cdef class MemoryResource: diff --git a/cuda_core/cuda/core/_memory/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index 756b8661488..ce31b9129f3 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -29,8 +29,6 @@ class Buffer: by calling :meth:`from_ipc_descriptor` and therefore performs an IPC import. Do not unpickle buffers from untrusted sources. """ - _size: int - def _clear(self) -> None: ... def __init__(self, *args, **kwargs) -> None: ... @classmethod diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py deleted file mode 100644 index ea1e2455c6f..00000000000 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ /dev/null @@ -1,636 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Iterable - -if TYPE_CHECKING: - from cuda.core._stream import Stream - from cuda.core.graph import GraphBuilder - -from cuda.core._device import Device -from cuda.core._memory._buffer import Buffer, MemoryResource -from cuda.core._utils.cuda_utils import ( - Transaction, - check_or_create_options, - driver, -) -from cuda.core._utils.cuda_utils import ( - _check_driver_error as raise_if_driver_error, -) -from cuda.core._utils.version import binding_version -from cuda.core.typing import ( - DevicePointerType, - VirtualMemoryAccessType, - VirtualMemoryAllocationType, - VirtualMemoryGranularityType, - VirtualMemoryHandleType, - VirtualMemoryLocationType, -) - -__all__ = ["VirtualMemoryResource", "VirtualMemoryResourceOptions"] - -# Location types whose physical backing lives in host memory. Shared by -# VirtualMemoryResource.__init__ and is_host_accessible so the two cannot drift. -_HOST_LOCATION_TYPES = frozenset( - { - VirtualMemoryLocationType.HOST, - VirtualMemoryLocationType.HOST_NUMA, - VirtualMemoryLocationType.HOST_NUMA_CURRENT, - } -) - - -@dataclass -class VirtualMemoryResourceOptions: - """A configuration object for the VirtualMemoryResource - Stores configuration information which tells the resource how to use the CUDA VMM APIs - - Attributes - ---------- - allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str - Controls the type of allocation. - location_type: :obj:`~_memory.VirtualMemoryLocationType` | str - Controls the location of the allocation. - handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str - Export handle type for the physical allocation. Use ``"posix_fd"`` on - Linux if you plan to import/export the allocation. Use `None` if you - don't need an exportable handle. - gpu_direct_rdma: bool - Hint that the allocation should be GDR-capable (if supported). - granularity: :obj:`~_memory.VirtualMemoryGranularityType` | str - Controls granularity query and size rounding. - addr_hint: int - A (optional) virtual address hint to try to reserve at. Setting it to 0 lets the CUDA driver decide. - addr_align: int - Alignment for the VA reservation. If `None`, use the queried granularity. - peers: Iterable[int] - Extra device IDs that should be granted access in addition to ``device``. - self_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str - Access flags for the owning device. - peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str - Access flags for peers. - """ - - allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED - location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE - handle_type: VirtualMemoryHandleType = VirtualMemoryHandleType.POSIX_FD - granularity: VirtualMemoryGranularityType = VirtualMemoryGranularityType.RECOMMENDED - gpu_direct_rdma: bool = False - addr_hint: int | None = 0 - addr_align: int | None = None - peers: Iterable[int] = field(default_factory=tuple) - self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE - peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE - - _a = driver.CUmemAccess_flags - _access_flags = { # noqa: RUF012 - VirtualMemoryAccessType.READ_WRITE: _a.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, - VirtualMemoryAccessType.READ: _a.CU_MEM_ACCESS_FLAGS_PROT_READ, - None: 0, - } - _h = driver.CUmemAllocationHandleType - _handle_types = { # noqa: RUF012 - None: _h.CU_MEM_HANDLE_TYPE_NONE, - VirtualMemoryHandleType.POSIX_FD: _h.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, - VirtualMemoryHandleType.WIN32_KMT: _h.CU_MEM_HANDLE_TYPE_WIN32_KMT, - VirtualMemoryHandleType.FABRIC: _h.CU_MEM_HANDLE_TYPE_FABRIC, - } - _g = driver.CUmemAllocationGranularity_flags - _granularity = { # noqa: RUF012 - VirtualMemoryGranularityType.RECOMMENDED: _g.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, - VirtualMemoryGranularityType.MINIMUM: _g.CU_MEM_ALLOC_GRANULARITY_MINIMUM, - } - _l = driver.CUmemLocationType - _location_type = { # noqa: RUF012 - VirtualMemoryLocationType.DEVICE: _l.CU_MEM_LOCATION_TYPE_DEVICE, - VirtualMemoryLocationType.HOST: _l.CU_MEM_LOCATION_TYPE_HOST, - VirtualMemoryLocationType.HOST_NUMA: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA, - VirtualMemoryLocationType.HOST_NUMA_CURRENT: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, - } - _t = driver.CUmemAllocationType - # CUDA 13+ exposes MANAGED in CUmemAllocationType; older 12.x does not - _allocation_type = {VirtualMemoryAllocationType.PINNED: _t.CU_MEM_ALLOCATION_TYPE_PINNED} # noqa: RUF012 - if binding_version() >= (13, 0, 0): - _allocation_type[VirtualMemoryAllocationType.MANAGED] = _t.CU_MEM_ALLOCATION_TYPE_MANAGED - - @staticmethod - def _access_to_flags(spec: VirtualMemoryAccessType | None) -> int: - flags = VirtualMemoryResourceOptions._access_flags.get(spec) - if flags is None: - raise ValueError(f"Unknown access spec: {spec!r}") - return flags # type: ignore[no-any-return] - - @staticmethod - def _allocation_type_to_driver(spec: VirtualMemoryAllocationType) -> int: - alloc_type = VirtualMemoryResourceOptions._allocation_type.get(spec) - if alloc_type is None: - raise ValueError(f"Unsupported allocation_type: {spec!r}") - return alloc_type # type: ignore[no-any-return] - - @staticmethod - def _location_type_to_driver(spec: VirtualMemoryLocationType) -> int: - loc_type = VirtualMemoryResourceOptions._location_type.get(spec) - if loc_type is None: - raise ValueError(f"Unsupported location_type: {spec!r}") - return loc_type # type: ignore[no-any-return] - - @staticmethod - def _handle_type_to_driver(spec: VirtualMemoryHandleType | None) -> int: - if spec == "win32": - raise NotImplementedError("win32 is currently not supported, please reach out to the CUDA Python team") - handle_type = VirtualMemoryResourceOptions._handle_types.get(spec) - if handle_type is None: - raise ValueError(f"Unsupported handle_type: {spec!r}") - return handle_type # type: ignore[no-any-return] - - @staticmethod - def _granularity_to_driver(spec: VirtualMemoryGranularityType) -> int: - granularity = VirtualMemoryResourceOptions._granularity.get(spec) - if granularity is None: - raise ValueError(f"Unsupported granularity: {spec!r}") - return granularity # type: ignore[no-any-return] - - -class VirtualMemoryResource(MemoryResource): - """Create a device memory resource that uses the CUDA VMM APIs to allocate memory. - - Parameters - ---------- - device_id : Device | int - Device for which a memory resource is constructed. - - config : VirtualMemoryResourceOptions, optional - A configuration object for the VirtualMemoryResource - - - Warning - ------- - This is a low-level API that is provided only for convenience. Make sure you fully understand - how CUDA Virtual Memory Management works before using this. Other MemoryResource subclasses - in cuda.core should already meet the common needs. - """ - - def __init__(self, device_id: Device | int, config: VirtualMemoryResourceOptions | None = None) -> None: - self.device: Device | None = Device(device_id) - self.config: VirtualMemoryResourceOptions = check_or_create_options( # type: ignore[assignment] - VirtualMemoryResourceOptions, config, "VirtualMemoryResource options", keep_none=False - ) - if self.config.location_type in _HOST_LOCATION_TYPES: - self.device = None - - if not self.device and self.config.location_type == "device": - raise RuntimeError("VirtualMemoryResource requires a device for device memory allocations") - - if self.device and not self.device.properties.virtual_memory_management_supported: - raise RuntimeError("VirtualMemoryResource requires CUDA VMM API support") - - # Validate RDMA support if requested - if ( - self.config.gpu_direct_rdma - and self.device is not None - and not self.device.properties.gpu_direct_rdma_supported - ): - raise RuntimeError("GPU Direct RDMA is not supported on this device") - - @staticmethod - def _align_up(size: int, gran: int) -> int: - """ - Align a size up to the nearest multiple of a granularity. - """ - return (size + gran - 1) & ~(gran - 1) - - def modify_allocation( - self, buf: Buffer, new_size: int, config: VirtualMemoryResourceOptions | None = None - ) -> Buffer: - """ - Grow an existing allocation using CUDA VMM, with a configurable policy. - - This implements true growing allocations that preserve the base pointer - by extending the virtual address range and mapping additional physical memory. - - This function uses transactional allocation: if any step fails, the original buffer is not modified and - all steps the function took are rolled back so a new allocation is not created. - - Parameters - ---------- - buf : Buffer - The existing buffer to grow - new_size : int - The new total size for the allocation - config : VirtualMemoryResourceOptions, optional - Configuration for the new physical memory chunks. If None, uses current config. - - Returns - ------- - Buffer - The same buffer with updated size and properties, preserving the original pointer - """ - if not isinstance(buf, Buffer): - raise TypeError(f"buf must be a Buffer, got {type(buf).__name__}") - if buf.is_closed: - raise RuntimeError("Buffer has been closed") - if config is not None: - self.config = config - - # Build allocation properties for new chunks - prop = driver.CUmemAllocationProp() - prop.type = VirtualMemoryResourceOptions._allocation_type_to_driver(self.config.allocation_type) - prop.location.type = VirtualMemoryResourceOptions._location_type_to_driver(self.config.location_type) - # Caller must not invoke modify_allocation on a host-located resource; - # we rely on the dataclass invariant that self.device is non-None for - # device-located resources (it's only None when location is host). - assert self.device is not None, "modify_allocation requires a device-located resource" - prop.location.id = self.device.device_id - prop.allocFlags.gpuDirectRDMACapable = 1 if self.config.gpu_direct_rdma else 0 - prop.requestedHandleTypes = VirtualMemoryResourceOptions._handle_type_to_driver(self.config.handle_type) - prop.win32HandleMetaData = 0 - - # Query granularity - gran_flag = VirtualMemoryResourceOptions._granularity_to_driver(self.config.granularity) - res, gran = driver.cuMemGetAllocationGranularity(prop, gran_flag) - raise_if_driver_error(res) - - # Calculate sizes - additional_size = new_size - buf.size - if additional_size <= 0: - # Same size: only update access policy if needed; avoid zero-sized driver calls - descs = self._build_access_descriptors(prop) - if descs: - (res,) = driver.cuMemSetAccess(int(buf.handle), buf.size, descs, len(descs)) - raise_if_driver_error(res) - return buf - - aligned_additional_size = VirtualMemoryResource._align_up(additional_size, gran) - total_aligned_size = VirtualMemoryResource._align_up(new_size, gran) - aligned_prev_size = total_aligned_size - aligned_additional_size - addr_align = self.config.addr_align or gran - - # Try to extend the existing VA range first - res, new_ptr = driver.cuMemAddressReserve( - aligned_additional_size, - addr_align, - int(buf.handle) + aligned_prev_size, # fixedAddr hint - aligned end of current range - 0, - ) - - if res != driver.CUresult.CUDA_SUCCESS or new_ptr != (int(buf.handle) + aligned_prev_size): - # Check for specific errors that are not recoverable with the slow path - if res in ( - driver.CUresult.CUDA_ERROR_INVALID_VALUE, - driver.CUresult.CUDA_ERROR_NOT_PERMITTED, - driver.CUresult.CUDA_ERROR_NOT_INITIALIZED, - driver.CUresult.CUDA_ERROR_NOT_SUPPORTED, - ): - raise_if_driver_error(res) - (res2,) = driver.cuMemAddressFree(new_ptr, aligned_additional_size) - raise_if_driver_error(res2) - # Fallback: couldn't extend contiguously, need full remapping - return self._grow_allocation_slow_path( - buf, new_size, prop, aligned_additional_size, total_aligned_size, addr_align - ) - else: - # Success! We can extend the VA range contiguously - return self._grow_allocation_fast_path(buf, new_size, prop, aligned_additional_size, new_ptr) - - def _grow_allocation_fast_path( - self, buf: Buffer, new_size: int, prop: driver.CUmemAllocationProp, aligned_additional_size: int, new_ptr: int - ) -> Buffer: - """ - Fast path for growing a virtual memory allocation when the new region can be - reserved contiguously after the existing buffer. - - This function creates and maps new physical memory for the additional size, - sets access permissions, and updates the buffer size in place (the pointer - remains unchanged). - - Args: - buf (Buffer): - The buffer to grow. - - new_size (int): - The new total size in bytes. - - prop (driver.CUmemAllocationProp): - Allocation properties for the new memory. - - aligned_additional_size (int): - The size of the new region to allocate, aligned to granularity. - - new_ptr (int): - The address of the newly reserved contiguous VA region (should - be at the end of the current buffer). - - Returns: - Buffer: The same buffer object with its size updated to `new_size`. - """ - with Transaction() as trans: - # Create new physical memory for the additional size - trans.on_failure( - lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0]) - ) - res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) - raise_if_driver_error(res) - trans.on_exit(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - - # Map the new physical memory to the extended VA range - (res,) = driver.cuMemMap(new_ptr, aligned_additional_size, 0, new_handle, 0) - raise_if_driver_error(res) - # Register undo for mapping - trans.on_failure( - lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0]) - ) - - # Set access permissions for the new portion - descs = self._build_access_descriptors(prop) - if descs: - (res,) = driver.cuMemSetAccess(new_ptr, aligned_additional_size, descs, len(descs)) - raise_if_driver_error(res) - - # All succeeded, cancel undo actions - trans.commit() - - # Update the buffer size (pointer stays the same). `Buffer.size` has - # no public setter, so this reaches into the private attribute. - buf._size = new_size - return buf - - def _grow_allocation_slow_path( - self, - buf: Buffer, - new_size: int, - prop: driver.CUmemAllocationProp, - aligned_additional_size: int, - total_aligned_size: int, - addr_align: int, - ) -> Buffer: - """ - Slow path for growing a virtual memory allocation when the new region cannot be - reserved contiguously after the existing buffer. - - This function reserves a new, larger virtual address (VA) range, remaps the old - physical memory to the beginning of the new VA range, creates and maps new physical - memory for the additional size, sets access permissions, and updates the buffer's - pointer and size. - - Args: - buf (Buffer): The buffer to grow. - new_size (int): The new total size in bytes. - prop (driver.CUmemAllocationProp): Allocation properties for the new memory. - aligned_additional_size (int): The size of the new region to allocate, aligned to granularity. - total_aligned_size (int): The total new size to reserve, aligned to granularity. - addr_align (int): The required address alignment for the new VA range. - - Returns: - Buffer: The buffer object updated with the new pointer and size. - """ - with Transaction() as trans: - # Reserve a completely new, larger VA range - res, new_ptr = driver.cuMemAddressReserve(total_aligned_size, addr_align, 0, 0) - raise_if_driver_error(res) - # Register undo for VA reservation - trans.on_failure( - lambda np=new_ptr, s=total_aligned_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0]) - ) - - # Get the old allocation handle for remapping - result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle) - raise_if_driver_error(result) - trans.on_exit(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - - # Unmap the old VA range (aligned previous size) - aligned_prev_size = total_aligned_size - aligned_additional_size - (result,) = driver.cuMemUnmap(int(buf.handle), aligned_prev_size) - raise_if_driver_error(result) - - def _remap_old() -> None: - # Try to remap the old physical memory back to the original VA range - try: - (res,) = driver.cuMemMap(int(buf.handle), aligned_prev_size, 0, old_handle, 0) - raise_if_driver_error(res) - except Exception: # noqa: S110 - # TODO: consider logging this exception - pass - - trans.on_failure(_remap_old) - - # Remap the old physical memory to the new VA range (aligned previous size) - (res,) = driver.cuMemMap(int(new_ptr), aligned_prev_size, 0, old_handle, 0) - raise_if_driver_error(res) - - # Register undo for mapping - trans.on_failure(lambda np=new_ptr, s=aligned_prev_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0])) - - # Create new physical memory for the additional size - res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) - raise_if_driver_error(res) - trans.on_exit(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - - # Map the new physical memory to the extended portion (aligned offset) - (res,) = driver.cuMemMap(int(new_ptr) + aligned_prev_size, aligned_additional_size, 0, new_handle, 0) - raise_if_driver_error(res) - - # Register undo for mapping - trans.on_failure( - lambda base=int(new_ptr), offs=aligned_prev_size, s=aligned_additional_size: raise_if_driver_error( - driver.cuMemUnmap(base + offs, s)[0] - ) - ) - - # Set access permissions for the entire new range - descs = self._build_access_descriptors(prop) - if descs: - (res,) = driver.cuMemSetAccess(new_ptr, total_aligned_size, descs, len(descs)) - raise_if_driver_error(res) - - # All succeeded, cancel undo actions - trans.commit() - - # Free the old VA range (aligned previous size) - (res2,) = driver.cuMemAddressFree(int(buf.handle), aligned_prev_size) - raise_if_driver_error(res2) - - # Invalidate the old buffer so its destructor won't try to free again - buf._clear() - - # Return a new Buffer for the new mapping - return Buffer.from_handle(ptr=new_ptr, size=new_size, mr=self) - - def _build_access_descriptors(self, prop: driver.CUmemAllocationProp) -> list[driver.CUmemAccessDesc]: - """ - Build access descriptors for memory access permissions. - - Returns - ------- - list - List of CUmemAccessDesc objects for setting memory access - """ - descs = [] - - # Owner access - owner_flags = VirtualMemoryResourceOptions._access_to_flags(self.config.self_access) - if owner_flags: - d = driver.CUmemAccessDesc() - d.location.type = prop.location.type - d.location.id = prop.location.id - d.flags = owner_flags - descs.append(d) - - # Peer device access - peer_flags = VirtualMemoryResourceOptions._access_to_flags(self.config.peer_access) - if peer_flags: - for peer_dev in self.config.peers: - d = driver.CUmemAccessDesc() - d.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - d.location.id = int(peer_dev) - d.flags = peer_flags - descs.append(d) - - return descs - - def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> Buffer: - """ - Allocate a buffer of the given size using CUDA virtual memory. - - Parameters - ---------- - size : int - The size in bytes of the buffer to allocate. - stream : Stream, optional - Keyword-only. Unused because virtual memory operations are - synchronous. - - Returns - ------- - Buffer - A Buffer object representing the allocated virtual memory. - - Raises - ------ - CUDAError - If any CUDA driver API call fails during allocation. - - Notes - ----- - This method uses transactional allocation: if any step fails, all resources - allocated so far are automatically cleaned up. The allocation is performed - with the configured granularity, access permissions, and peer access as - specified in the resource's configuration. - """ - if stream is not None: - from cuda.core._stream import Stream_accept - - Stream_accept(stream) - - config = self.config - # ---- Build allocation properties ---- - prop = driver.CUmemAllocationProp() - prop.type = VirtualMemoryResourceOptions._allocation_type_to_driver(config.allocation_type) - prop.location.type = VirtualMemoryResourceOptions._location_type_to_driver(config.location_type) - prop.location.id = self.device.device_id if self.device is not None else -1 - prop.allocFlags.gpuDirectRDMACapable = 1 if config.gpu_direct_rdma else 0 - prop.requestedHandleTypes = VirtualMemoryResourceOptions._handle_type_to_driver(config.handle_type) - prop.win32HandleMetaData = 0 - - # ---- Query and apply granularity ---- - # Choose min vs recommended granularity per config - gran_flag = VirtualMemoryResourceOptions._granularity_to_driver(config.granularity) - res, gran = driver.cuMemGetAllocationGranularity(prop, gran_flag) - raise_if_driver_error(res) - - aligned_size = VirtualMemoryResource._align_up(size, gran) - addr_align = config.addr_align or gran - - # ---- Transactional allocation ---- - with Transaction() as trans: - # ---- Create physical memory ---- - res, handle = driver.cuMemCreate(aligned_size, prop, 0) - raise_if_driver_error(res) - # Drop the creation reference on either outcome; a successful mapping keeps the allocation alive. - trans.on_exit(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - - # ---- Reserve VA space ---- - # Potentially, use a separate size for the VA reservation from the physical allocation size - res, ptr = driver.cuMemAddressReserve(aligned_size, addr_align, config.addr_hint, 0) - raise_if_driver_error(res) - # Register undo for VA reservation - trans.on_failure(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemAddressFree(p, s)[0])) - - # ---- Map physical memory into VA ---- - (res,) = driver.cuMemMap(ptr, aligned_size, 0, handle, 0) - raise_if_driver_error(res) - trans.on_failure(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0])) - - # ---- Set access for owner + peers ---- - descs = self._build_access_descriptors(prop) - if descs: - (res,) = driver.cuMemSetAccess(ptr, aligned_size, descs, len(descs)) - raise_if_driver_error(res) - - trans.commit() - - # Done — return a Buffer that tracks this VA range - buf = Buffer.from_handle(ptr=ptr, size=aligned_size, mr=self) - return buf - - def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | GraphBuilder | None = None) -> None: - """ - Deallocate memory on the device using CUDA VMM APIs. - - Parameters - ---------- - ptr : DevicePointerType - The pointer to the memory to deallocate. - size : int - The size in bytes of the memory to deallocate. - stream : Stream, optional - Keyword-only. Unused because virtual memory operations are - synchronous. - """ - ptr = 0 if ptr is None else int(ptr) - - if stream is not None: - from cuda.core._stream import Stream_accept - - Stream_accept(stream) - # The mapping owns the allocation; unmapping frees its backing memory when no external references remain. - (result,) = driver.cuMemUnmap(ptr, size) - raise_if_driver_error(result) - (result,) = driver.cuMemAddressFree(ptr, size) - raise_if_driver_error(result) - - @property - def is_device_accessible(self) -> bool: - """ - Indicates whether the allocated memory is accessible from the device. - """ - return self.config.location_type == "device" - - @property - def is_host_accessible(self) -> bool: - """ - Indicates whether the allocated memory is accessible from the host. - """ - return self.config.location_type in _HOST_LOCATION_TYPES - - @property - def device_id(self) -> int: - """ - Get the device ID associated with this memory resource. - - Returns: - int: CUDA device ID. -1 if the memory resource allocates host memory - """ - return self.device.device_id if self.device is not None else -1 - - def __repr__(self) -> str: - """ - Return a string representation of the VirtualMemoryResource. - - Returns: - str: A string describing the object - """ - return f"" diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi new file mode 100644 index 00000000000..992771794cd --- /dev/null +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -0,0 +1,250 @@ +# This file was generated by stubgen-pyx v0.2.22 from cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx + +from dataclasses import dataclass, field +from typing import Iterable + +from cuda.core._device import Device +from cuda.core._memory._buffer import Buffer, MemoryResource +from cuda.core._stream import Stream +from cuda.core._utils.cuda_utils import driver +from cuda.core.graph import GraphBuilder +from cuda.core.typing import (DevicePointerType, VirtualMemoryAccessType, + VirtualMemoryAllocationType, + VirtualMemoryGranularityType, + VirtualMemoryHandleType, + VirtualMemoryLocationType) + +__all__ = ['VirtualMemoryBuffer', 'VirtualMemoryResource', 'VirtualMemoryResourceOptions'] +_HOST_LOCATION_TYPES = frozenset({VirtualMemoryLocationType.HOST, VirtualMemoryLocationType.HOST_NUMA, VirtualMemoryLocationType.HOST_NUMA_CURRENT}) + +@dataclass +class VirtualMemoryResourceOptions: + """A configuration object for the VirtualMemoryResource + Stores configuration information which tells the resource how to use the CUDA VMM APIs + + Attributes + ---------- + allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str + Controls the type of allocation. + location_type: :obj:`~_memory.VirtualMemoryLocationType` | str + Controls the location of the allocation. + handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str + Export handle type for the physical allocation. Use ``"posix_fd"`` on + Linux if you plan to import/export the allocation. Use `None` if you + don't need an exportable handle. Host-located allocations require + `None`. + gpu_direct_rdma: bool + Hint that the allocation should be GDR-capable (if supported). + granularity: :obj:`~_memory.VirtualMemoryGranularityType` | str + Controls granularity query and size rounding. + addr_hint: int + A (optional) virtual address hint to try to reserve at. Setting it to 0 lets the CUDA driver decide. + addr_align: int + Alignment for the VA reservation. If `None`, use the queried granularity. + peers: Iterable[int] + Extra device IDs that should be granted access in addition to ``device``. + self_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str + Access flags for the owning device. + peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str + Access flags for peers. + """ + allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED + location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE + handle_type: VirtualMemoryHandleType = VirtualMemoryHandleType.POSIX_FD + granularity: VirtualMemoryGranularityType = VirtualMemoryGranularityType.RECOMMENDED + gpu_direct_rdma: bool = False + addr_hint: int | None = 0 + addr_align: int | None = None + peers: Iterable[int] = field(default_factory=tuple) + self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE + peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE + _a = driver.CUmemAccess_flags + _access_flags = {VirtualMemoryAccessType.READ_WRITE: _a.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, VirtualMemoryAccessType.READ: _a.CU_MEM_ACCESS_FLAGS_PROT_READ, None: 0} + _h = driver.CUmemAllocationHandleType + _handle_types = {None: _h.CU_MEM_HANDLE_TYPE_NONE, VirtualMemoryHandleType.POSIX_FD: _h.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, VirtualMemoryHandleType.WIN32_KMT: _h.CU_MEM_HANDLE_TYPE_WIN32_KMT, VirtualMemoryHandleType.FABRIC: _h.CU_MEM_HANDLE_TYPE_FABRIC} + _g = driver.CUmemAllocationGranularity_flags + _granularity = {VirtualMemoryGranularityType.RECOMMENDED: _g.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, VirtualMemoryGranularityType.MINIMUM: _g.CU_MEM_ALLOC_GRANULARITY_MINIMUM} + _l = driver.CUmemLocationType + _location_type = {VirtualMemoryLocationType.DEVICE: _l.CU_MEM_LOCATION_TYPE_DEVICE, VirtualMemoryLocationType.HOST: _l.CU_MEM_LOCATION_TYPE_HOST, VirtualMemoryLocationType.HOST_NUMA: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA, VirtualMemoryLocationType.HOST_NUMA_CURRENT: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT} + _t = driver.CUmemAllocationType + _allocation_type = {VirtualMemoryAllocationType.PINNED: _t.CU_MEM_ALLOCATION_TYPE_PINNED} + + @staticmethod + def _access_to_flags(spec: VirtualMemoryAccessType | None) -> int: ... + @staticmethod + def _allocation_type_to_driver(spec: VirtualMemoryAllocationType) -> int: ... + @staticmethod + def _location_type_to_driver(spec: VirtualMemoryLocationType) -> int: ... + @staticmethod + def _handle_type_to_driver(spec: VirtualMemoryHandleType | None) -> int: ... + @staticmethod + def _granularity_to_driver(spec: VirtualMemoryGranularityType) -> int: ... + +class VirtualMemoryBuffer(Buffer): + """A :class:`Buffer` returned by :class:`VirtualMemoryResource`. + + The buffer owns its address reservations, physical allocations and + mappings through its device pointer handle; closing it is the only way + to release them. A buffer returned by + :meth:`VirtualMemoryResource.modify_allocation` aliases the buffer it + was grown from: the two share their physical memory, and that memory is + freed when the last buffer that maps it closes. + """ + def close(self, stream: Stream | GraphBuilder | None=None) -> None: + """Release this buffer's share of its address range. + + The mappings, reservations and physical allocations go away when the + last buffer that maps them closes. Before it unmaps, the resource + synchronizes every deallocation stream the buffers of the range + recorded. Virtual memory deallocation is synchronous and cannot be + captured, so closing on a capturing stream raises and leaves the + buffer open. + + Parameters + ---------- + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional + If given, replaces the recorded deallocation stream, as for + :meth:`Buffer.close`. + """ + +class VirtualMemoryResource(MemoryResource): + """Create a device memory resource that uses the CUDA VMM APIs to allocate memory. + + Parameters + ---------- + device_id : Device | int + Device for which a memory resource is constructed. + + config : VirtualMemoryResourceOptions, optional + A configuration object for the VirtualMemoryResource + + + Warning + ------- + This is a low-level API that is provided only for convenience. Make sure you fully understand + how CUDA Virtual Memory Management works before using this. Other MemoryResource subclasses + in cuda.core should already meet the common needs. + + Notes + ----- + Every buffer this resource returns is a :class:`VirtualMemoryBuffer` that + owns its address reservations, physical allocations and mappings; closing + the buffer releases them. :meth:`deallocate` is not involved in that path. + """ + device: object + config: object + + def __init__(self, device_id: Device | int, config: VirtualMemoryResourceOptions | None=None) -> None: ... + def allocate(self, size: int, *, stream: Stream | GraphBuilder | None=None) -> VirtualMemoryBuffer: + """ + Allocate a buffer of the given size using CUDA virtual memory. + + Parameters + ---------- + size : int + The size in bytes of the buffer to allocate. It is rounded up to the + allocation granularity; the returned buffer reports the rounded size. + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional + Keyword-only. The allocation itself is synchronous. A real stream is + recorded as the buffer's deallocation stream and synchronized when + the buffer closes; with `None` or a default-stream token the legacy + default stream of the resource's device is recorded instead. + + Returns + ------- + VirtualMemoryBuffer + A buffer that owns its reservation, physical allocation and mapping. + + Raises + ------ + CUDAError + If any CUDA driver API call fails during allocation. Nothing is + left allocated when this method raises. + """ + def modify_allocation(self, buf: Buffer, new_size: int, config: VirtualMemoryResourceOptions | None=None) -> VirtualMemoryBuffer: + """ + Grow a buffer of this resource to at least ``new_size`` bytes. + + The buffer passed in stays open and usable. The returned buffer aliases + it: both map the same physical memory, which is freed when the last of + the two closes. When the driver can extend the address range in place, + the returned buffer has the same pointer; otherwise it has a new one and + the existing contents are reachable through both. + + This method is not thread-safe with respect to two buffers that share an + address range. + + Parameters + ---------- + buf : VirtualMemoryBuffer + A buffer returned by :meth:`allocate` or by this method. + new_size : int + The requested total size in bytes; rounded up to the granularity. + config : VirtualMemoryResourceOptions, optional + Configuration for the new physical memory chunk only. Existing + chunks keep the access they were created with, and the resource's + own configuration is unchanged. + + Returns + ------- + VirtualMemoryBuffer + ``buf`` itself when it already covers ``new_size``; otherwise a new + buffer of the rounded size. + + Raises + ------ + TypeError + If ``buf`` did not come from this resource. + CUDAError + If a driver call fails. ``buf`` is untouched when this method raises. + """ + def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | GraphBuilder | None=None) -> None: + """ + Unmap and free one address range that was reserved and mapped outside this resource. + + Buffers returned by :meth:`allocate` and :meth:`modify_allocation` free + themselves when they close and never call this method. It exists for + raw pointers wrapped with :meth:`Buffer.from_handle` with ``mr`` set to + this resource: the range must be exactly one reservation, and the caller + must already have released its own ``cuMemCreate`` handle, so the + physical memory is freed by the unmap. + + Parameters + ---------- + ptr : DevicePointerType + The start of the reservation. + size : int + The size of the reservation in bytes. + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional + Keyword-only. If given, ``stream.sync()`` is called before the + range is unmapped, except for a default-stream token on a + host-located resource, which has no context to synchronize in. + """ + @property + def is_device_accessible(self) -> bool: + """ + Indicates whether the allocated memory is accessible from the device. + """ + @property + def is_host_accessible(self) -> bool: + """ + Indicates whether the allocated memory is accessible from the host. + """ + @property + def is_ipc_enabled(self) -> bool: + """Return False. Buffers of this resource cannot be shared through IPC descriptors.""" + @property + def device_id(self) -> int: + """ + Get the device ID associated with this memory resource. + + Returns: + int: CUDA device ID. -1 if the memory resource allocates host memory + """ + def __repr__(self) -> str: + """ + Return a string representation of the VirtualMemoryResource. + + Returns: + str: A string describing the object + """ diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx new file mode 100644 index 00000000000..3ab62d43ec4 --- /dev/null +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -0,0 +1,701 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from libc.stdint cimport uintptr_t +from libc.string cimport memset +from libcpp.vector cimport vector + +from cuda.bindings cimport cydriver +from cuda.core._memory._buffer cimport ( + Buffer, + Buffer_check_open, + Buffer_from_deviceptr_handle, + MemoryResource, +) +from cuda.core._rt cimport ( + ContextHandle, + DevicePtrHandle, + MemAllocationHandle, + StreamHandle, + VaMappingHandle, + VaReservationHandle, + VmmRangeHandle, + as_cu, + create_context_bound_legacy_stream, + create_mem_allocation_handle, + create_va_mapping_handle, + create_va_reservation_handle, + create_vmm_range, + deallocation_stream, + deviceptr_create_ref, + deviceptr_create_vmm, + get_last_error, + get_primary_context, + mem_allocation_size, + set_deallocation_stream, + va_mapping_allocation, + vmm_range, + vmm_range_append, + vmm_range_count, + vmm_range_mapping, + vmm_range_reserve, + vmm_range_total, +) +from cuda.core._stream cimport Stream, Stream_accept, Stream_is_default_token +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, check_or_create_options + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Iterable + +from cuda.core._device import Device +from cuda.core._utils.cuda_utils import driver +from cuda.core._utils.version import binding_version +from cuda.core.typing import ( + DevicePointerType, + VirtualMemoryAccessType, + VirtualMemoryAllocationType, + VirtualMemoryGranularityType, + VirtualMemoryHandleType, + VirtualMemoryLocationType, +) + +if TYPE_CHECKING: + from cuda.core.graph import GraphBuilder + +__all__ = ["VirtualMemoryBuffer", "VirtualMemoryResource", "VirtualMemoryResourceOptions"] + +# Location types whose physical backing lives in host memory. Shared by +# VirtualMemoryResource.__init__ and is_host_accessible so the two cannot drift. +_HOST_LOCATION_TYPES = frozenset( + { + VirtualMemoryLocationType.HOST, + VirtualMemoryLocationType.HOST_NUMA, + VirtualMemoryLocationType.HOST_NUMA_CURRENT, + } +) + + +@dataclass +class VirtualMemoryResourceOptions: + """A configuration object for the VirtualMemoryResource + Stores configuration information which tells the resource how to use the CUDA VMM APIs + + Attributes + ---------- + allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str + Controls the type of allocation. + location_type: :obj:`~_memory.VirtualMemoryLocationType` | str + Controls the location of the allocation. + handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str + Export handle type for the physical allocation. Use ``"posix_fd"`` on + Linux if you plan to import/export the allocation. Use `None` if you + don't need an exportable handle. Host-located allocations require + `None`. + gpu_direct_rdma: bool + Hint that the allocation should be GDR-capable (if supported). + granularity: :obj:`~_memory.VirtualMemoryGranularityType` | str + Controls granularity query and size rounding. + addr_hint: int + A (optional) virtual address hint to try to reserve at. Setting it to 0 lets the CUDA driver decide. + addr_align: int + Alignment for the VA reservation. If `None`, use the queried granularity. + peers: Iterable[int] + Extra device IDs that should be granted access in addition to ``device``. + self_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str + Access flags for the owning device. + peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str + Access flags for peers. + """ + + allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED + location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE + handle_type: VirtualMemoryHandleType = VirtualMemoryHandleType.POSIX_FD + granularity: VirtualMemoryGranularityType = VirtualMemoryGranularityType.RECOMMENDED + gpu_direct_rdma: bool = False + addr_hint: int | None = 0 + addr_align: int | None = None + peers: Iterable[int] = field(default_factory=tuple) + self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE + peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE + + _a = driver.CUmemAccess_flags + _access_flags = { # noqa: RUF012 + VirtualMemoryAccessType.READ_WRITE: _a.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, + VirtualMemoryAccessType.READ: _a.CU_MEM_ACCESS_FLAGS_PROT_READ, + None: 0, + } + _h = driver.CUmemAllocationHandleType + _handle_types = { # noqa: RUF012 + None: _h.CU_MEM_HANDLE_TYPE_NONE, + VirtualMemoryHandleType.POSIX_FD: _h.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, + VirtualMemoryHandleType.WIN32_KMT: _h.CU_MEM_HANDLE_TYPE_WIN32_KMT, + VirtualMemoryHandleType.FABRIC: _h.CU_MEM_HANDLE_TYPE_FABRIC, + } + _g = driver.CUmemAllocationGranularity_flags + _granularity = { # noqa: RUF012 + VirtualMemoryGranularityType.RECOMMENDED: _g.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + VirtualMemoryGranularityType.MINIMUM: _g.CU_MEM_ALLOC_GRANULARITY_MINIMUM, + } + _l = driver.CUmemLocationType + _location_type = { # noqa: RUF012 + VirtualMemoryLocationType.DEVICE: _l.CU_MEM_LOCATION_TYPE_DEVICE, + VirtualMemoryLocationType.HOST: _l.CU_MEM_LOCATION_TYPE_HOST, + VirtualMemoryLocationType.HOST_NUMA: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA, + VirtualMemoryLocationType.HOST_NUMA_CURRENT: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, + } + _t = driver.CUmemAllocationType + # CUDA 13+ exposes MANAGED in CUmemAllocationType; older 12.x does not + _allocation_type = {VirtualMemoryAllocationType.PINNED: _t.CU_MEM_ALLOCATION_TYPE_PINNED} # noqa: RUF012 + if binding_version() >= (13, 0, 0): + _allocation_type[VirtualMemoryAllocationType.MANAGED] = _t.CU_MEM_ALLOCATION_TYPE_MANAGED + + @staticmethod + def _access_to_flags(spec: VirtualMemoryAccessType | None) -> int: + flags = VirtualMemoryResourceOptions._access_flags.get(spec) + if flags is None: + raise ValueError(f"Unknown access spec: {spec!r}") + return flags # type: ignore[no-any-return] + + @staticmethod + def _allocation_type_to_driver(spec: VirtualMemoryAllocationType) -> int: + alloc_type = VirtualMemoryResourceOptions._allocation_type.get(spec) + if alloc_type is None: + raise ValueError(f"Unsupported allocation_type: {spec!r}") + return alloc_type # type: ignore[no-any-return] + + @staticmethod + def _location_type_to_driver(spec: VirtualMemoryLocationType) -> int: + loc_type = VirtualMemoryResourceOptions._location_type.get(spec) + if loc_type is None: + raise ValueError(f"Unsupported location_type: {spec!r}") + return loc_type # type: ignore[no-any-return] + + @staticmethod + def _handle_type_to_driver(spec: VirtualMemoryHandleType | None) -> int: + if spec == "win32": + raise NotImplementedError("win32 is currently not supported, please reach out to the CUDA Python team") + handle_type = VirtualMemoryResourceOptions._handle_types.get(spec) + if handle_type is None: + raise ValueError(f"Unsupported handle_type: {spec!r}") + return handle_type # type: ignore[no-any-return] + + @staticmethod + def _granularity_to_driver(spec: VirtualMemoryGranularityType) -> int: + granularity = VirtualMemoryResourceOptions._granularity.get(spec) + if granularity is None: + raise ValueError(f"Unsupported granularity: {spec!r}") + return granularity # type: ignore[no-any-return] + + +cdef inline size_t _align_up(size_t size, size_t gran) noexcept nogil: + return (size + gran - 1) // gran * gran + + +cdef inline bint _is_default_token(cydriver.CUstream s) noexcept nogil: + cdef uintptr_t h = s + return h == 0 or h == cydriver.CU_STREAM_LEGACY or h == cydriver.CU_STREAM_PER_THREAD + + +cdef bint _stream_is_capturing(cydriver.CUstream s) except -1: + cdef cydriver.CUstreamCaptureStatus cap_status + IF CUDA_CORE_BUILD_MAJOR >= 13: + HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(s, &cap_status, NULL, NULL, NULL, NULL, NULL)) + ELSE: + HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(s, &cap_status, NULL, NULL, NULL, NULL)) + return cap_status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE + + +cdef int _raise_last_error() except -1: + """Raise the status a handle factory recorded when it returned empty.""" + HANDLE_RETURN(get_last_error()) + raise RuntimeError( + "internal cuda.core error, please report: a virtual memory handle factory " + "returned an empty handle without recording a CUDA error" + ) + + +cdef class VirtualMemoryBuffer(Buffer): + """A :class:`Buffer` returned by :class:`VirtualMemoryResource`. + + The buffer owns its address reservations, physical allocations and + mappings through its device pointer handle; closing it is the only way + to release them. A buffer returned by + :meth:`VirtualMemoryResource.modify_allocation` aliases the buffer it + was grown from: the two share their physical memory, and that memory is + freed when the last buffer that maps it closes. + """ + + def close(self, stream: Stream | GraphBuilder | None = None) -> None: + """Release this buffer's share of its address range. + + The mappings, reservations and physical allocations go away when the + last buffer that maps them closes. Before it unmaps, the resource + synchronizes every deallocation stream the buffers of the range + recorded. Virtual memory deallocation is synchronous and cannot be + captured, so closing on a capturing stream raises and leaves the + buffer open. + + Parameters + ---------- + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional + If given, replaces the recorded deallocation stream, as for + :meth:`Buffer.close`. + """ + cdef Stream s + cdef StreamHandle h + cdef cydriver.CUstream raw + if not self._h_ptr: + return + if stream is not None: + s = Stream_accept(stream) + raw = as_cu(s._h_stream) + else: + h = deallocation_stream(self._h_ptr) + raw = as_cu(h) + # Default-stream tokens are checked by the range deleter under their + # bound context; a real stream can be checked here and refused. + if not _is_default_token(raw) and _stream_is_capturing(raw): + raise RuntimeError( + "cannot close a VirtualMemoryResource buffer on a capturing stream: " + "virtual memory deallocation is synchronous and cannot be captured" + ) + Buffer.close(self, stream) + + +cdef class VirtualMemoryResource(MemoryResource): + """Create a device memory resource that uses the CUDA VMM APIs to allocate memory. + + Parameters + ---------- + device_id : Device | int + Device for which a memory resource is constructed. + + config : VirtualMemoryResourceOptions, optional + A configuration object for the VirtualMemoryResource + + + Warning + ------- + This is a low-level API that is provided only for convenience. Make sure you fully understand + how CUDA Virtual Memory Management works before using this. Other MemoryResource subclasses + in cuda.core should already meet the common needs. + + Notes + ----- + Every buffer this resource returns is a :class:`VirtualMemoryBuffer` that + owns its address reservations, physical allocations and mappings; closing + the buffer releases them. :meth:`deallocate` is not involved in that path. + """ + + cdef: + public object device + public object config + + def __init__(self, device_id: Device | int, config: VirtualMemoryResourceOptions | None = None) -> None: + self.device = Device(device_id) + self.config = check_or_create_options( + VirtualMemoryResourceOptions, config, "VirtualMemoryResource options", keep_none=False + ) + if self.config.location_type in _HOST_LOCATION_TYPES: + self.device = None + # The driver rejects an exportable handle type for host memory. + if self.config.handle_type is not None: + raise ValueError( + "host-located virtual memory cannot have an exportable handle type; " + "pass handle_type=None" + ) + + if self.device is not None and not self.device.properties.virtual_memory_management_supported: + raise RuntimeError("VirtualMemoryResource requires CUDA VMM API support") + + # Validate RDMA support if requested + if ( + self.config.gpu_direct_rdma + and self.device is not None + and not self.device.properties.gpu_direct_rdma_supported + ): + raise RuntimeError("GPU Direct RDMA is not supported on this device") + + cdef int _fill_prop(self, object cfg, cydriver.CUmemAllocationProp* prop) except -1: + # The location comes from the resource; the rest may come from a + # per-call configuration. + cdef int alloc_type = int(VirtualMemoryResourceOptions._allocation_type_to_driver(cfg.allocation_type)) + cdef int loc_type = int(VirtualMemoryResourceOptions._location_type_to_driver(self.config.location_type)) + cdef int handle_type = int(VirtualMemoryResourceOptions._handle_type_to_driver(cfg.handle_type)) + memset(prop, 0, sizeof(cydriver.CUmemAllocationProp)) + prop.type = alloc_type + prop.location.type = loc_type + prop.location.id = self.device.device_id if self.device is not None else -1 + prop.allocFlags.gpuDirectRDMACapable = 1 if cfg.gpu_direct_rdma else 0 + prop.requestedHandleTypes = handle_type + prop.win32HandleMetaData = NULL + return 0 + + cdef int _fill_access( + self, object cfg, const cydriver.CUmemAllocationProp* prop, + vector[cydriver.CUmemAccessDesc]& descs, + ) except -1: + cdef cydriver.CUmemAccessDesc d + cdef int owner_flags = int(VirtualMemoryResourceOptions._access_to_flags(cfg.self_access)) + cdef int peer_flags = int(VirtualMemoryResourceOptions._access_to_flags(cfg.peer_access)) + if owner_flags: + memset(&d, 0, sizeof(d)) + d.location.type = prop.location.type + d.location.id = prop.location.id + d.flags = owner_flags + descs.push_back(d) + if peer_flags: + for peer_dev in cfg.peers: + memset(&d, 0, sizeof(d)) + d.location.type = cydriver.CU_MEM_LOCATION_TYPE_DEVICE + d.location.id = int(peer_dev) + d.flags = peer_flags + descs.push_back(d) + return 0 + + cdef size_t _granularity(self, object cfg, const cydriver.CUmemAllocationProp* prop) except? 0: + cdef size_t gran = 0 + cdef int flag = int(VirtualMemoryResourceOptions._granularity_to_driver(cfg.granularity)) + with nogil: + HANDLE_RETURN(cydriver.cuMemGetAllocationGranularity( + &gran, prop, flag)) + return gran + + cdef int _record_deallocation_stream(self, const DevicePtrHandle& h_ptr, Stream s) except -1: + cdef StreamHandle h + cdef ContextHandle h_ctx + if s is not None and not Stream_is_default_token(s): + h = s._h_stream + elif self.device is None: + # Host-located memory needs no context to free, and a default-stream + # token has no context to bind to: record nothing. + return 0 + else: + # Bind the legacy default-stream token to the device's primary + # context so the free is ordered correctly no matter what is current + # then, and so allocate() never depends on a current context. + # get_primary_context keeps its own cache; nothing is cached here. + h_ctx = get_primary_context(self.device.device_id) + if not h_ctx: + _raise_last_error() + h = create_context_bound_legacy_stream(h_ctx) + if not h: + _raise_last_error() + HANDLE_RETURN(set_deallocation_stream(h_ptr, h)) + return 0 + + cdef int _copy_deallocation_stream(self, const DevicePtrHandle& dst, const DevicePtrHandle& src) except -1: + cdef StreamHandle h = deallocation_stream(src) + if h: + HANDLE_RETURN(set_deallocation_stream(dst, h)) + return 0 + + def allocate(self, size_t size, *, stream: Stream | GraphBuilder | None = None) -> VirtualMemoryBuffer: + """ + Allocate a buffer of the given size using CUDA virtual memory. + + Parameters + ---------- + size : int + The size in bytes of the buffer to allocate. It is rounded up to the + allocation granularity; the returned buffer reports the rounded size. + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional + Keyword-only. The allocation itself is synchronous. A real stream is + recorded as the buffer's deallocation stream and synchronized when + the buffer closes; with `None` or a default-stream token the legacy + default stream of the resource's device is recorded instead. + + Returns + ------- + VirtualMemoryBuffer + A buffer that owns its reservation, physical allocation and mapping. + + Raises + ------ + CUDAError + If any CUDA driver API call fails during allocation. Nothing is + left allocated when this method raises. + """ + cdef Stream s = None + if stream is not None: + s = Stream_accept(stream) + return self._allocate(self.config, size, s) + + cdef Buffer _allocate(self, object cfg, size_t size, Stream s): + """Allocate ``size`` bytes with ``cfg``; ``s`` is the accepted stream or None.""" + cdef cydriver.CUmemAllocationProp prop + cdef vector[cydriver.CUmemAccessDesc] descs + cdef size_t gran, aligned, addr_align + cdef cydriver.CUdeviceptr hint + cdef MemAllocationHandle h_alloc + cdef VaReservationHandle h_res + cdef VaMappingHandle h_map + cdef VmmRangeHandle rng + cdef DevicePtrHandle h_ptr + + if size == 0: + # Nothing to reserve or map; an empty buffer with a non-owning handle. + return Buffer_from_deviceptr_handle(deviceptr_create_ref(0), 0, self, None, VirtualMemoryBuffer) + + self._fill_prop(cfg, &prop) + self._fill_access(cfg, &prop, descs) + gran = self._granularity(cfg, &prop) + aligned = _align_up(size, gran) + addr_align = cfg.addr_align or gran + hint = cfg.addr_hint or 0 + + # Every handle below is a local: if a later step fails, the locals die + # in reverse order and undo everything created so far. + with nogil: + h_alloc = create_mem_allocation_handle(aligned, prop, descs.data(), descs.size()) + if not h_alloc: + _raise_last_error() + with nogil: + h_res = create_va_reservation_handle(aligned, addr_align, hint) + if not h_res: + _raise_last_error() + with nogil: + h_map = create_va_mapping_handle(as_cu(h_res), h_alloc, h_res) + if not h_map: + _raise_last_error() + + rng = create_vmm_range(as_cu(h_res)) + vmm_range_append(rng, h_map) + h_ptr = deviceptr_create_vmm(as_cu(h_res), rng) + if not h_ptr: + _raise_last_error() + self._record_deallocation_stream(h_ptr, s) + return Buffer_from_deviceptr_handle(h_ptr, aligned, self, None, VirtualMemoryBuffer) + + def modify_allocation( + self, buf: Buffer, size_t new_size, config: VirtualMemoryResourceOptions | None = None + ) -> VirtualMemoryBuffer: + """ + Grow a buffer of this resource to at least ``new_size`` bytes. + + The buffer passed in stays open and usable. The returned buffer aliases + it: both map the same physical memory, which is freed when the last of + the two closes. When the driver can extend the address range in place, + the returned buffer has the same pointer; otherwise it has a new one and + the existing contents are reachable through both. + + This method is not thread-safe with respect to two buffers that share an + address range. + + Parameters + ---------- + buf : VirtualMemoryBuffer + A buffer returned by :meth:`allocate` or by this method. + new_size : int + The requested total size in bytes; rounded up to the granularity. + config : VirtualMemoryResourceOptions, optional + Configuration for the new physical memory chunk only. Existing + chunks keep the access they were created with, and the resource's + own configuration is unchanged. + + Returns + ------- + VirtualMemoryBuffer + ``buf`` itself when it already covers ``new_size``; otherwise a new + buffer of the rounded size. + + Raises + ------ + TypeError + If ``buf`` did not come from this resource. + CUDAError + If a driver call fails. ``buf`` is untouched when this method raises. + """ + cdef Buffer b + cdef VmmRangeHandle rng, rng_new + cdef object cfg + cdef cydriver.CUmemAllocationProp prop + cdef vector[cydriver.CUmemAccessDesc] descs + cdef size_t gran, req, total, add, count, addr_align, offset, i, chunk + cdef cydriver.CUdeviceptr base, base_new + cdef MemAllocationHandle h_alloc, a + cdef VaReservationHandle h_res, h_res_new + cdef VaMappingHandle h_map, m, m2 + cdef DevicePtrHandle h_ptr2, h_ptr_new + cdef object new_buf + + if not isinstance(buf, Buffer): + raise TypeError(f"buf must be a Buffer, got {type(buf).__name__}") + b = buf + Buffer_check_open(b) + if b.memory_resource is not self: + raise TypeError("buf was not allocated by this VirtualMemoryResource") + cfg = self.config if config is None else check_or_create_options( + VirtualMemoryResourceOptions, config, "VirtualMemoryResource options", keep_none=False + ) + if b._size == 0: + # An empty buffer maps nothing; the request is a fresh allocation. + return self._allocate(cfg, new_size, None) + rng = vmm_range(b._h_ptr) + if not rng: + raise TypeError("buf was not allocated by VirtualMemoryResource.allocate") + + self._fill_prop(cfg, &prop) + self._fill_access(cfg, &prop, descs) + gran = self._granularity(cfg, &prop) + req = _align_up(new_size, gran) + total = vmm_range_total(rng) + base = as_cu(b._h_ptr) + + if req <= b._size: + return buf + if req <= total: + # A shorter alias asking for what the range already maps. + h_ptr2 = deviceptr_create_vmm(base, rng) + if not h_ptr2: + _raise_last_error() + self._copy_deallocation_stream(h_ptr2, b._h_ptr) + return Buffer_from_deviceptr_handle(h_ptr2, req, self, None, VirtualMemoryBuffer) + + # The new chunk is a whole number of granules; the result covers it. + add = _align_up(req - total, gran) + req = total + add + count = vmm_range_count(rng) + + # Grow in place: reserve the range right after the current one. The + # driver raises alignment 0 to its default, so the hint is well formed. + with nogil: + h_res = create_va_reservation_handle(add, 0, base + total) + if not h_res: + get_last_error() # the probe may fail; that is not an error here + elif as_cu(h_res) != base + total: + h_res.reset() # granted elsewhere: free it and move instead + else: + with nogil: + h_alloc = create_mem_allocation_handle(add, prop, descs.data(), descs.size()) + if not h_alloc: + _raise_last_error() + with nogil: + h_map = create_va_mapping_handle(base + total, h_alloc, h_res) + if not h_map: + _raise_last_error() + vmm_range_reserve(rng, count + 1) + h_ptr2 = deviceptr_create_vmm(base, rng) + if not h_ptr2: + _raise_last_error() + self._copy_deallocation_stream(h_ptr2, b._h_ptr) + new_buf = Buffer_from_deviceptr_handle(h_ptr2, req, self, None, VirtualMemoryBuffer) + # Last step, and it cannot throw after the reserve above: the input + # buffer is untouched if anything before this raised. + vmm_range_append(rng, h_map) + return new_buf + + # Move: a new range that maps every existing allocation, then the new one. + # The allocations are shared with the input buffer's range. + addr_align = cfg.addr_align or gran + with nogil: + h_res_new = create_va_reservation_handle(req, addr_align, 0) + if not h_res_new: + _raise_last_error() + base_new = as_cu(h_res_new) + rng_new = create_vmm_range(base_new) + vmm_range_reserve(rng_new, count + 1) + offset = 0 + for i in range(count): + m = vmm_range_mapping(rng, i) + a = va_mapping_allocation(m) + chunk = mem_allocation_size(a) + with nogil: + m2 = create_va_mapping_handle(base_new + offset, a, h_res_new) + if not m2: + _raise_last_error() + vmm_range_append(rng_new, m2) + offset += chunk + with nogil: + h_alloc = create_mem_allocation_handle(add, prop, descs.data(), descs.size()) + if not h_alloc: + _raise_last_error() + with nogil: + m2 = create_va_mapping_handle(base_new + offset, h_alloc, h_res_new) + if not m2: + _raise_last_error() + vmm_range_append(rng_new, m2) + h_ptr_new = deviceptr_create_vmm(base_new, rng_new) + if not h_ptr_new: + _raise_last_error() + self._copy_deallocation_stream(h_ptr_new, b._h_ptr) + return Buffer_from_deviceptr_handle(h_ptr_new, req, self, None, VirtualMemoryBuffer) + + def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | GraphBuilder | None = None) -> None: + """ + Unmap and free one address range that was reserved and mapped outside this resource. + + Buffers returned by :meth:`allocate` and :meth:`modify_allocation` free + themselves when they close and never call this method. It exists for + raw pointers wrapped with :meth:`Buffer.from_handle` with ``mr`` set to + this resource: the range must be exactly one reservation, and the caller + must already have released its own ``cuMemCreate`` handle, so the + physical memory is freed by the unmap. + + Parameters + ---------- + ptr : DevicePointerType + The start of the reservation. + size : int + The size of the reservation in bytes. + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional + Keyword-only. If given, ``stream.sync()`` is called before the + range is unmapped, except for a default-stream token on a + host-located resource, which has no context to synchronize in. + """ + cdef cydriver.CUdeviceptr devptr = 0 if ptr is None else int(ptr) + cdef size_t nbytes = size + cdef Stream s + if stream is not None: + s = Stream_accept(stream) + # A host-located resource records no stream, so Buffer teardown + # passes an unbound default-stream token, which has no context to + # synchronize in. There is nothing queued on it to wait for. + if self.device is not None or not Stream_is_default_token(s): + s.sync() + if devptr == 0 or nbytes == 0: + return + with nogil: + HANDLE_RETURN(cydriver.cuMemUnmap(devptr, nbytes)) + HANDLE_RETURN(cydriver.cuMemAddressFree(devptr, nbytes)) + + @property + def is_device_accessible(self) -> bool: + """ + Indicates whether the allocated memory is accessible from the device. + """ + return self.config.location_type == "device" + + @property + def is_host_accessible(self) -> bool: + """ + Indicates whether the allocated memory is accessible from the host. + """ + return self.config.location_type in _HOST_LOCATION_TYPES + + @property + def is_ipc_enabled(self) -> bool: + """Return False. Buffers of this resource cannot be shared through IPC descriptors.""" + return False + + @property + def device_id(self) -> int: + """ + Get the device ID associated with this memory resource. + + Returns: + int: CUDA device ID. -1 if the memory resource allocates host memory + """ + return self.device.device_id if self.device is not None else -1 + + def __repr__(self) -> str: + """ + Return a string representation of the VirtualMemoryResource. + + Returns: + str: A string describing the object + """ + return f"" diff --git a/cuda_core/cuda/core/_rt.pxd b/cuda_core/cuda/core/_rt.pxd index 76082d7ec0e..b8b2ca5e416 100644 --- a/cuda_core/cuda/core/_rt.pxd +++ b/cuda_core/cuda/core/_rt.pxd @@ -58,6 +58,22 @@ cdef extern from "_cpp/rt/handles.hpp" namespace "cuda_core::rt": ctypedef shared_ptr[const TexObjectValue] TexObjectHandle ctypedef shared_ptr[const SurfObjectValue] SurfObjectHandle + # Virtual memory management (VMM_DESIGN.md): tagged values for the + # physical allocation, the address reservation and the mapping, plus the + # opaque range a buffer owns. + cppclass MemAllocationValue "cuda_core::rt::MemAllocationValue": + pass + cppclass VaReservationValue "cuda_core::rt::VaReservationValue": + pass + cppclass VaMappingValue "cuda_core::rt::VaMappingValue": + pass + cppclass VmmRange "cuda_core::rt::VmmRange": + pass + ctypedef shared_ptr[const MemAllocationValue] MemAllocationHandle + ctypedef shared_ptr[const VaReservationValue] VaReservationHandle + ctypedef shared_ptr[const VaMappingValue] VaMappingHandle + ctypedef shared_ptr[VmmRange] VmmRangeHandle + # Type-erased shared owner for resources attached to graph node slots. # Typed handles above assign directly to an OpaqueHandle (shared control # block); make_opaque_py / make_opaque_malloc cover the two cases needing a @@ -108,6 +124,9 @@ cdef extern from "_cpp/rt/handles.hpp" namespace "cuda_core::rt": cydriver.CUmipmappedArray as_cu(MipmappedArrayHandle h) noexcept nogil cydriver.CUtexObject as_cu(TexObjectHandle h) noexcept nogil cydriver.CUsurfObject as_cu(SurfObjectHandle h) noexcept nogil + cydriver.CUmemGenericAllocationHandle as_cu(MemAllocationHandle h) noexcept nogil + cydriver.CUdeviceptr as_cu(VaReservationHandle h) noexcept nogil + cydriver.CUdeviceptr as_cu(VaMappingHandle h) noexcept nogil # as_intptr() - extract handle as intptr_t for Python interop (inline C++) intptr_t as_intptr(ContextHandle h) noexcept nogil @@ -132,6 +151,9 @@ cdef extern from "_cpp/rt/handles.hpp" namespace "cuda_core::rt": intptr_t as_intptr(MipmappedArrayHandle h) noexcept nogil intptr_t as_intptr(TexObjectHandle h) noexcept nogil intptr_t as_intptr(SurfObjectHandle h) noexcept nogil + intptr_t as_intptr(MemAllocationHandle h) noexcept nogil + intptr_t as_intptr(VaReservationHandle h) noexcept nogil + intptr_t as_intptr(VaMappingHandle h) noexcept nogil # as_py() - convert handle to Python wrapper object (inline C++; requires GIL) object as_py(ContextHandle h) @@ -156,6 +178,9 @@ cdef extern from "_cpp/rt/handles.hpp" namespace "cuda_core::rt": object as_py(MipmappedArrayHandle h) object as_py(TexObjectHandle h) object as_py(SurfObjectHandle h) + object as_py(MemAllocationHandle h) + object as_py(VaReservationHandle h) + object as_py(VaMappingHandle h) # ============================================================================= @@ -268,6 +293,29 @@ cdef DevicePtrHandle deviceptr_import_ipc( cdef StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept nogil cdef cydriver.CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil +# Virtual memory management (VMM_DESIGN.md) +cdef MemAllocationHandle create_mem_allocation_handle( + size_t size, const cydriver.CUmemAllocationProp& prop, + const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil +cdef size_t mem_allocation_size(const MemAllocationHandle& h) noexcept nogil +cdef VaReservationHandle create_va_reservation_handle( + size_t size, size_t alignment, cydriver.CUdeviceptr hint) except+ nogil +cdef size_t va_reservation_size(const VaReservationHandle& h) noexcept nogil +cdef VaMappingHandle create_va_mapping_handle( + cydriver.CUdeviceptr ptr, const MemAllocationHandle& h_alloc, + const VaReservationHandle& h_res) except+ nogil +cdef size_t va_mapping_size(const VaMappingHandle& h) noexcept nogil +cdef MemAllocationHandle va_mapping_allocation(const VaMappingHandle& h) noexcept nogil +cdef VmmRangeHandle create_vmm_range(cydriver.CUdeviceptr base) except+ nogil +cdef VmmRangeHandle vmm_range(const DevicePtrHandle& h) except+ nogil +cdef size_t vmm_range_count(const VmmRangeHandle& range) noexcept nogil +cdef VaMappingHandle vmm_range_mapping(const VmmRangeHandle& range, size_t index) noexcept nogil +cdef size_t vmm_range_total(const VmmRangeHandle& range) noexcept nogil +cdef void vmm_range_reserve(const VmmRangeHandle& range, size_t count) except+ nogil +cdef void vmm_range_append(const VmmRangeHandle& range, const VaMappingHandle& mapping) except+ nogil +cdef DevicePtrHandle deviceptr_create_vmm( + cydriver.CUdeviceptr base, const VmmRangeHandle& range) except+ nogil + # Library handles cdef LibraryHandle create_library_handle_from_file(const char* path) except+ nogil cdef LibraryHandle create_library_handle_from_data(const void* data) except+ nogil diff --git a/cuda_core/cuda/core/_rt.pyi b/cuda_core/cuda/core/_rt.pyi index b8044d7ea72..49f3ab9dd67 100644 --- a/cuda_core/cuda/core/_rt.pyi +++ b/cuda_core/cuda/core/_rt.pyi @@ -27,6 +27,10 @@ OpaqueArrayHandle: TypeAlias = Incomplete MipmappedArrayHandle: TypeAlias = Incomplete TexObjectHandle: TypeAlias = Incomplete SurfObjectHandle: TypeAlias = Incomplete +MemAllocationHandle: TypeAlias = Incomplete +VaReservationHandle: TypeAlias = Incomplete +VaMappingHandle: TypeAlias = Incomplete +VmmRangeHandle: TypeAlias = Incomplete OpaqueHandle: TypeAlias = Incomplete PreparedAttachment: TypeAlias = Incomplete PreparedChildGraphUpdate: TypeAlias = Incomplete @@ -36,6 +40,10 @@ NvvmProgramValue: TypeAlias = Incomplete NvJitLinkValue: TypeAlias = Incomplete TexObjectValue: TypeAlias = Incomplete SurfObjectValue: TypeAlias = Incomplete +MemAllocationValue: TypeAlias = Incomplete +VaReservationValue: TypeAlias = Incomplete +VaMappingValue: TypeAlias = Incomplete +VmmRange: TypeAlias = Incomplete PreparedAttachmentState: TypeAlias = Incomplete PreparedAttachmentDeleter: TypeAlias = Incomplete PreparedChildGraphUpdateState: TypeAlias = Incomplete diff --git a/cuda_core/cuda/core/_rt.pyx b/cuda_core/cuda/core/_rt.pyx index f92b051c0d7..69fcf4f702b 100644 --- a/cuda_core/cuda/core/_rt.pyx +++ b/cuda_core/cuda/core/_rt.pyx @@ -168,6 +168,36 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt": cydriver.CUresult set_deallocation_stream "cuda_core::rt::set_deallocation_stream" ( const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil + # Virtual memory management (VMM_DESIGN.md) + MemAllocationHandle create_mem_allocation_handle "cuda_core::rt::create_mem_allocation_handle" ( + size_t size, const cydriver.CUmemAllocationProp& prop, + const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil + size_t mem_allocation_size "cuda_core::rt::mem_allocation_size" ( + const MemAllocationHandle& h) noexcept nogil + VaReservationHandle create_va_reservation_handle "cuda_core::rt::create_va_reservation_handle" ( + size_t size, size_t alignment, cydriver.CUdeviceptr hint) except+ nogil + size_t va_reservation_size "cuda_core::rt::va_reservation_size" ( + const VaReservationHandle& h) noexcept nogil + VaMappingHandle create_va_mapping_handle "cuda_core::rt::create_va_mapping_handle" ( + cydriver.CUdeviceptr ptr, const MemAllocationHandle& h_alloc, + const VaReservationHandle& h_res) except+ nogil + size_t va_mapping_size "cuda_core::rt::va_mapping_size" (const VaMappingHandle& h) noexcept nogil + MemAllocationHandle va_mapping_allocation "cuda_core::rt::va_mapping_allocation" ( + const VaMappingHandle& h) noexcept nogil + VmmRangeHandle create_vmm_range "cuda_core::rt::create_vmm_range" ( + cydriver.CUdeviceptr base) except+ nogil + VmmRangeHandle vmm_range "cuda_core::rt::vmm_range" (const DevicePtrHandle& h) except+ nogil + size_t vmm_range_count "cuda_core::rt::vmm_range_count" (const VmmRangeHandle& range) noexcept nogil + VaMappingHandle vmm_range_mapping "cuda_core::rt::vmm_range_mapping" ( + const VmmRangeHandle& range, size_t index) noexcept nogil + size_t vmm_range_total "cuda_core::rt::vmm_range_total" (const VmmRangeHandle& range) noexcept nogil + void vmm_range_reserve "cuda_core::rt::vmm_range_reserve" ( + const VmmRangeHandle& range, size_t count) except+ nogil + void vmm_range_append "cuda_core::rt::vmm_range_append" ( + const VmmRangeHandle& range, const VaMappingHandle& mapping) except+ nogil + DevicePtrHandle deviceptr_create_vmm "cuda_core::rt::deviceptr_create_vmm" ( + cydriver.CUdeviceptr base, const VmmRangeHandle& range) except+ nogil + # Library handles LibraryHandle create_library_handle_from_file "cuda_core::rt::create_library_handle_from_file" ( const char* path) except+ nogil @@ -384,6 +414,17 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt": # IPC void* p_cuMemPoolImportPointer "reinterpret_cast(cuda_core::rt::p_cuMemPoolImportPointer)" + # Virtual memory management + void* p_cuMemCreate "reinterpret_cast(cuda_core::rt::p_cuMemCreate)" + void* p_cuMemRelease "reinterpret_cast(cuda_core::rt::p_cuMemRelease)" + void* p_cuMemAddressReserve "reinterpret_cast(cuda_core::rt::p_cuMemAddressReserve)" + void* p_cuMemAddressFree "reinterpret_cast(cuda_core::rt::p_cuMemAddressFree)" + void* p_cuMemMap "reinterpret_cast(cuda_core::rt::p_cuMemMap)" + void* p_cuMemUnmap "reinterpret_cast(cuda_core::rt::p_cuMemUnmap)" + void* p_cuMemSetAccess "reinterpret_cast(cuda_core::rt::p_cuMemSetAccess)" + void* p_cuStreamSynchronize "reinterpret_cast(cuda_core::rt::p_cuStreamSynchronize)" + void* p_cuStreamGetCaptureInfo "reinterpret_cast(cuda_core::rt::p_cuStreamGetCaptureInfo)" + # Library void* p_cuLibraryLoadFromFile "reinterpret_cast(cuda_core::rt::p_cuLibraryLoadFromFile)" void* p_cuLibraryLoadData "reinterpret_cast(cuda_core::rt::p_cuLibraryLoadData)" @@ -465,6 +506,9 @@ cdef void _init_driver_fn_pointers() noexcept: global p_cuMemAllocFromPoolAsync, p_cuMemAllocAsync, p_cuMemAlloc, p_cuMemAllocHost global p_cuMemFreeAsync, p_cuMemFree, p_cuMemFreeHost global p_cuMemPoolImportPointer + global p_cuMemCreate, p_cuMemRelease, p_cuMemAddressReserve, p_cuMemAddressFree + global p_cuMemMap, p_cuMemUnmap, p_cuMemSetAccess + global p_cuStreamSynchronize, p_cuStreamGetCaptureInfo global p_cuLibraryLoadFromFile, p_cuLibraryLoadData, p_cuLibraryUnload, p_cuLibraryGetKernel global p_cuGraphDestroy, p_cuGraphInstantiateWithParams global p_cuGraphExecUpdate, p_cuGraphExecDestroy @@ -534,6 +578,16 @@ cdef void _init_driver_fn_pointers() noexcept: # IPC p_cuMemPoolImportPointer = _get_driver_fn("cuMemPoolImportPointer") + p_cuMemCreate = _get_driver_fn("cuMemCreate") + p_cuMemRelease = _get_driver_fn("cuMemRelease") + p_cuMemAddressReserve = _get_driver_fn("cuMemAddressReserve") + p_cuMemAddressFree = _get_driver_fn("cuMemAddressFree") + p_cuMemMap = _get_driver_fn("cuMemMap") + p_cuMemUnmap = _get_driver_fn("cuMemUnmap") + p_cuMemSetAccess = _get_driver_fn("cuMemSetAccess") + p_cuStreamSynchronize = _get_driver_fn("cuStreamSynchronize") + p_cuStreamGetCaptureInfo = _get_driver_fn("cuStreamGetCaptureInfo") + # Library p_cuLibraryLoadFromFile = _get_driver_fn("cuLibraryLoadFromFile") p_cuLibraryLoadData = _get_driver_fn("cuLibraryLoadData") diff --git a/cuda_core/cuda/core/_utils/cuda_utils.pyi b/cuda_core/cuda/core/_utils/cuda_utils.pyi index f5b5134d98b..1a1fbf8be5f 100644 --- a/cuda_core/cuda/core/_utils/cuda_utils.pyi +++ b/cuda_core/cuda/core/_utils/cuda_utils.pyi @@ -52,38 +52,6 @@ class ComputeCapability(NamedTuple): major: int minor: int -class Transaction: - """ - A context manager for transactional operations with failure and exit callbacks. - - Failure callbacks are executed in LIFO order if the transaction exits without being committed. - Exit callbacks always run: in LIFO order on rollback or FIFO order during commit. - - Usage: - with Transaction() as txn: - txn.on_failure(some_cleanup_function, arg1, arg2) - txn.on_exit(some_finalize_function, arg1, arg2) - # ... perform operations ... - txn.commit() - - Methods: - on_failure(fn, *args, **kwargs): Register a callback to be called on rollback. - on_exit(fn, *args, **kwargs): Register a callback to be called on rollback or commit. - commit(): Disarm failure callbacks and run exit callbacks. - """ - def __init__(self) -> None: ... - def __enter__(self): ... - def __exit__(self, exc_type, exc, tb): ... - def _register(self, callback: Callable[[], Any], on_commit: bool) -> None: ... - def on_failure(self, fn: Callable[..., Any], /, *args: Any, **kwargs) -> None: - """Register a failure callback (runs if the with-block exits without commit()).""" - def on_exit(self, fn: Callable[..., Any], /, *args: Any, **kwargs) -> None: - """Register an exit callback (runs exactly once, on rollback or during commit()).""" - def commit(self) -> None: - """ - Disarm all failure callbacks, then run exit callbacks in FIFO order. - """ - def cast_to_3_tuple(label: str, cfg: int | tuple[int, ...]) -> tuple[int, int, int]: ... def _check_driver_error(error: cydriver.CUresult) -> int: ... def _check_runtime_error(error) -> int: ... diff --git a/cuda_core/cuda/core/_utils/cuda_utils.pyx b/cuda_core/cuda/core/_utils/cuda_utils.pyx index 14a66072c13..d1431fc3bf3 100644 --- a/cuda_core/cuda/core/_utils/cuda_utils.pyx +++ b/cuda_core/cuda/core/_utils/cuda_utils.pyx @@ -3,12 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 import functools -from functools import partial import multiprocessing import platform import warnings from collections.abc import Sequence -from contextlib import ExitStack from typing import Any, Callable, NamedTuple from cuda.bindings import driver as driver, nvrtc as nvrtc, runtime as runtime @@ -329,71 +327,6 @@ def is_nested_sequence(obj: object) -> bool: return is_sequence(obj) and any(is_sequence(elem) for elem in obj) - -class Transaction: - """ - A context manager for transactional operations with failure and exit callbacks. - - Failure callbacks are executed in LIFO order if the transaction exits without being committed. - Exit callbacks always run: in LIFO order on rollback or FIFO order during commit. - - Usage: - with Transaction() as txn: - txn.on_failure(some_cleanup_function, arg1, arg2) - txn.on_exit(some_finalize_function, arg1, arg2) - # ... perform operations ... - txn.commit() - - Methods: - on_failure(fn, *args, **kwargs): Register a callback to be called on rollback. - on_exit(fn, *args, **kwargs): Register a callback to be called on rollback or commit. - commit(): Disarm failure callbacks and run exit callbacks. - """ - def __init__(self) -> None: - self._stack = ExitStack() - self._on_exit: list[Callable[[], Any]] = [] - self._entered = False - - def __enter__(self): - self._stack.__enter__() - self._entered = True - return self - - def __exit__(self, exc_type, exc, tb): - # If exit callbacks remain, they'll run in LIFO order. - self._entered = False - self._on_exit.clear() - return self._stack.__exit__(exc_type, exc, tb) - - def _register(self, callback: Callable[[], Any], on_commit: bool) -> None: - if not self._entered: - raise RuntimeError("Transaction must be entered before registering callbacks") - # The ExitStack copy runs on rollback (LIFO, interleaved with the failure - # callbacks); the _on_exit copy runs at commit(). commit() disarms the stack - # before running _on_exit, so exactly one of the two ever fires. - self._stack.callback(callback) - if on_commit: - self._on_exit.append(callback) - - def on_failure(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: - """Register a failure callback (runs if the with-block exits without commit()).""" - self._register(partial(fn, *args, **kwargs), on_commit=False) - - def on_exit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: - """Register an exit callback (runs exactly once, on rollback or during commit()).""" - self._register(partial(fn, *args, **kwargs), on_commit=True) - - def commit(self) -> None: - """ - Disarm all failure callbacks, then run exit callbacks in FIFO order. - """ - # pop_all() empties this stack so no callbacks are triggered on exit. - self._stack.pop_all() - for fn in self._on_exit: - fn() - self._on_exit.clear() - - # Track whether we've already warned about fork method _fork_warning_checked = False diff --git a/cuda_core/docs/source/api.rst b/cuda_core/docs/source/api.rst index 76a228625e8..1705e1dc41e 100644 --- a/cuda_core/docs/source/api.rst +++ b/cuda_core/docs/source/api.rst @@ -62,6 +62,7 @@ Memory management Buffer ManagedBuffer + VirtualMemoryBuffer MemoryResource DeviceMemoryResource GraphMemoryResource diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index 31598150775..069d8346595 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -31,6 +31,10 @@ New features top-level ``cuda.core`` namespace, alongside :class:`CUDAWarning`. Previously they were only available from a private module. +- Added :class:`VirtualMemoryBuffer`, the :class:`Buffer` subclass that + :class:`VirtualMemoryResource` returns. Closing it on a capturing stream + raises, because virtual memory deallocation cannot be captured. + Fixes and enhancements ---------------------- @@ -134,3 +138,37 @@ Fixes and enhancements :attr:`graph.GraphNode.succ` views with an object that is not a :class:`~graph.GraphNode` is now a no-op instead of raising ``TypeError``, consistent with :class:`collections.abc.MutableSet` semantics. + +- :class:`VirtualMemoryResource` now holds each physical allocation, address + reservation and mapping in a reference-counted handle, so everything a + buffer maps is released when the last buffer that maps it closes. The + physical memory was leaked on every ``allocate()`` and grow + (`#2344 `__), the address + range added by an in-place grow could not be freed + (`#2887 `__), a second + grow of a buffer failed + (`#2908 `__), an + unaligned grow recorded a size the driver rejects at close + (`#2907 `__), and a + buffer was unmapped without waiting for the work queued on its stream + (`#2886 `__). No + reservation is freed by hand any more + (`#2345 `__, + `#2388 `__). The + behavior of the resource changes as follows. + :meth:`VirtualMemoryResource.modify_allocation` returns a new + :class:`VirtualMemoryBuffer` and leaves the buffer passed in open; the two + alias the same physical memory, which is freed when the last of them closes, + and the pointer is preserved when the driver grants the adjacent address + range. A request the buffer already covers returns the buffer itself, and + :attr:`Buffer.size` after a grow is the aligned total. The ``config`` + argument applies to the chunk that call adds and is no longer stored on the + resource (`#2909 `__). + Buffers from ``allocate()`` free themselves when they close and no longer + call :meth:`VirtualMemoryResource.deallocate`, which now serves pointers + wrapped with :meth:`Buffer.from_handle`. A buffer records the stream passed + to ``allocate()`` and the last close of an aliased range synchronizes every + recorded stream before it unmaps. ``modify_allocation`` accepts only buffers + this resource returned. ``location_type="host"`` requires + ``handle_type=None``, which the driver requires. ``allocate(0)`` returns an + empty buffer without a driver call. diff --git a/cuda_core/tests/graph/test_graph_definition_lifetime.py b/cuda_core/tests/graph/test_graph_definition_lifetime.py index 0f174ae240e..a84e0d2c207 100644 --- a/cuda_core/tests/graph/test_graph_definition_lifetime.py +++ b/cuda_core/tests/graph/test_graph_definition_lifetime.py @@ -16,7 +16,7 @@ from helpers.memory import xfail_on_graph_mempool_oom from helpers.misc import try_create_condition -from cuda_python_test_helpers import under_compute_sanitizer +from cuda_python_test_helpers import IS_WINDOWS, under_compute_sanitizer # Resource finalization triggered by graph destruction is not synchronous. A # CUDA user-object callback transfers each node attachment bundle to a @@ -75,7 +75,16 @@ def _wait_until(predicate, timeout=None, interval=0.02): raise AssertionError(f"condition not satisfied within {timeout}s") -from cuda.core import Device, DeviceMemoryResource, EventOptions, Kernel, LaunchConfig, LegacyPinnedMemoryResource +from cuda.core import ( + Device, + DeviceMemoryResource, + EventOptions, + Kernel, + LaunchConfig, + LegacyPinnedMemoryResource, + VirtualMemoryResource, + VirtualMemoryResourceOptions, +) from cuda.core._utils.cuda_utils import CUDAError from cuda.core._utils.version import driver_version from cuda.core.graph import ( @@ -96,6 +105,26 @@ def _skip_if_no_mempool(): pytest.skip("Device does not support mempool operations") +def _device_memory_resource(dev): + _skip_if_no_mempool() + return DeviceMemoryResource(dev) + + +def _virtual_memory_resource(dev): + if not dev.properties.virtual_memory_management_supported: + pytest.skip("Device does not support virtual memory management") + handle_type = "win32_kmt" if IS_WINDOWS else "posix_fd" + return VirtualMemoryResource(dev, config=VirtualMemoryResourceOptions(handle_type=handle_type)) + + +# Memory resources whose buffers a graph node can retain through its +# attachment. Each factory skips when the device lacks the feature. +_MEMORY_RESOURCES = [ + pytest.param(_device_memory_resource, id="device_mr"), + pytest.param(_virtual_memory_resource, id="vmm"), +] + + # ============================================================================= # Conditional body graph lifetime # ============================================================================= @@ -865,13 +894,18 @@ def callback(): @pytest.mark.agent_authored(model="gpt-5.6") -def test_inflight_launch_retains_attachments_until_completion(init_cuda): - """An in-flight launch retains the final allocation reference.""" +@pytest.mark.parametrize("make_mr", _MEMORY_RESOURCES) +def test_inflight_launch_retains_attachments_until_completion(init_cuda, make_mr): + """An in-flight launch retains the final allocation reference. + + The memcpy node attaches the buffer's allocation handle, so the memory + outlives ``buf.close()`` until the launch completes and the graph is gone. + For a virtual memory buffer that release also unmaps the range. + """ from cuda.core._utils._weak_handles import weak_handle - _skip_if_no_mempool() dev = Device() - mr = DeviceMemoryResource(dev) + mr = make_mr(dev) buf = mr.allocate(8, stream=dev.default_stream) dev.default_stream.sync() dptr = int(buf.handle) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index c20e263262b..25d1b6b7513 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -3,7 +3,9 @@ import ctypes import multiprocessing as mp +import subprocess import sys +import textwrap from cuda.bindings import driver @@ -49,6 +51,7 @@ MemoryResource, PinnedMemoryResource, PinnedMemoryResourceOptions, + VirtualMemoryBuffer, VirtualMemoryResource, VirtualMemoryResourceOptions, ) @@ -129,6 +132,7 @@ def test_package_contents(): "MemoryResource", "PinnedMemoryResource", "PinnedMemoryResourceOptions", + "VirtualMemoryBuffer", "VirtualMemoryResource", "VirtualMemoryResourceOptions", ] @@ -1249,6 +1253,74 @@ def get_handle_type(): return (("posix_fd", None),) +VMM_HANDLE_TYPE = "win32_kmt" if IS_WINDOWS else "posix_fd" +MIB = 1024 * 1024 + + +def _vmm_device_or_skip(): + device = Device() + device.set_current() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + return device + + +def _vmm_resource(device, **options): + options.setdefault("handle_type", VMM_HANDLE_TYPE) + return VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(**options)) + + +def _vmm_fill(buf, value, *, offset=0, size=None): + """Write ``value`` to ``size`` bytes of ``buf`` at ``offset`` and wait for the write.""" + size = buf.size - offset if size is None else size + handle_return(driver.cuMemsetD8(int(buf.handle) + offset, value, size)) + handle_return(driver.cuCtxSynchronize()) + + +def _vmm_read(buf, offset, size): + """Copy ``size`` bytes of ``buf`` at ``offset`` to the host.""" + host = (ctypes.c_ubyte * size)() + handle_return(driver.cuMemcpyDtoH(ctypes.addressof(host), int(buf.handle) + offset, size)) + return bytes(host) + + +def _vmm_free_memory(): + return handle_return(driver.cuMemGetInfo())[0] + + +def _vmm_reserve_at(addr, size): + """Reserve ``size`` bytes at ``addr``. Return the reservation, or None if the driver placed it elsewhere.""" + ptr = handle_return(driver.cuMemAddressReserve(size, 0, addr, 0)) + if int(ptr) != addr: + handle_return(driver.cuMemAddressFree(ptr, size)) + return None + return ptr + + +def _vmm_allocate_in_free_hole(device, size, extra): + """Allocate ``size`` bytes at the start of a free address range ``extra`` bytes longer. + + Reserving and freeing a range finds free address space; a resource with + that address as its hint then places the buffer there, so the range right + after the buffer is known to be free. Sizes round up to the granularity. + Skips when the driver places the buffer elsewhere. Returns the resource + and the buffer. + """ + probe = _vmm_resource(device).allocate(1) + gran = probe.size + probe.close() + size = -(-size // gran) * gran + total = size + -(-extra // gran) * gran + hole = handle_return(driver.cuMemAddressReserve(total, 0, 0, 0)) + handle_return(driver.cuMemAddressFree(hole, total)) + mr = _vmm_resource(device, addr_hint=int(hole)) + buf = mr.allocate(size) + if int(buf.handle) != int(hole): + buf.close() + pytest.skip("the driver did not place the allocation at the requested address") + return mr, buf + + @pytest.mark.parametrize("use_device_object", [True, False]) @pytest.mark.parametrize("handle_type", get_handle_type()) def test_vmm_allocator_basic_allocation(use_device_object, handle_type): @@ -1257,12 +1329,7 @@ def test_vmm_allocator_basic_allocation(use_device_object, handle_type): This test verifies that VirtualMemoryResource can allocate memory using CUDA VMM APIs with default configuration. """ - device = Device() - device.set_current() - - # Skip if virtual memory management is not supported - if not device.properties.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") + device = _vmm_device_or_skip() handle_type, security_attribute = handle_type # unpack options = VirtualMemoryResourceOptions(handle_type=handle_type) @@ -1276,12 +1343,15 @@ def test_vmm_allocator_basic_allocation(use_device_object, handle_type): except NotImplementedError: assert handle_type == "win32" return + assert isinstance(buffer, VirtualMemoryBuffer) assert buffer.size >= 4096 # May be aligned up assert buffer.device_id == device.device_id assert buffer.memory_resource == vmm_mr + assert vmm_mr.is_ipc_enabled is False # Test deallocation buffer.close() + assert buffer.size == 0 # Test multiple allocations buffers = [] @@ -1295,19 +1365,15 @@ def test_vmm_allocator_basic_allocation(use_device_object, handle_type): buf.close() +@pytest.mark.agent_authored(model="claude-fable-5-1") def test_vmm_allocator_policy_configuration(): """Test VMM allocator with different policy configurations. - This test verifies that VirtualMemoryResource can be configured - with different allocation policies and that the configuration affects - the allocation behavior. + The resource applies its configuration to every allocation. A + configuration passed to ``modify_allocation`` applies to the chunk that + call adds, and the resource keeps its own (#2909). """ - device = Device() - device.set_current() - - # Skip if virtual memory management is not supported - if not device.properties.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") + device = _vmm_device_or_skip() # Skip if GPU Direct RDMA is not supported if not device.properties.gpu_direct_rdma_supported: @@ -1343,7 +1409,7 @@ def test_vmm_allocator_policy_configuration(): assert buffer.size >= 8192 assert buffer.device_id == device.device_id - # Test policy modification + # A read-only configuration for the chunk the grow adds. new_config = VirtualMemoryResourceOptions( allocation_type="pinned", location_type="device", @@ -1355,35 +1421,39 @@ def test_vmm_allocator_policy_configuration(): peer_access="r", ) - # Modify allocation policy try: - modified_buffer = vmm_mr.modify_allocation(buffer, 16384, config=new_config) + grown = vmm_mr.modify_allocation(buffer, buffer.size + 8192, config=new_config) except CUDAError as exc: msg = str(exc) if "CUDA_ERROR_UNKNOWN" in msg: pytest.xfail("TODO(#1300): Known to fail already with CTK 13.0 (Windows)") raise - assert modified_buffer.size >= 16384 - assert vmm_mr.config == new_config - assert vmm_mr.config.self_access == "r" + assert grown.size >= buffer.size + 8192 + # The resource keeps its own configuration. + assert vmm_mr.config == custom_config + # The chunk the input already mapped stays writable; the new chunk is readable. + _vmm_fill(grown, 0x42, size=buffer.size) + assert _vmm_read(grown, 0, 16) == bytes([0x42]) * 16 + assert len(_vmm_read(grown, buffer.size, 16)) == 16 + # A later allocation uses the resource's configuration, which is writable. + later = vmm_mr.allocate(4096) + _vmm_fill(later, 0x24) # Clean up - modified_buffer.close() + later.close() + buffer.close() + grown.close() +@pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.parametrize("handle_type", get_handle_type()) def test_vmm_allocator_grow_allocation(handle_type): """Test VMM allocator's ability to grow existing allocations. - This test verifies that VirtualMemoryResource can grow existing - allocations while preserving the base pointer when possible. + ``modify_allocation`` returns a new buffer that aliases the input, which + stays open. A request the buffer already covers returns the buffer itself. """ - device = Device() - device.set_current() - - # Skip if virtual memory management is not supported (we need it for VMM) - if not device.properties.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") + device = _vmm_device_or_skip() handle_type, security_attribute = handle_type # unpack options = VirtualMemoryResourceOptions(handle_type=handle_type) @@ -1391,151 +1461,339 @@ def test_vmm_allocator_grow_allocation(handle_type): # Create initial allocation try: - buffer = vmm_mr.allocate(2 * 1024 * 1024) + buffer = vmm_mr.allocate(2 * MIB) except NotImplementedError: assert handle_type == "win32" return original_size = buffer.size # Grow the allocation - grown_buffer = vmm_mr.modify_allocation(buffer, 4 * 1024 * 1024) + grown_buffer = vmm_mr.modify_allocation(buffer, 4 * MIB) # Verify growth - assert grown_buffer.size >= 4 * 1024 * 1024 + assert grown_buffer is not buffer + assert grown_buffer.size >= 4 * MIB assert grown_buffer.size > original_size - # Because of the slow path, the pointer may change - # We cannot assert that the new pointer is the same, - # but we can assert that a new pointer was assigned - assert grown_buffer.handle is not None + # The input stays open, at its old size and address. + assert buffer.size == original_size + assert int(buffer.handle) != 0 - # Test growing to same size (should return original buffer) - same_buffer = vmm_mr.modify_allocation(grown_buffer, 4 * 1024 * 1024) - assert same_buffer.size == grown_buffer.size - - # Test growing to smaller size (should return original buffer) - smaller_buffer = vmm_mr.modify_allocation(grown_buffer, 2 * 1024 * 1024) - assert smaller_buffer.size == grown_buffer.size + # Requests the buffer already covers return the same object. + assert vmm_mr.modify_allocation(grown_buffer, 4 * MIB) is grown_buffer + assert vmm_mr.modify_allocation(grown_buffer, 2 * MIB) is grown_buffer # Clean up + buffer.close() grown_buffer.close() -def test_vmm_allocator_grow_allocation_fast_path(init_cuda, monkeypatch): - """Exercise the VMM grow fast path with mocked driver calls. +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_grow_in_place_keeps_pointer(init_cuda): + """With free address space after the buffer, a grow extends it in place and keeps the pointer.""" + device = _vmm_device_or_skip() + mr, buf = _vmm_allocate_in_free_hole(device, 2 * MIB, 4 * MIB) + size = buf.size + grown = mr.modify_allocation(buf, 2 * size) + if int(grown.handle) != int(buf.handle): + grown.close() + buf.close() + pytest.skip("the driver did not grant the adjacent address range") + assert grown.size == 2 * size + _vmm_fill(grown, 0x10) + assert _vmm_read(buf, 0, size) == bytes([0x10]) * size + buf.close() + grown.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_grow_preserves_contents_and_aliases_input(init_cuda): + """A grown buffer sees the input's contents, and a write through either buffer is visible through the other.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + buf = mr.allocate(2 * MIB) + size = buf.size + _vmm_fill(buf, 0xAB) + + grown = mr.modify_allocation(buf, size + 2 * MIB) + assert isinstance(grown, VirtualMemoryBuffer) + assert grown.size == 2 * size + assert _vmm_read(grown, 0, size) == bytes([0xAB]) * size + + _vmm_fill(grown, 0xCD, size=size) + assert _vmm_read(buf, 0, size) == bytes([0xCD]) * size + _vmm_fill(grown, 0xEF, offset=size) + assert _vmm_read(grown, size, size) == bytes([0xEF]) * size + + # Closing the input first leaves the result usable. + buf.close() + assert _vmm_read(grown, 0, 16) == bytes([0xCD]) * 16 + grown.close() - The real driver usually rejects the adjacent reservation that reaches this - path, so the test supplies that precondition by construction and verifies - the successful commit bookkeeping. - """ - device = Device() - if not device.properties.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") - vmm_mr = VirtualMemoryResource( - device, - config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), - ) +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_grow_moves_when_adjacent_range_is_taken(init_cuda): + """When the address range after the buffer is taken, the grow maps the existing memory at a new address.""" + device = _vmm_device_or_skip() + mr, buf = _vmm_allocate_in_free_hole(device, 2 * MIB, 4 * MIB) + size = buf.size + _vmm_fill(buf, 0x5A) + + decoy = _vmm_reserve_at(int(buf.handle) + size, size) + if decoy is None: + buf.close() + pytest.skip("the driver did not grant a reservation right after the buffer") + try: + grown = mr.modify_allocation(buf, 2 * size) + finally: + handle_return(driver.cuMemAddressFree(decoy, size)) + + # The range after the buffer was taken, so the grow had to move. + assert int(grown.handle) != int(buf.handle) + assert grown.size == 2 * size + assert _vmm_read(grown, 0, size) == bytes([0x5A]) * size + _vmm_fill(buf, 0x3C) + assert _vmm_read(grown, 0, size) == bytes([0x3C]) * size + + # Closing the result first leaves the input usable. + grown.close() + assert _vmm_read(buf, 0, 16) == bytes([0x3C]) * 16 + buf.close() + - # Build the prop the same shape modify_allocation does, so the helper's - # cuMemCreate / _build_access_descriptors path sees production-like input. +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings") +def test_vmm_unaligned_and_repeated_grow_close_cleanly(init_cuda): + """An unaligned request rounds up to the granularity (#2907), a grown buffer + grows again (#2908), and every alias closes without a warning.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + with assert_no_cuda_warning(): + buf = mr.allocate(1) + gran = buf.size + assert gran >= 1 + + second = mr.modify_allocation(buf, gran + 1) + assert second.size == 2 * gran + third = mr.modify_allocation(second, 3 * gran) + assert third.size == 3 * gran + + # A shorter alias asking for what its range already maps keeps the base. + alias = mr.modify_allocation(buf, 2 * gran) + assert alias.size >= 2 * gran + if int(second.handle) == int(buf.handle): + assert int(alias.handle) == int(buf.handle) + + # Every buffer maps the first chunk. + _vmm_fill(third, 0x77) + assert _vmm_read(alias, 0, gran) == bytes([0x77]) * gran + + third.close() + buf.close() + alias.close() + second.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_modify_allocation_rejects_foreign_or_closed_buffers(init_cuda): + """modify_allocation accepts only open buffers this resource returned.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + other_mr = _vmm_resource(device) + foreign = DummyDeviceMemoryResource(device).allocate(64) + sibling = other_mr.allocate(4096) + closed = mr.allocate(4096) + closed.close() + try: + with pytest.raises(TypeError): + mr.modify_allocation(foreign, 4096) + with pytest.raises(TypeError): + mr.modify_allocation(sibling, 8192) + with pytest.raises(RuntimeError, match="closed"): + mr.modify_allocation(closed, 8192) + finally: + foreign.close() + sibling.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_allocate_zero_size(init_cuda): + """allocate(0) returns an empty buffer without a driver call; it can be grown.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + buf = mr.allocate(0) + assert isinstance(buf, VirtualMemoryBuffer) + assert buf.size == 0 + assert int(buf.handle) == 0 + grown = mr.modify_allocation(buf, 4096) + assert grown.size >= 4096 + grown.close() + buf.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_allocate_without_access_descriptors(init_cuda): + """self_access=None with no peers creates and maps memory; no access call is made.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device, self_access=None) + buf = mr.allocate(4096) + assert buf.size >= 4096 + grown = mr.modify_allocation(buf, 2 * buf.size) + grown.close() + buf.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_close_on_capturing_stream_raises(init_cuda): + """An explicit close on a capturing stream raises and leaves the buffer open.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + buf = mr.allocate(4096) + gb = device.create_graph_builder().begin_building() + try: + with pytest.raises(RuntimeError, match="capturing stream"): + buf.close(gb) + finally: + gb.end_building() + gb.close() + assert int(buf.handle) != 0 + buf.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings") +@pytest.mark.parametrize("close_input_first", [True, False]) +def test_vmm_close_synchronizes_recorded_streams(init_cuda, close_input_first): + """Aliases with different deallocation streams close under queued work on both streams (#2886).""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + s1 = device.create_stream() + s2 = device.create_stream() + try: + buf = mr.allocate(8 * MIB, stream=s1) + grown = mr.modify_allocation(buf, 2 * buf.size) + grown.set_deallocation_stream(s2) + with assert_no_cuda_warning(): + for _ in range(8): + handle_return(driver.cuMemsetD8Async(int(buf.handle), 1, buf.size, s1.handle)) + handle_return(driver.cuMemsetD8Async(int(grown.handle), 2, grown.size, s2.handle)) + first, second = (buf, grown) if close_input_first else (grown, buf) + first.close() + second.close() + s1.sync() + s2.sync() + finally: + s1.close() + s2.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="cuMemGetInfo measures process-wide free memory") +def test_vmm_deallocate_frees_wrapped_pointer(init_cuda): + """deallocate() unmaps and frees one reservation wrapped with Buffer.from_handle.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) prop = driver.CUmemAllocationProp() prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE prop.location.id = device.device_id + prop.requestedHandleTypes = VirtualMemoryResourceOptions._handle_type_to_driver(VMM_HANDLE_TYPE) + size = handle_return( + driver.cuMemGetAllocationGranularity( + prop, driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED + ) + ) + desc = driver.CUmemAccessDesc() + desc.location.type = prop.location.type + desc.location.id = prop.location.id + desc.flags = driver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + + baseline = _vmm_free_memory() + handle = handle_return(driver.cuMemCreate(size, prop, 0)) + ptr = handle_return(driver.cuMemAddressReserve(size, 0, 0, 0)) + handle_return(driver.cuMemMap(ptr, size, 0, handle, 0)) + handle_return(driver.cuMemSetAccess(ptr, size, [desc], 1)) + handle_return(driver.cuMemRelease(handle)) # the mapping keeps the memory alive + + buf = Buffer.from_handle(ptr=int(ptr), size=size, mr=mr) + _vmm_fill(buf, 0x11) + buf.close() + assert baseline - _vmm_free_memory() < size - SUCCESS = driver.CUresult.CUDA_SUCCESS - NEW_HANDLE = 0xBEEF - calls = [] - - def fake_create(size, p, flags): - calls.append(("create", size)) - return (SUCCESS, NEW_HANDLE) - - def fake_map(ptr, size, offset, handle, flags): - calls.append(("map", ptr, size, handle)) - return (SUCCESS,) - - def fake_set_access(ptr, size, descs, count): - calls.append(("set_access", ptr, size, count)) - return (SUCCESS,) - - # Cleanup entry points. Release runs on commit; unmap and address free - # remain rollback-only. - def fake_unmap(ptr, size): - calls.append(("unmap", ptr, size)) - return (SUCCESS,) - - def fake_release(handle): - calls.append(("release", handle)) - return (SUCCESS,) - - def fake_addr_free(ptr, size): - calls.append(("addr_free", ptr, size)) - return (SUCCESS,) - - monkeypatch.setattr(driver, "cuMemCreate", fake_create) - monkeypatch.setattr(driver, "cuMemMap", fake_map) - monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) - monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) - monkeypatch.setattr(driver, "cuMemRelease", fake_release) - monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) - - # A real Buffer carries C++-owned handle state we must not fabricate; the - # helper only reads and writes buf._size, so a light stand-in suffices. - class FakeBuffer: - def __init__(self, size): - self._size = size - - original_size = 2 * 1024 * 1024 - aligned_additional = 2 * 1024 * 1024 - new_size = original_size + aligned_additional - new_ptr = 0x10_0000 # stand-in VA for the (mocked) contiguous extension - - buf = FakeBuffer(original_size) - result = vmm_mr._grow_allocation_fast_path(buf, new_size, prop, aligned_additional, new_ptr) - - # Fast-path contract: same buffer object, size updated in place. - assert result is buf - assert buf._size == new_size - - # Successful commit: create, map, set access, and release the creation handle. - assert [c[0] for c in calls] == ["create", "map", "set_access", "release"] - assert ("create", aligned_additional) in calls - assert ("map", new_ptr, aligned_additional, NEW_HANDLE) in calls - assert ("set_access", new_ptr, aligned_additional, 1) in calls - assert ("release", NEW_HANDLE) in calls + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_buffers_alive_at_shutdown_are_freed_quietly(init_cuda): + """Buffers still referenced at interpreter exit are freed without a warning or an error.""" + device = _vmm_device_or_skip() + code = textwrap.dedent( + f""" + import warnings + + from cuda.core import CUDAWarning, Device, VirtualMemoryResource, VirtualMemoryResourceOptions + + warnings.simplefilter("error", CUDAWarning) + dev = Device({device.device_id}) + dev.set_current() + mr = VirtualMemoryResource(dev, config=VirtualMemoryResourceOptions(handle_type={VMM_HANDLE_TYPE!r})) + keep = [mr.allocate(4096)] + keep.append(mr.modify_allocation(keep[0], 2 * keep[0].size)) + """ + ) + result = subprocess.run( # noqa: S603 + [sys.executable, "-c", code], + capture_output=True, + text=True, + timeout=CHILD_TIMEOUT_SEC, + check=False, + ) + assert result.returncode == 0, result.stderr + assert result.stderr == "" @pytest.mark.thread_unsafe(reason="cuMemGetInfo measures process-wide free memory") -@pytest.mark.parametrize("grow", [False, True], ids=["allocate", "grow"]) -def test_vmm_allocate_close_does_not_leak(init_cuda, grow): +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.parametrize("mode", ["allocate", "grow", "grow_moved"]) +def test_vmm_allocate_close_does_not_leak(init_cuda, mode): device = Device() if not device.properties.virtual_memory_management_supported: pytest.skip("Virtual memory management is not supported on this device") - mr = VirtualMemoryResource( - device, - config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), - ) - requested_size = 8 * 1024 * 1024 + mr = _vmm_resource(device) + requested_size = 8 * MIB def allocate_and_close(): - buf = mr.allocate(requested_size) - if grow: - buf = mr.modify_allocation(buf, 2 * buf.size) - aligned_size = buf.size - buf.close() - return aligned_size - - aligned_size = allocate_and_close() # Warm up and learn the aligned allocation size. - - baseline = handle_return(driver.cuMemGetInfo())[0] + if mode == "grow_moved": + # Place the buffer at the start of a free hole and take the range + # right after it, so the grow has to move. + hole_mr, buf = _vmm_allocate_in_free_hole(device, requested_size, 2 * requested_size) + decoy = _vmm_reserve_at(int(buf.handle) + buf.size, buf.size) + if decoy is None: + buf.close() + pytest.skip("the driver did not grant a reservation right after the buffer") + buffers = [buf] + try: + buffers.append(hole_mr.modify_allocation(buf, 2 * buf.size)) + finally: + handle_return(driver.cuMemAddressFree(decoy, buf.size)) + assert int(buffers[1].handle) != int(buf.handle) + else: + buf = mr.allocate(requested_size) + buffers = [buf] + if mode == "grow": + buffers.append(mr.modify_allocation(buf, 2 * buf.size)) + mapped_size = buffers[-1].size + for b in buffers: + b.close() + return mapped_size + + mapped_size = allocate_and_close() # Warm up and learn the aligned mapped size. + + baseline = _vmm_free_memory() for _ in range(8): allocate_and_close() - free = handle_return(driver.cuMemGetInfo())[0] + free = _vmm_free_memory() - # Current main leaks aligned_size per iteration; the fixed path stays near baseline. - assert baseline - free < aligned_size + # A leak would cost mapped_size per iteration; the fixed path stays near baseline. + assert baseline - free < mapped_size def test_vmm_allocator_rdma_unsupported_exception(): @@ -1561,6 +1819,21 @@ def test_vmm_allocator_rdma_unsupported_exception(): VirtualMemoryResource(device, config=options) +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings and mutates the context stack") +def test_vmm_host_location_allocate_without_current_context(init_cuda): + """A host-located resource allocates and closes with no CUDA context current.""" + device = Device() + if not device.properties.host_virtual_memory_management_supported: + pytest.skip("Host virtual memory management is not supported on this device") + mr = VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type="host", handle_type=None)) + with no_current_context(), assert_no_cuda_warning(): + buf = mr.allocate(4096) + assert buf.size >= 4096 + assert mr.is_host_accessible + buf.close() + + def test_device_memory_resource_with_options(init_cuda): device = Device() if not device.properties.memory_pools_supported: @@ -2456,21 +2729,27 @@ def test_vmm_options_handle_type_win32_raises(): VirtualMemoryResourceOptions._handle_type_to_driver("win32") -@pytest.mark.agent_authored(model="claude-opus-5") +@pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.parametrize("location_type", ["host", "host_numa", "host_numa_current"]) def test_vmm_host_location_types_report_host_accessible(location_type): - """Every host-backed location type reports is_host_accessible. + """Every host-backed location type reports is_host_accessible and needs handle_type=None. __init__ classifies "host", "host_numa" and "host_numa_current" alike when deciding the resource is not bound to a device, so is_host_accessible must agree; otherwise a NUMA-located resource claims to be neither host- nor - device-accessible. + device-accessible. The driver rejects an exportable handle type for host + memory, so the default handle type is rejected at construction. """ device = Device() device.set_current() - mr = VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type)) + with pytest.raises(ValueError, match="handle_type=None"): + VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type)) + mr = VirtualMemoryResource( + device, config=VirtualMemoryResourceOptions(location_type=location_type, handle_type=None) + ) assert mr.device is None assert mr.is_host_accessible is True + assert mr.device_id == -1 def test_device_memory_resource_peer_accessible_by_non_owned(mempool_device): From e7cbaccd414648c83d47e2ef06b75da1877b7386 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Fri, 18 Sep 2026 14:29:30 -0700 Subject: [PATCH 2/6] test(cuda.core): run the VMM shutdown test from an empty directory The child interpreter inherited pytest's working directory, cuda_core/, so `import cuda.core` resolved to the uncompiled source tree in CI and failed on `cuda.core._version`. Use the shared run_python_snippet helper, which starts the child in an empty temporary directory. Co-Authored-By: Claude Fable 5.1 --- cuda_core/tests/test_memory.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 25d1b6b7513..15ec6b19b57 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -3,7 +3,6 @@ import ctypes import multiprocessing as mp -import subprocess import sys import textwrap @@ -17,6 +16,7 @@ import re import pytest +from cuda_python_test_helpers.subprocess_runner import run_python_snippet from helpers import supports_ipc_mempool from helpers.buffers import ( DummyDeviceMemoryResource, @@ -1738,14 +1738,9 @@ def test_vmm_buffers_alive_at_shutdown_are_freed_quietly(init_cuda): keep.append(mr.modify_allocation(keep[0], 2 * keep[0].size)) """ ) - result = subprocess.run( # noqa: S603 - [sys.executable, "-c", code], - capture_output=True, - text=True, - timeout=CHILD_TIMEOUT_SEC, - check=False, - ) - assert result.returncode == 0, result.stderr + # The runner starts the child in an empty directory, so the installed + # package is imported rather than the source tree the test runs from. + result = run_python_snippet(code, timeout=CHILD_TIMEOUT_SEC) assert result.stderr == "" From f9e5940bb6211521b67c29a083105777a1da5dce Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Fri, 18 Sep 2026 15:36:52 -0700 Subject: [PATCH 3/6] fix(cuda.core): harden VirtualMemoryResource after review - Run the range deleter's stream sync in relaxed capture mode, so a capture on an unrelated stream is not invalidated. - Record a real stream on allocate(0) and inherit it on the grow. - Require handle_type=None for location "host" only. - Apply the constructor's option checks to a per-call modify_allocation config, including the RDMA support check. - Narrow the close() capture contract to non-default streams in the docstring, design doc and release note. - Tests: failed grow leaves the input intact, close during an unrelated capture, GC release during capture, deterministic stream sync with a sleep kernel, cuMemGetAccess on both chunks, graph retention across a grow, forced-move leak on 2 MiB that fails rather than skips. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 33 ++- cuda_core/cuda/core/_cpp/rt/driver_api.cpp | 1 + cuda_core/cuda/core/_cpp/rt/driver_api.hpp | 1 + .../cuda/core/_cpp/rt/virtual_memory.cpp | 12 + .../core/_memory/_virtual_memory_resource.pyi | 19 +- .../core/_memory/_virtual_memory_resource.pyx | 69 +++-- cuda_core/cuda/core/_rt.pyx | 4 +- cuda_core/docs/source/release/1.3.0-notes.rst | 10 +- .../graph/test_graph_definition_lifetime.py | 43 ++++ cuda_core/tests/test_memory.py | 242 +++++++++++++++--- 10 files changed, 364 insertions(+), 70 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index d7d1ee3b107..ae448ecb67d 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -30,7 +30,13 @@ its local handles die. The module is written in Cython, because the handles are Writes through one are visible through the other in stream order and at kernel boundaries; the driver adds no synchronization of its own. - **Context and synchronization.** No VMM entry point needs a current context. `cuMemUnmap` does - not synchronize. `cuStreamSynchronize` is rejected on a capturing stream. + not synchronize. `cuStreamSynchronize` is rejected on a capturing stream. It is also one of the + calls the driver treats as unsafe while any capture is active: in the calling thread's default + (global) capture mode it invalidates every global-mode capture in the process, and any + non-relaxed capture the thread began, before it returns an error. Switching the thread to + relaxed mode with `cuThreadExchangeStreamCaptureMode` disables that interaction; the stream's + own capture state is unaffected by the mode. The VMM entry points do not have this + interaction. So a mapping depends on exactly one reservation and one allocation, mappings never depend on other mappings, and reservations and allocations are independent of each other. A buffer is a @@ -97,8 +103,9 @@ Deleters: - `VmmRange`: release the GIL; unless the interpreter is finalizing, for each forwarded stream check the capture status and skip the stream with one report when a sync would disturb a capture (the stream is capturing, or it is the legacy stream while a blocking stream in its - context is capturing), otherwise synchronize it under its bound context. Then, whether or not the syncs succeeded, - destroy the mappings. Each mapping unmaps; the reservations free and the allocations release as + context is capturing), otherwise synchronize it under its bound context with the thread's + capture mode switched to relaxed for the call, so a capture on an unrelated stream is not + invalidated. Then, whether or not the syncs succeeded, destroy the mappings. Each mapping unmaps; the reservations free and the allocations release as their last references go. Every forwarded stream is synchronized because two aliases may have recorded different streams; synchronizing only the last one to die would unmap under work queued on the other. This is the first blocking deleter in the layer, and it may run inside @@ -109,14 +116,18 @@ what makes that safe: the allocation is released exactly once, when its last map ## The resource -- `VirtualMemoryResourceOptions` describes the allocations. `__init__` rejects `location_type="host"` with a - handle type other than `None`, which the driver rejects, and keeps the RDMA and VMM-support - checks. The resource reports `is_ipc_enabled = False`, which `Buffer.ipc_descriptor` reads. +- `VirtualMemoryResourceOptions` describes the allocations. `_check_config` rejects + `location_type="host"` with a handle type other than `None`, which the driver rejects, and a + request for GPUDirect RDMA on a device without support. `__init__` runs it after the VMM-support + check; `modify_allocation` runs it on a per-call configuration, which must also name the + resource's location. The resource reports `is_ipc_enabled = False`, which + `Buffer.ipc_descriptor` reads. - `cdef class VirtualMemoryBuffer(Buffer)` carries no extra state. It is created with `Buffer_from_deviceptr_handle(h_ptr, size, self, cls=VirtualMemoryBuffer)` and documented in - `api.rst` like `ManagedBuffer`. It overrides `close(stream=None)` to reject a capturing stream, - since VMM deallocation is synchronous and cannot be captured. `allocate(0)` returns one with no - mapping. + `api.rst` like `ManagedBuffer`. It overrides `close(stream=None)` to reject a capturing stream + other than a default stream, since VMM deallocation is synchronous and cannot be captured; a + default stream is checked by the range deleter under its bound context, which reports and skips + the sync instead of raising. `allocate(0)` returns one with no mapping. - `allocate(size, *, stream=None)`: 1. `size == 0` returns an empty buffer without a driver call, like the other resources. 2. Build `CUmemAllocationProp` and the access descriptors from the options; query the @@ -131,8 +142,8 @@ what makes that safe: the allocation is released exactly once, when its last map 6. Return a `VirtualMemoryBuffer` whose `size` is the aligned size. - `modify_allocation(buf, new_size, config=None)`: - `Buffer_check_open(buf)`; `range = vmm_range(buf._h_ptr)`; an empty range means the buffer - did not come from this resource: `TypeError`. `cfg = config or self.config` governs the new - chunk only and is not stored on the resource. Let `req = align_up(new_size)` and `total` be + did not come from this resource: `TypeError`. `cfg = config or self.config` passes + `_check_config`, governs the new chunk only and is not stored on the resource. Let `req = align_up(new_size)` and `total` be the range total. - `req <= buf.size`: return `buf`. The buffer already covers the request. - `buf.size < req <= total`: return a new `VirtualMemoryBuffer` over the same range with size diff --git a/cuda_core/cuda/core/_cpp/rt/driver_api.cpp b/cuda_core/cuda/core/_cpp/rt/driver_api.cpp index 3a98387dd36..9620bfd4586 100644 --- a/cuda_core/cuda/core/_cpp/rt/driver_api.cpp +++ b/cuda_core/cuda/core/_cpp/rt/driver_api.cpp @@ -70,6 +70,7 @@ decltype(&cuMemUnmap) p_cuMemUnmap = nullptr; decltype(&cuMemSetAccess) p_cuMemSetAccess = nullptr; decltype(&cuStreamSynchronize) p_cuStreamSynchronize = nullptr; decltype(&cuStreamGetCaptureInfo) p_cuStreamGetCaptureInfo = nullptr; +decltype(&cuThreadExchangeStreamCaptureMode) p_cuThreadExchangeStreamCaptureMode = nullptr; decltype(&cuLibraryLoadFromFile) p_cuLibraryLoadFromFile = nullptr; decltype(&cuLibraryLoadData) p_cuLibraryLoadData = nullptr; diff --git a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp index eec3a7fbd32..ee40418a4fc 100644 --- a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp @@ -73,6 +73,7 @@ extern decltype(&cuMemUnmap) p_cuMemUnmap; extern decltype(&cuMemSetAccess) p_cuMemSetAccess; extern decltype(&cuStreamSynchronize) p_cuStreamSynchronize; extern decltype(&cuStreamGetCaptureInfo) p_cuStreamGetCaptureInfo; +extern decltype(&cuThreadExchangeStreamCaptureMode) p_cuThreadExchangeStreamCaptureMode; // Library extern decltype(&cuLibraryLoadFromFile) p_cuLibraryLoadFromFile; diff --git a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp index 76be333cabd..340a2f62787 100644 --- a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp @@ -215,6 +215,13 @@ static bool sync_would_disturb_capture(CUstream stream) noexcept { // skip message that fits this use: when the sync cannot run, nothing leaks, // because the range is unmapped regardless. Sets `capture_skipped` instead of // synchronizing when the sync would disturb a capture. +// +// The sync itself runs with the calling thread in relaxed capture mode. +// cuStreamSynchronize is one of the calls the driver treats as unsafe while a +// capture is active: in the thread's default (global) mode it invalidates +// every global-mode capture in the process, and any non-relaxed capture this +// thread began, on streams unrelated to `s`. Relaxed mode disables that +// interaction; the stream's own capture state is still checked above. static void sync_recorded_stream(const DeallocationStream& ds, bool& capture_skipped) noexcept { const CUstream s = as_cu(ds.h_stream); CUcontext previous = nullptr; @@ -226,7 +233,12 @@ static void sync_recorded_stream(const DeallocationStream& ds, bool& capture_ski } else if (sync_would_disturb_capture(s)) { capture_skipped = true; } else { + CUstreamCaptureMode mode = CU_STREAM_CAPTURE_MODE_RELAXED; + const CUresult swapped = p_cuThreadExchangeStreamCaptureMode(&mode); // `mode` now holds the previous mode status = p_cuStreamSynchronize(s); + if (swapped == CUDA_SUCCESS) { + p_cuThreadExchangeStreamCaptureMode(&mode); // restore the previous mode + } } const CUresult restore = exit_context(previous, changed, CUDA_SUCCESS); if (restore != CUDA_SUCCESS) { diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index 992771794cd..38a1224a2b1 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -31,7 +31,7 @@ class VirtualMemoryResourceOptions: handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str Export handle type for the physical allocation. Use ``"posix_fd"`` on Linux if you plan to import/export the allocation. Use `None` if you - don't need an exportable handle. Host-located allocations require + don't need an exportable handle. ``location_type="host"`` requires `None`. gpu_direct_rdma: bool Hint that the allocation should be GDR-capable (if supported). @@ -97,8 +97,12 @@ class VirtualMemoryBuffer(Buffer): last buffer that maps them closes. Before it unmaps, the resource synchronizes every deallocation stream the buffers of the range recorded. Virtual memory deallocation is synchronous and cannot be - captured, so closing on a capturing stream raises and leaves the - buffer open. + captured. When the stream this close uses, given or recorded, is not + a default stream and is capturing, the call raises and leaves the + buffer open. A default stream is checked when the range is released + instead: if synchronizing it would disturb a capture in its context, + the release reports a :class:`CUDAWarning` and unmaps without + synchronizing that stream. Parameters ---------- @@ -183,7 +187,8 @@ class VirtualMemoryResource(MemoryResource): config : VirtualMemoryResourceOptions, optional Configuration for the new physical memory chunk only. Existing chunks keep the access they were created with, and the resource's - own configuration is unchanged. + own configuration is unchanged. It must name the resource's + ``location_type`` and passes the same checks as the constructor. Returns ------- @@ -195,6 +200,12 @@ class VirtualMemoryResource(MemoryResource): ------ TypeError If ``buf`` did not come from this resource. + ValueError + If ``config`` names a different location than the resource, or + the constructor would reject it. + RuntimeError + If ``buf`` is closed, or ``config`` requests GPUDirect RDMA on a + device without support. CUDAError If a driver call fails. ``buf`` is untouched when this method raises. """ diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 3ab62d43ec4..b1db96da35c 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -92,7 +92,7 @@ class VirtualMemoryResourceOptions: handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str Export handle type for the physical allocation. Use ``"posix_fd"`` on Linux if you plan to import/export the allocation. Use `None` if you - don't need an exportable handle. Host-located allocations require + don't need an exportable handle. ``location_type="host"`` requires `None`. gpu_direct_rdma: bool Hint that the allocation should be GDR-capable (if supported). @@ -235,8 +235,12 @@ cdef class VirtualMemoryBuffer(Buffer): last buffer that maps them closes. Before it unmaps, the resource synchronizes every deallocation stream the buffers of the range recorded. Virtual memory deallocation is synchronous and cannot be - captured, so closing on a capturing stream raises and leaves the - buffer open. + captured. When the stream this close uses, given or recorded, is not + a default stream and is capturing, the call raises and leaves the + buffer open. A default stream is checked when the range is released + instead: if synchronizing it would disturb a capture in its context, + the release reports a :class:`CUDAWarning` and unmaps without + synchronizing that stream. Parameters ---------- @@ -301,23 +305,31 @@ cdef class VirtualMemoryResource(MemoryResource): ) if self.config.location_type in _HOST_LOCATION_TYPES: self.device = None - # The driver rejects an exportable handle type for host memory. - if self.config.handle_type is not None: - raise ValueError( - "host-located virtual memory cannot have an exportable handle type; " - "pass handle_type=None" - ) - if self.device is not None and not self.device.properties.virtual_memory_management_supported: raise RuntimeError("VirtualMemoryResource requires CUDA VMM API support") + self._check_config(self.config) + + cdef int _check_config(self, object cfg) except -1: + """Reject options the driver would reject, before any driver call. - # Validate RDMA support if requested - if ( - self.config.gpu_direct_rdma - and self.device is not None - and not self.device.properties.gpu_direct_rdma_supported - ): + Shared by ``__init__`` and ``modify_allocation``, so a per-call + configuration is held to the same rules as the resource's own. + """ + if cfg.location_type != self.config.location_type: + raise ValueError( + f"config.location_type {str(cfg.location_type)!r} does not match the resource's " + f"{str(self.config.location_type)!r}; the location of a buffer cannot change" + ) + # The driver rejects an exportable handle type for HOST memory only; + # the NUMA location types may carry one. + if cfg.location_type == VirtualMemoryLocationType.HOST and cfg.handle_type is not None: + raise ValueError( + 'virtual memory with location_type="host" cannot have an exportable handle type; ' + "pass handle_type=None" + ) + if cfg.gpu_direct_rdma and self.device is not None and not self.device.properties.gpu_direct_rdma_supported: raise RuntimeError("GPU Direct RDMA is not supported on this device") + return 0 cdef int _fill_prop(self, object cfg, cydriver.CUmemAllocationProp* prop) except -1: # The location comes from the resource; the rest may come from a @@ -438,7 +450,13 @@ cdef class VirtualMemoryResource(MemoryResource): if size == 0: # Nothing to reserve or map; an empty buffer with a non-owning handle. - return Buffer_from_deviceptr_handle(deviceptr_create_ref(0), 0, self, None, VirtualMemoryBuffer) + # A real stream is still recorded so that a later grow inherits it; + # a default-stream token is not, which keeps this path free of + # driver calls. + h_ptr = deviceptr_create_ref(0) + if s is not None and not Stream_is_default_token(s): + HANDLE_RETURN(set_deallocation_stream(h_ptr, s._h_stream)) + return Buffer_from_deviceptr_handle(h_ptr, 0, self, None, VirtualMemoryBuffer) self._fill_prop(cfg, &prop) self._fill_access(cfg, &prop, descs) @@ -494,7 +512,8 @@ cdef class VirtualMemoryResource(MemoryResource): config : VirtualMemoryResourceOptions, optional Configuration for the new physical memory chunk only. Existing chunks keep the access they were created with, and the resource's - own configuration is unchanged. + own configuration is unchanged. It must name the resource's + ``location_type`` and passes the same checks as the constructor. Returns ------- @@ -506,6 +525,12 @@ cdef class VirtualMemoryResource(MemoryResource): ------ TypeError If ``buf`` did not come from this resource. + ValueError + If ``config`` names a different location than the resource, or + the constructor would reject it. + RuntimeError + If ``buf`` is closed, or ``config`` requests GPUDirect RDMA on a + device without support. CUDAError If a driver call fails. ``buf`` is untouched when this method raises. """ @@ -531,9 +556,13 @@ cdef class VirtualMemoryResource(MemoryResource): cfg = self.config if config is None else check_or_create_options( VirtualMemoryResourceOptions, config, "VirtualMemoryResource options", keep_none=False ) + self._check_config(cfg) if b._size == 0: - # An empty buffer maps nothing; the request is a fresh allocation. - return self._allocate(cfg, new_size, None) + # An empty buffer maps nothing; the request is a fresh allocation + # that inherits the stream the empty buffer recorded, if any. + new_buf = self._allocate(cfg, new_size, None) + self._copy_deallocation_stream((new_buf)._h_ptr, b._h_ptr) + return new_buf rng = vmm_range(b._h_ptr) if not rng: raise TypeError("buf was not allocated by VirtualMemoryResource.allocate") diff --git a/cuda_core/cuda/core/_rt.pyx b/cuda_core/cuda/core/_rt.pyx index 69fcf4f702b..70483960a9e 100644 --- a/cuda_core/cuda/core/_rt.pyx +++ b/cuda_core/cuda/core/_rt.pyx @@ -424,6 +424,7 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt": void* p_cuMemSetAccess "reinterpret_cast(cuda_core::rt::p_cuMemSetAccess)" void* p_cuStreamSynchronize "reinterpret_cast(cuda_core::rt::p_cuStreamSynchronize)" void* p_cuStreamGetCaptureInfo "reinterpret_cast(cuda_core::rt::p_cuStreamGetCaptureInfo)" + void* p_cuThreadExchangeStreamCaptureMode "reinterpret_cast(cuda_core::rt::p_cuThreadExchangeStreamCaptureMode)" # Library void* p_cuLibraryLoadFromFile "reinterpret_cast(cuda_core::rt::p_cuLibraryLoadFromFile)" @@ -508,7 +509,7 @@ cdef void _init_driver_fn_pointers() noexcept: global p_cuMemPoolImportPointer global p_cuMemCreate, p_cuMemRelease, p_cuMemAddressReserve, p_cuMemAddressFree global p_cuMemMap, p_cuMemUnmap, p_cuMemSetAccess - global p_cuStreamSynchronize, p_cuStreamGetCaptureInfo + global p_cuStreamSynchronize, p_cuStreamGetCaptureInfo, p_cuThreadExchangeStreamCaptureMode global p_cuLibraryLoadFromFile, p_cuLibraryLoadData, p_cuLibraryUnload, p_cuLibraryGetKernel global p_cuGraphDestroy, p_cuGraphInstantiateWithParams global p_cuGraphExecUpdate, p_cuGraphExecDestroy @@ -587,6 +588,7 @@ cdef void _init_driver_fn_pointers() noexcept: p_cuMemSetAccess = _get_driver_fn("cuMemSetAccess") p_cuStreamSynchronize = _get_driver_fn("cuStreamSynchronize") p_cuStreamGetCaptureInfo = _get_driver_fn("cuStreamGetCaptureInfo") + p_cuThreadExchangeStreamCaptureMode = _get_driver_fn("cuThreadExchangeStreamCaptureMode") # Library p_cuLibraryLoadFromFile = _get_driver_fn("cuLibraryLoadFromFile") diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index 069d8346595..afc18d60b40 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -33,7 +33,10 @@ New features - Added :class:`VirtualMemoryBuffer`, the :class:`Buffer` subclass that :class:`VirtualMemoryResource` returns. Closing it on a capturing stream - raises, because virtual memory deallocation cannot be captured. + other than a default stream raises, because virtual memory deallocation + cannot be captured. A release ordered on a default stream that would + disturb a capture in its context is reported as a :class:`CUDAWarning` + and proceeds without the synchronization. Fixes and enhancements ---------------------- @@ -163,8 +166,9 @@ Fixes and enhancements range. A request the buffer already covers returns the buffer itself, and :attr:`Buffer.size` after a grow is the aligned total. The ``config`` argument applies to the chunk that call adds and is no longer stored on the - resource (`#2909 `__). - Buffers from ``allocate()`` free themselves when they close and no longer + resource (`#2909 `__); + it must keep the resource's location and passes the constructor's option + checks. Buffers from ``allocate()`` free themselves when they close and no longer call :meth:`VirtualMemoryResource.deallocate`, which now serves pointers wrapped with :meth:`Buffer.from_handle`. A buffer records the stream passed to ``allocate()`` and the last close of an aliased range synchronizes every diff --git a/cuda_core/tests/graph/test_graph_definition_lifetime.py b/cuda_core/tests/graph/test_graph_definition_lifetime.py index a84e0d2c207..adf7470d232 100644 --- a/cuda_core/tests/graph/test_graph_definition_lifetime.py +++ b/cuda_core/tests/graph/test_graph_definition_lifetime.py @@ -958,6 +958,49 @@ def close_graph(graph_to_close): _wait_until(lambda: not allocation_weak) +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_memcpy_node_retains_vmm_range_across_grow(init_cuda): + """A memcpy node keeps a virtual memory range mapped across a grow and the close of every alias. + + The node attaches the input buffer's device pointer handle. Growing the + buffer creates an alias; closing both buffers leaves the node as the last + owner, so the range stays mapped until the graph that retains it is gone. + """ + from cuda.core._utils._weak_handles import weak_handle + + dev = Device() + mr = _virtual_memory_resource(dev) + buf = mr.allocate(8, stream=dev.default_stream) + dev.default_stream.sync() + dptr = int(buf.handle) + + graph_def = GraphDefinition() + copy_node = graph_def.memcpy(dptr, dptr + 4, 4, dst_owner=buf, src_owner=buf) + grown = mr.modify_allocation(buf, 2 * buf.size) + input_weak = weak_handle(buf) + grown_weak = weak_handle(grown) + + buf.close() + grown.close() + gc.collect() + # The node still owns the input's handle; the alias released its own. + assert input_weak + assert not grown_weak + + graph = graph_def.instantiate() + del copy_node, graph_def + gc.collect() + assert input_weak + + # The copy runs against the range the node kept mapped. + stream = dev.create_stream() + graph.launch(stream) + stream.sync() + graph.close() + del graph + _wait_until(lambda: not input_weak) + + @pytest.mark.agent_authored(model="gpt-5.6") def test_callback_survives_source_node_deletion_after_clone(init_cuda): """A clone independently retains a callback removed from its source graph.""" diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 15ec6b19b57..dae775de61c 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -2,9 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 import ctypes +import gc import multiprocessing as mp import sys import textwrap +import warnings from cuda.bindings import driver @@ -36,6 +38,7 @@ skip_if_managed_memory_unsupported, skip_if_pinned_memory_unsupported, ) +from helpers.nanosleep_kernel import NanosleepKernel from cuda.core import ( Buffer, @@ -1297,14 +1300,14 @@ def _vmm_reserve_at(addr, size): return ptr -def _vmm_allocate_in_free_hole(device, size, extra): +def _vmm_allocate_in_free_hole(device, size, extra, *, miss=pytest.skip): """Allocate ``size`` bytes at the start of a free address range ``extra`` bytes longer. Reserving and freeing a range finds free address space; a resource with that address as its hint then places the buffer there, so the range right after the buffer is known to be free. Sizes round up to the granularity. - Skips when the driver places the buffer elsewhere. Returns the resource - and the buffer. + Calls ``miss`` (``pytest.skip`` by default) when the driver places the + buffer elsewhere. Returns the resource and the buffer. """ probe = _vmm_resource(device).allocate(1) gran = probe.size @@ -1317,7 +1320,7 @@ def _vmm_allocate_in_free_hole(device, size, extra): buf = mr.allocate(size) if int(buf.handle) != int(hole): buf.close() - pytest.skip("the driver did not place the allocation at the requested address") + miss("the driver did not place the allocation at the requested address") return mr, buf @@ -1431,7 +1434,19 @@ def test_vmm_allocator_policy_configuration(): assert grown.size >= buffer.size + 8192 # The resource keeps its own configuration. assert vmm_mr.config == custom_config - # The chunk the input already mapped stays writable; the new chunk is readable. + # The chunk the input already mapped keeps read-write access; the chunk + # the grow added has the read-only access the per-call configuration asked for. + location = driver.CUmemLocation() + location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + location.id = device.device_id + access = driver.CUmemAccess_flags + assert ( + handle_return(driver.cuMemGetAccess(location, int(grown.handle))) == access.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + ) + assert ( + handle_return(driver.cuMemGetAccess(location, int(grown.handle) + buffer.size)) + == access.CU_MEM_ACCESS_FLAGS_PROT_READ + ) _vmm_fill(grown, 0x42, size=buffer.size) assert _vmm_read(grown, 0, 16) == bytes([0x42]) * 16 assert len(_vmm_read(grown, buffer.size, 16)) == 16 @@ -1615,19 +1630,72 @@ def test_vmm_modify_allocation_rejects_foreign_or_closed_buffers(init_cuda): sibling.close() +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_modify_allocation_validates_per_call_config(init_cuda): + """A per-call config passes the constructor's option checks and must keep the resource's location.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + buf = mr.allocate(4096) + ptr, size = int(buf.handle), buf.size + try: + with pytest.raises(ValueError, match="location_type"): + mr.modify_allocation( + buf, 2 * size, config=VirtualMemoryResourceOptions(location_type="host", handle_type=None) + ) + if not device.properties.gpu_direct_rdma_supported: + with pytest.raises(RuntimeError, match="GPU Direct RDMA"): + mr.modify_allocation( + buf, + 2 * size, + config=VirtualMemoryResourceOptions(handle_type=VMM_HANDLE_TYPE, gpu_direct_rdma=True), + ) + # The rejected requests leave the buffer untouched. + assert int(buf.handle) == ptr + assert buf.size == size + finally: + buf.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_host_modify_allocation_rejects_exportable_handle_type(init_cuda): + """On a host-located resource, a per-call config with an exportable handle type is rejected before any driver call.""" + device = Device() + if not device.properties.host_virtual_memory_management_supported: + pytest.skip("Host virtual memory management is not supported on this device") + mr = VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type="host", handle_type=None)) + buf = mr.allocate(4096) + try: + with pytest.raises(ValueError, match="handle_type=None"): + mr.modify_allocation(buf, 2 * buf.size, config=VirtualMemoryResourceOptions(location_type="host")) + finally: + buf.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") def test_vmm_allocate_zero_size(init_cuda): - """allocate(0) returns an empty buffer without a driver call; it can be grown.""" + """allocate(0) returns an empty buffer without a driver call; a grow of it inherits its stream.""" device = _vmm_device_or_skip() mr = _vmm_resource(device) - buf = mr.allocate(0) - assert isinstance(buf, VirtualMemoryBuffer) - assert buf.size == 0 - assert int(buf.handle) == 0 - grown = mr.modify_allocation(buf, 4096) - assert grown.size >= 4096 - grown.close() - buf.close() + s = device.create_stream() + try: + buf = mr.allocate(0, stream=s) + assert isinstance(buf, VirtualMemoryBuffer) + assert buf.size == 0 + assert int(buf.handle) == 0 + grown = mr.modify_allocation(buf, 4096) + assert grown.size >= 4096 + # The grown buffer's deallocation is ordered on the stream passed to + # allocate(0): its close waits for the work queued there. The sleep + # kernel does not touch the buffer, so a missing wait fails the event + # check instead of faulting. + NanosleepKernel(device, sleep_duration_ms=200).launch(s) + done = s.record() + grown.close() + assert done.is_done + buf.close() + finally: + s.sync() + s.close() @pytest.mark.agent_authored(model="claude-fable-5-1") @@ -1663,25 +1731,36 @@ def test_vmm_close_on_capturing_stream_raises(init_cuda): @pytest.mark.thread_unsafe(reason="records process-global warnings") @pytest.mark.parametrize("close_input_first", [True, False]) def test_vmm_close_synchronizes_recorded_streams(init_cuda, close_input_first): - """Aliases with different deallocation streams close under queued work on both streams (#2886).""" + """The closes that unmap wait for the work queued on every recorded stream (#2886). + + A sleep kernel holds each stream. Whether the grow shared the range (one + close synchronizes both streams) or moved it (each close synchronizes its + own stream), both streams must be idle once both buffers are closed. The + kernels do not touch the buffers, so a missing wait fails the event checks + instead of faulting. + """ device = _vmm_device_or_skip() mr = _vmm_resource(device) + sleeper = NanosleepKernel(device, sleep_duration_ms=200) s1 = device.create_stream() s2 = device.create_stream() try: buf = mr.allocate(8 * MIB, stream=s1) grown = mr.modify_allocation(buf, 2 * buf.size) grown.set_deallocation_stream(s2) + first, second = (buf, grown) if close_input_first else (grown, buf) with assert_no_cuda_warning(): - for _ in range(8): - handle_return(driver.cuMemsetD8Async(int(buf.handle), 1, buf.size, s1.handle)) - handle_return(driver.cuMemsetD8Async(int(grown.handle), 2, grown.size, s2.handle)) - first, second = (buf, grown) if close_input_first else (grown, buf) + sleeper.launch(s1) + done1 = s1.record() + sleeper.launch(s2) + done2 = s2.record() first.close() second.close() + assert done1.is_done + assert done2.is_done + finally: s1.sync() s2.sync() - finally: s1.close() s2.close() @@ -1744,6 +1823,99 @@ def test_vmm_buffers_alive_at_shutdown_are_freed_quietly(init_cuda): assert result.stderr == "" +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings and cuMemGetInfo measures process-wide free memory") +def test_vmm_failed_grow_leaves_input_intact(init_cuda): + """A grow the driver rejects raises, unwinds what it created, and leaves the input untouched (#2345).""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + buf = mr.allocate(4096) + _vmm_fill(buf, 0x5C) + ptr, size = int(buf.handle), buf.size + baseline = _vmm_free_memory() + # No address space of this size exists: the adjacent probe fails, and so + # does the reservation for the move. Nothing may be freed by hand. + with assert_no_cuda_warning(), pytest.raises(CUDAError): + mr.modify_allocation(buf, 1 << 62) + assert int(buf.handle) == ptr + assert buf.size == size + assert _vmm_read(buf, 0, 16) == bytes([0x5C]) * 16 + assert abs(baseline - _vmm_free_memory()) < size + buf.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings") +@pytest.mark.parametrize("mode", ["global", "thread_local"]) +def test_vmm_close_during_unrelated_capture_keeps_capture_valid(init_cuda, mode): + """Closing a buffer synchronizes its stream without invalidating a capture on another stream. + + The driver treats cuStreamSynchronize as unsafe while a global-mode capture + is active anywhere, or a thread-local capture is active on this thread; the + range deleter runs the sync in relaxed mode so the unrelated capture + survives and end_building succeeds. + """ + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + s = device.create_stream() + buf = mr.allocate(4096, stream=s) + gb = device.create_graph_builder().begin_building(mode=mode) + try: + with assert_no_cuda_warning(): + buf.close() + finally: + # Raises if the close invalidated the capture. + gb.end_building() + gb.close() + s.close() + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings and cuMemGetInfo measures process-wide free memory") +def test_vmm_release_from_gc_on_capturing_stream_reports_and_unmaps(init_cuda): + """A release from garbage collection while the recorded stream is capturing warns once and unmaps. + + Only an explicit close can raise; garbage collection skips the sync, + reports it, unmaps the range, and leaves the capture intact. + """ + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + baseline = _vmm_free_memory() + buf = mr.allocate(4096) + size = buf.size + gb = device.create_graph_builder().begin_building() + try: + buf.set_deallocation_stream(gb) + with warnings.catch_warnings(record=True) as records: + warnings.simplefilter("always", CUDAWarning) + del buf + gc.collect() + reports = [str(r.message) for r in records if issubclass(r.category, CUDAWarning)] + assert len(reports) == 1 and "capturing" in reports[0], reports + finally: + gb.end_building() + gb.close() + assert baseline - _vmm_free_memory() < size + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_close_without_stream_argument_on_capturing_stream_raises(init_cuda): + """close() with no argument checks the recorded deallocation stream and refuses to close during its capture.""" + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + buf = mr.allocate(4096) + gb = device.create_graph_builder().begin_building() + try: + buf.set_deallocation_stream(gb) + with pytest.raises(RuntimeError, match="capturing stream"): + buf.close() + finally: + gb.end_building() + gb.close() + assert int(buf.handle) != 0 + buf.close() + + @pytest.mark.thread_unsafe(reason="cuMemGetInfo measures process-wide free memory") @pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.parametrize("mode", ["allocate", "grow", "grow_moved"]) @@ -1755,15 +1927,19 @@ def test_vmm_allocate_close_does_not_leak(init_cuda, mode): mr = _vmm_resource(device) requested_size = 8 * MIB - def allocate_and_close(): + def allocate_and_close(warm_up=False): if mode == "grow_moved": - # Place the buffer at the start of a free hole and take the range - # right after it, so the grow has to move. - hole_mr, buf = _vmm_allocate_in_free_hole(device, requested_size, 2 * requested_size) + # Place a 2 MiB buffer at the start of a free hole and take the + # range right after it, so the grow has to move. The driver honors + # this placement for 2 MiB where it declined 8 MiB on Linux. The + # warm-up proves placement works here; a later miss would hide a + # leaked reservation behind a skip, so it fails instead. + miss = pytest.skip if warm_up else pytest.fail + hole_mr, buf = _vmm_allocate_in_free_hole(device, 2 * MIB, 4 * MIB, miss=miss) decoy = _vmm_reserve_at(int(buf.handle) + buf.size, buf.size) if decoy is None: buf.close() - pytest.skip("the driver did not grant a reservation right after the buffer") + miss("the driver did not grant a reservation right after the buffer") buffers = [buf] try: buffers.append(hole_mr.modify_allocation(buf, 2 * buf.size)) @@ -1780,14 +1956,14 @@ def allocate_and_close(): b.close() return mapped_size - mapped_size = allocate_and_close() # Warm up and learn the aligned mapped size. + mapped_size = allocate_and_close(warm_up=True) # Warm up and learn the aligned mapped size. baseline = _vmm_free_memory() for _ in range(8): allocate_and_close() free = _vmm_free_memory() - # A leak would cost mapped_size per iteration; the fixed path stays near baseline. + # A leak would cost at least one chunk per iteration; the fixed path stays near baseline. assert baseline - free < mapped_size @@ -2727,17 +2903,21 @@ def test_vmm_options_handle_type_win32_raises(): @pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.parametrize("location_type", ["host", "host_numa", "host_numa_current"]) def test_vmm_host_location_types_report_host_accessible(location_type): - """Every host-backed location type reports is_host_accessible and needs handle_type=None. + """Every host-backed location type reports is_host_accessible; only "host" needs handle_type=None. __init__ classifies "host", "host_numa" and "host_numa_current" alike when deciding the resource is not bound to a device, so is_host_accessible must agree; otherwise a NUMA-located resource claims to be neither host- nor - device-accessible. The driver rejects an exportable handle type for host - memory, so the default handle type is rejected at construction. + device-accessible. The driver rejects an exportable handle type for HOST + memory only, so the default handle type is rejected at construction for + "host" and accepted for the NUMA location types. """ device = Device() device.set_current() - with pytest.raises(ValueError, match="handle_type=None"): + if location_type == "host": + with pytest.raises(ValueError, match="handle_type=None"): + VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type)) + else: VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type)) mr = VirtualMemoryResource( device, config=VirtualMemoryResourceOptions(location_type=location_type, handle_type=None) From d5b79770cdd171d5f0bfd261c976e28b1630b913 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Fri, 18 Sep 2026 16:04:13 -0700 Subject: [PATCH 4/6] test(cuda.core): pass the handle-type enum to the CUDA 13.0 bindings The struct setter in cuda-bindings 13.0 accepts only the enum, and the Cython helper returns a plain int. Co-Authored-By: Claude Fable 5.1 --- cuda_core/tests/test_memory.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index dae775de61c..093be73e5d7 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1775,7 +1775,11 @@ def test_vmm_deallocate_frees_wrapped_pointer(init_cuda): prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE prop.location.id = device.device_id - prop.requestedHandleTypes = VirtualMemoryResourceOptions._handle_type_to_driver(VMM_HANDLE_TYPE) + # The helper returns a plain int; the CUDA 13.0 bindings' struct setter + # accepts only the enum. + prop.requestedHandleTypes = driver.CUmemAllocationHandleType( + VirtualMemoryResourceOptions._handle_type_to_driver(VMM_HANDLE_TYPE) + ) size = handle_return( driver.cuMemGetAllocationGranularity( prop, driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED From 5f07a73ca9a21890aedbdeb8fc6cc8396de1817b Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 21 Sep 2026 15:52:38 -0700 Subject: [PATCH 5/6] test(cuda.core): cover host VMM grow, default-stream capture skip, and size overflow Add tests for a host-located grow that moves, for a release ordered on the legacy default stream while a blocking stream in its context is capturing, and for a size whose rounding to the granularity does not fit in size_t. The forced-move test now asserts that neither the grow nor the closes warn (#2877). _align_up raises OverflowError instead of wrapping. The docstrings and the release note say that config has no effect when the buffer already covers the request and never changes the access of mapped memory, and that a host-located resource records no default stream. VMM_DESIGN.md describes the close() override. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 7 +- .../core/_memory/_virtual_memory_resource.pyi | 12 +- .../core/_memory/_virtual_memory_resource.pyx | 21 ++- cuda_core/docs/source/release/1.3.0-notes.rst | 4 +- cuda_core/tests/test_memory.py | 136 ++++++++++++++++-- 5 files changed, 159 insertions(+), 21 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index ae448ecb67d..e46c0562858 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -89,7 +89,9 @@ VmmRangeHandle vmm_range(const DevicePtrHandle& h); // empty for a non-VMM or There is exactly one ownership chain: `Buffer._h_ptr` -> `DevicePtrBox` (holds the range) -> `VmmRange` -> mappings -> reservations and allocations. The Cython buffer keeps no other reference; grow operations call `Buffer_check_open` and then `vmm_range(buf._h_ptr)`. -`Buffer.close()` stays `_h_ptr.reset()`. A buffer's `size` is always a prefix of its range. +`VirtualMemoryBuffer.close()` first refuses a capturing stream other than a default stream +(see "The resource") and then resets `_h_ptr`, as `Buffer.close()` does; the release itself is +the deleter chain below. A buffer's `size` is always a prefix of its range. The `DevicePtrHandle` must own the memory because graph memcpy nodes retain `buf._h_ptr` as an opaque owner; a non-owning handle would let a launched graph outlive its buffer. Any number of @@ -145,7 +147,8 @@ what makes that safe: the allocation is released exactly once, when its last map did not come from this resource: `TypeError`. `cfg = config or self.config` passes `_check_config`, governs the new chunk only and is not stored on the resource. Let `req = align_up(new_size)` and `total` be the range total. - - `req <= buf.size`: return `buf`. The buffer already covers the request. + - `req <= buf.size`: return `buf`. The buffer already covers the request; `cfg` is not applied, + and the access of memory that is already mapped never changes. - `buf.size < req <= total`: return a new `VirtualMemoryBuffer` over the same range with size `req`; no driver call. This serves a shorter alias asking for what the range already maps. - `req > total`, in place: probe `cuMemAddressReserve(req - total, align=0, hint=base+total)`. diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index 38a1224a2b1..cced7cf0e56 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -152,7 +152,9 @@ class VirtualMemoryResource(MemoryResource): Keyword-only. The allocation itself is synchronous. A real stream is recorded as the buffer's deallocation stream and synchronized when the buffer closes; with `None` or a default-stream token the legacy - default stream of the resource's device is recorded instead. + default stream of the resource's device is recorded instead. A + host-located resource records no default stream: its buffers close + without a synchronization unless a real stream was given. Returns ------- @@ -164,6 +166,8 @@ class VirtualMemoryResource(MemoryResource): CUDAError If any CUDA driver API call fails during allocation. Nothing is left allocated when this method raises. + OverflowError + If ``size`` rounded up to the granularity does not fit in ``size_t``. """ def modify_allocation(self, buf: Buffer, new_size: int, config: VirtualMemoryResourceOptions | None=None) -> VirtualMemoryBuffer: """ @@ -189,6 +193,9 @@ class VirtualMemoryResource(MemoryResource): chunks keep the access they were created with, and the resource's own configuration is unchanged. It must name the resource's ``location_type`` and passes the same checks as the constructor. + When ``buf`` already covers ``new_size`` there is no new chunk, so + ``config`` has no effect. This method never changes the access of + memory that is already mapped. Returns ------- @@ -206,6 +213,9 @@ class VirtualMemoryResource(MemoryResource): RuntimeError If ``buf`` is closed, or ``config`` requests GPUDirect RDMA on a device without support. + OverflowError + If ``new_size`` rounded up to the granularity does not fit in + ``size_t``. CUDAError If a driver call fails. ``buf`` is untouched when this method raises. """ diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index b1db96da35c..9cab016348b 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -190,7 +190,14 @@ class VirtualMemoryResourceOptions: return granularity # type: ignore[no-any-return] -cdef inline size_t _align_up(size_t size, size_t gran) noexcept nogil: +cdef inline size_t _align_up(size_t size, size_t gran) except? 0: + """Round ``size`` up to a multiple of ``gran``. + + Raises ``OverflowError`` instead of wrapping when the rounded size does not + fit in ``size_t``. + """ + if size > -1 - (gran - 1): + raise OverflowError(f"size {size} rounded up to the {gran}-byte granularity does not fit in size_t") return (size + gran - 1) // gran * gran @@ -418,7 +425,9 @@ cdef class VirtualMemoryResource(MemoryResource): Keyword-only. The allocation itself is synchronous. A real stream is recorded as the buffer's deallocation stream and synchronized when the buffer closes; with `None` or a default-stream token the legacy - default stream of the resource's device is recorded instead. + default stream of the resource's device is recorded instead. A + host-located resource records no default stream: its buffers close + without a synchronization unless a real stream was given. Returns ------- @@ -430,6 +439,8 @@ cdef class VirtualMemoryResource(MemoryResource): CUDAError If any CUDA driver API call fails during allocation. Nothing is left allocated when this method raises. + OverflowError + If ``size`` rounded up to the granularity does not fit in ``size_t``. """ cdef Stream s = None if stream is not None: @@ -514,6 +525,9 @@ cdef class VirtualMemoryResource(MemoryResource): chunks keep the access they were created with, and the resource's own configuration is unchanged. It must name the resource's ``location_type`` and passes the same checks as the constructor. + When ``buf`` already covers ``new_size`` there is no new chunk, so + ``config`` has no effect. This method never changes the access of + memory that is already mapped. Returns ------- @@ -531,6 +545,9 @@ cdef class VirtualMemoryResource(MemoryResource): RuntimeError If ``buf`` is closed, or ``config`` requests GPUDirect RDMA on a device without support. + OverflowError + If ``new_size`` rounded up to the granularity does not fit in + ``size_t``. CUDAError If a driver call fails. ``buf`` is untouched when this method raises. """ diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index afc18d60b40..cf0c1c422d6 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -168,7 +168,9 @@ Fixes and enhancements argument applies to the chunk that call adds and is no longer stored on the resource (`#2909 `__); it must keep the resource's location and passes the constructor's option - checks. Buffers from ``allocate()`` free themselves when they close and no longer + checks. It never changes the access of memory that is already mapped, so a + request the buffer already covers has no effect even with ``config``. + Buffers from ``allocate()`` free themselves when they close and no longer call :meth:`VirtualMemoryResource.deallocate`, which now serves pointers wrapped with :meth:`Buffer.from_handle`. A buffer records the stream passed to ``allocate()`` and the last close of an aliased range synchronizes every diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 093be73e5d7..7bd0f673d4d 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -54,6 +54,7 @@ MemoryResource, PinnedMemoryResource, PinnedMemoryResourceOptions, + StreamOptions, VirtualMemoryBuffer, VirtualMemoryResource, VirtualMemoryResourceOptions, @@ -1546,8 +1547,12 @@ def test_vmm_grow_preserves_contents_and_aliases_input(init_cuda): @pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings") def test_vmm_grow_moves_when_adjacent_range_is_taken(init_cuda): - """When the address range after the buffer is taken, the grow maps the existing memory at a new address.""" + """When the address range after the buffer is taken, the grow maps the existing memory at a new address. + + Neither the grow nor the closes emit a warning (#2877). + """ device = _vmm_device_or_skip() mr, buf = _vmm_allocate_in_free_hole(device, 2 * MIB, 4 * MIB) size = buf.size @@ -1557,22 +1562,23 @@ def test_vmm_grow_moves_when_adjacent_range_is_taken(init_cuda): if decoy is None: buf.close() pytest.skip("the driver did not grant a reservation right after the buffer") - try: - grown = mr.modify_allocation(buf, 2 * size) - finally: - handle_return(driver.cuMemAddressFree(decoy, size)) + with assert_no_cuda_warning(): + try: + grown = mr.modify_allocation(buf, 2 * size) + finally: + handle_return(driver.cuMemAddressFree(decoy, size)) - # The range after the buffer was taken, so the grow had to move. - assert int(grown.handle) != int(buf.handle) - assert grown.size == 2 * size - assert _vmm_read(grown, 0, size) == bytes([0x5A]) * size - _vmm_fill(buf, 0x3C) - assert _vmm_read(grown, 0, size) == bytes([0x3C]) * size + # The range after the buffer was taken, so the grow had to move. + assert int(grown.handle) != int(buf.handle) + assert grown.size == 2 * size + assert _vmm_read(grown, 0, size) == bytes([0x5A]) * size + _vmm_fill(buf, 0x3C) + assert _vmm_read(grown, 0, size) == bytes([0x3C]) * size - # Closing the result first leaves the input usable. - grown.close() - assert _vmm_read(buf, 0, 16) == bytes([0x3C]) * 16 - buf.close() + # Closing the result first leaves the input usable. + grown.close() + assert _vmm_read(buf, 0, 16) == bytes([0x3C]) * 16 + buf.close() @pytest.mark.agent_authored(model="claude-fable-5-1") @@ -1848,6 +1854,28 @@ def test_vmm_failed_grow_leaves_input_intact(init_cuda): buf.close() +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_vmm_size_rounding_overflow_raises(init_cuda): + """A size whose rounding to the granularity does not fit in size_t raises instead of wrapping. + + Without the check the rounded size wraps to zero: ``allocate`` would ask the + driver for zero bytes, and ``modify_allocation`` would treat the request as + already covered and return the input buffer. + """ + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + too_large = 2**64 - 1 + with pytest.raises(OverflowError): + mr.allocate(too_large) + buf = mr.allocate(4096) + try: + with pytest.raises(OverflowError): + mr.modify_allocation(buf, too_large) + assert buf.size >= 4096 + finally: + buf.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.thread_unsafe(reason="records process-global warnings") @pytest.mark.parametrize("mode", ["global", "thread_local"]) @@ -1902,6 +1930,40 @@ def test_vmm_release_from_gc_on_capturing_stream_reports_and_unmaps(init_cuda): assert baseline - _vmm_free_memory() < size +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings and cuMemGetInfo measures process-wide free memory") +def test_vmm_release_on_default_stream_during_blocking_capture_reports_and_unmaps(init_cuda): + """A release ordered on the default stream while a blocking stream in its context is capturing warns once and unmaps. + + ``allocate()`` without a stream records the legacy default stream. While a + blocking stream in the same context is capturing, the driver reports any + legacy-stream operation as an implicit capture dependency, and a + synchronization would invalidate the capture. The release skips the + synchronization, reports it, unmaps the range, and leaves the capture + intact. Streams from ``Device.create_stream()`` are non-blocking by default + and never interact with the legacy stream, so this needs a blocking one. + """ + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + baseline = _vmm_free_memory() + buf = mr.allocate(4096) + size = buf.size + blocking = device.create_stream(options=StreamOptions(nonblocking=False)) + gb = blocking.create_graph_builder().begin_building() + try: + with warnings.catch_warnings(record=True) as records: + warnings.simplefilter("always", CUDAWarning) + buf.close() + reports = [str(r.message) for r in records if issubclass(r.category, CUDAWarning)] + assert len(reports) == 1 and "legacy stream" in reports[0], reports + finally: + # Raises if the release invalidated the capture. + gb.end_building() + gb.close() + blocking.close() + assert baseline - _vmm_free_memory() < size + + @pytest.mark.agent_authored(model="claude-fable-5-1") def test_vmm_close_without_stream_argument_on_capturing_stream_raises(init_cuda): """close() with no argument checks the recorded deallocation stream and refuses to close during its capture.""" @@ -2009,6 +2071,50 @@ def test_vmm_host_location_allocate_without_current_context(init_cuda): buf.close() +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="records process-global warnings") +def test_vmm_host_location_grow_and_close(init_cuda): + """A host-located resource grows a buffer twice and every alias closes without a warning. + + The first grow extends the range in place or moves it. The second grow is + forced to move when the driver grants a decoy reservation right after the + buffer, so the existing host allocations are remapped at a new address. + Host-located buffers record no default stream, so the closes unmap without + a synchronization. + """ + device = Device() + if not device.properties.host_virtual_memory_management_supported: + pytest.skip("Host virtual memory management is not supported on this device") + mr = VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type="host", handle_type=None)) + with assert_no_cuda_warning(): + buf = mr.allocate(4096) + size = buf.size + grown = mr.modify_allocation(buf, 2 * size) + assert isinstance(grown, VirtualMemoryBuffer) + assert grown.size == 2 * size + assert grown.is_host_accessible and not grown.is_device_accessible + assert grown.device_id == -1 + # The input stays open at its old size. + assert buf.size == size + assert int(buf.handle) != 0 + + decoy = _vmm_reserve_at(int(grown.handle) + grown.size, size) + try: + moved = mr.modify_allocation(grown, 3 * size) + finally: + if decoy is not None: + handle_return(driver.cuMemAddressFree(decoy, size)) + assert moved.size == 3 * size + if decoy is not None: + # The range after the buffer was taken, so the grow had to move. + assert int(moved.handle) != int(grown.handle) + + buf.close() + grown.close() + moved.close() + assert moved.size == 0 + + def test_device_memory_resource_with_options(init_cuda): device = Device() if not device.properties.memory_pools_supported: From 1f9c014b4c1f2d35fa404fd5bde2326b1ee96111 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 21 Sep 2026 16:37:08 -0700 Subject: [PATCH 6/6] docs(cuda.core): state the VirtualMemoryResource invariants in VMM_DESIGN.md List the fifteen properties the handle-based VirtualMemoryResource maintains: once-only and ordered release of reservations and allocations, what a failed or successful grow leaves behind, stream ordering and graph capture, ownership by graph nodes and aliases, per-chunk access, range layout and rounding, the base-address registry, the deallocate() contract, context independence, and interpreter shutdown. The wording names no mechanism, so the list stays valid if the release is made stream-ordered. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 38 +++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index e46c0562858..d26a4b3c88c 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -201,3 +201,41 @@ pointer, was rejected: the driver needs no context for it, so nothing leaks. - Empty handles always carry a status; the in-place probe drains its status before falling back. - Factories that allocate are declared `except+` in `_rt.pxd`; deleters only destroy vectors. + +## Invariants + +Reservations and allocations are shared. After a move, one reservation holds every remapped +chunk, and the old and new ranges share their allocations. The invariants below hold under that +sharing. The rule that no deleter holds a C++ lock while it calls CUDA or Python is a property of +the whole layer; see [DESIGN.md](DESIGN.md). + +1. A reservation is freed exactly once, with its original pointer and size. +2. A reservation is freed only after every mapping inside it is unmapped. +3. An allocation is released exactly once. +4. An allocation is released only after its last mapping in any range is unmapped. +5. A failed grow leaves the input unchanged and leaks nothing. Pointer, size, contents, access, + and free memory are as before the call. +6. A grow leaves the input open. The input and the result see the same memory, whether the range + grew in place or moved. +7. Every stream recorded by any owner of a range finishes before the range is unmapped, except as + invariant 8 states. +8. A release never invalidates a graph capture. An explicit close on a capturing non-default + stream raises. A release from garbage collection, or one ordered on a default stream that + would disturb a capture, proceeds without ordering on that stream, warns once, and unmaps. +9. Graph nodes and aliases keep the range alive. The range dies with its last owner, in any close + order. +10. A chunk's access is fixed when the chunk is created and travels with its allocation. Every + mapping of the chunk applies the same descriptors. A grow's `config` governs only the new + chunk and never changes mapped memory. +11. The mappings of a range are contiguous and ascending, and the range total is the sum of their + sizes. A buffer's size is a multiple of the granularity and a prefix of its range. A size that + cannot be rounded raises. +12. A range is findable by its base address only while it is alive. The registry entry is removed + before the reservation is freed. +13. Buffers from `allocate()` free themselves and never call `deallocate()`. `deallocate()` serves + only pointers wrapped with `Buffer.from_handle`. +14. No operation needs a current context. A default-stream token is bound to the resource's + device context when it is recorded, and the release runs under that context. A host-located + resource records no default stream, so its release is ordered only on a stream the caller + passes. +15. Buffers alive at interpreter shutdown are freed without a warning or an error.