From 2b4c46f8ab5c314d254bdc842228b87ba3bc7a49 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Fri, 18 Sep 2026 12:40:37 -0700 Subject: [PATCH 01/14] 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 02/14] 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 03/14] 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 04/14] 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 05/14] 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 06/14] 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. From 7e088222a24fe79121aa7d96caef12853ad5921a Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 28 Sep 2026 13:50:37 -0700 Subject: [PATCH 07/14] refactor(cuda.core): make VMM resource attributes read-only and share the default-token check `VirtualMemoryResource.device` and `.config` are now `cdef readonly`; no other resource in the layer exposes writable attributes. The raw-handle default-stream check moves to the stream module as `Stream_handle_is_default_token`, and `Stream_is_default_token` delegates to it, so the two modules agree on one definition. `VirtualMemoryBuffer. close()` treats an empty handle as "no stream recorded" explicitly instead of folding it into the token check. Co-Authored-By: Claude Fable 5.1 --- .../core/_memory/_virtual_memory_resource.pyx | 19 ++++++------------- cuda_core/cuda/core/_stream.pxd | 2 ++ cuda_core/cuda/core/_stream.pyx | 11 ++++++++--- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 9cab016348b..007cd93e00a 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -4,7 +4,6 @@ from __future__ import annotations -from libc.stdint cimport uintptr_t from libc.string cimport memset from libcpp.vector cimport vector @@ -44,7 +43,7 @@ from cuda.core._rt cimport ( vmm_range_reserve, vmm_range_total, ) -from cuda.core._stream cimport Stream, Stream_accept, Stream_is_default_token +from cuda.core._stream cimport Stream, Stream_accept, Stream_handle_is_default_token, Stream_is_default_token from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, check_or_create_options from dataclasses import dataclass, field @@ -201,11 +200,6 @@ cdef inline size_t _align_up(size_t size, size_t gran) except? 0: 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: @@ -257,18 +251,17 @@ cdef class VirtualMemoryBuffer(Buffer): """ 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) + h = s._h_stream else: + # Empty when no stream was recorded (host-located memory). 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): + if h and not Stream_handle_is_default_token(as_cu(h)) and _stream_is_capturing(as_cu(h)): raise RuntimeError( "cannot close a VirtualMemoryResource buffer on a capturing stream: " "virtual memory deallocation is synchronous and cannot be captured" @@ -302,8 +295,8 @@ cdef class VirtualMemoryResource(MemoryResource): """ cdef: - public object device - public object config + readonly object device + readonly object config def __init__(self, device_id: Device | int, config: VirtualMemoryResourceOptions | None = None) -> None: self.device = Device(device_id) diff --git a/cuda_core/cuda/core/_stream.pxd b/cuda_core/cuda/core/_stream.pxd index 9ae01151785..1b1c9837eb4 100644 --- a/cuda_core/cuda/core/_stream.pxd +++ b/cuda_core/cuda/core/_stream.pxd @@ -2,6 +2,7 @@ # # SPDX-License-Identifier: Apache-2.0 +from cuda.bindings cimport cydriver from cuda.core._rt cimport ContextHandle, StreamHandle @@ -27,5 +28,6 @@ cdef inline int Stream_check_open(Stream self) except -1: if not self._h_stream: raise RuntimeError("Stream has been closed") return 0 +cdef bint Stream_handle_is_default_token(cydriver.CUstream s) noexcept nogil cdef bint Stream_is_default_token(Stream self) noexcept nogil cdef bint Stream_is_legacy_default_token(Stream self) noexcept nogil diff --git a/cuda_core/cuda/core/_stream.pyx b/cuda_core/cuda/core/_stream.pyx index 4c51b32f4f0..32872327835 100644 --- a/cuda_core/cuda/core/_stream.pyx +++ b/cuda_core/cuda/core/_stream.pyx @@ -517,16 +517,21 @@ cpdef Stream default_stream(): else: return LEGACY_DEFAULT_STREAM -cdef inline bint Stream_is_default_token(Stream self) noexcept nogil: - """Return True for CU_STREAM_LEGACY and CU_STREAM_PER_THREAD. +cdef inline bint Stream_handle_is_default_token(cydriver.CUstream s) noexcept nogil: + """Return True for the raw CU_STREAM_LEGACY and CU_STREAM_PER_THREAD tokens. These tokens carry no context of their own; they refer to whatever context is current, so nothing resolved from one may be cached on the object. """ - cdef uintptr_t h = as_cu(self._h_stream) + cdef uintptr_t h = s return h == cydriver.CU_STREAM_LEGACY or h == cydriver.CU_STREAM_PER_THREAD +cdef inline bint Stream_is_default_token(Stream self) noexcept nogil: + """Return True when this stream is CU_STREAM_LEGACY or CU_STREAM_PER_THREAD.""" + return Stream_handle_is_default_token(as_cu(self._h_stream)) + + cdef inline bint Stream_is_legacy_default_token(Stream self) noexcept nogil: """Return True only for CU_STREAM_LEGACY. From 343a4952d11d5c43cff09386823189947de62b97 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 28 Sep 2026 13:50:37 -0700 Subject: [PATCH 08/14] test(cuda.core): check VMM release per address instead of device-wide free memory The counter from cuMemGetInfo covers the whole device, so any other process moves it and the tests fail on shared machines. Each test now asks the driver about the exact addresses it used after close: the mapping lookup must fail and freeing the reservation must fail because it no longer exists. Closes run under assert_no_cuda_warning, so a failed unmap, address free or release fails the test. The leak test lists one reservation per allocate and one more per grow, for both the in-place and the moved case. One deterministic pass replaces the eight-iteration loop. Co-Authored-By: Claude Fable 5.1 --- cuda_core/tests/test_memory.py | 135 ++++++++++++++++++--------------- 1 file changed, 75 insertions(+), 60 deletions(-) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 7bd0f673d4d..bb68c17e95d 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1288,8 +1288,22 @@ def _vmm_read(buf, offset, size): return bytes(host) -def _vmm_free_memory(): - return handle_return(driver.cuMemGetInfo())[0] +def _vmm_assert_released(reservations): + """Assert that no mapping and no reservation remain for each ``(ptr, size)`` reservation. + + Both queries ask the driver about one address, so no other process can + disturb them, unlike a device-wide free-memory count. A mapping lookup + fails once the range is unmapped, and freeing an address range fails once + it is no longer reserved. With every mapping gone and every buffer's own + reference released, the driver frees the physical memory. + """ + for ptr, size in reservations: + status, handle = driver.cuMemRetainAllocationHandle(ptr) + if status == driver.CUresult.CUDA_SUCCESS: + handle_return(driver.cuMemRelease(handle)) + pytest.fail(f"address {ptr:#x} is still mapped") + (status,) = driver.cuMemAddressFree(ptr, size) + assert status != driver.CUresult.CUDA_SUCCESS, f"range {ptr:#x}+{size:#x} was still reserved" def _vmm_reserve_at(addr, size): @@ -1301,14 +1315,14 @@ def _vmm_reserve_at(addr, size): return ptr -def _vmm_allocate_in_free_hole(device, size, extra, *, miss=pytest.skip): +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. - Calls ``miss`` (``pytest.skip`` by default) when the driver places the - buffer elsewhere. Returns the resource and the buffer. + Skips when the driver places the buffer elsewhere. Returns the resource + and the buffer. """ probe = _vmm_resource(device).allocate(1) gran = probe.size @@ -1321,7 +1335,7 @@ def _vmm_allocate_in_free_hole(device, size, extra, *, miss=pytest.skip): buf = mr.allocate(size) if int(buf.handle) != int(hole): buf.close() - miss("the driver did not place the allocation at the requested address") + pytest.skip("the driver did not place the allocation at the requested address") return mr, buf @@ -1772,7 +1786,7 @@ def test_vmm_close_synchronizes_recorded_streams(init_cuda, close_input_first): @pytest.mark.agent_authored(model="claude-fable-5-1") -@pytest.mark.thread_unsafe(reason="cuMemGetInfo measures process-wide free memory") +@pytest.mark.thread_unsafe(reason="records process-global warnings") 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() @@ -1796,7 +1810,6 @@ def test_vmm_deallocate_frees_wrapped_pointer(init_cuda): 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)) @@ -1805,8 +1818,9 @@ def test_vmm_deallocate_frees_wrapped_pointer(init_cuda): buf = Buffer.from_handle(ptr=int(ptr), size=size, mr=mr) _vmm_fill(buf, 0x11) - buf.close() - assert baseline - _vmm_free_memory() < size + with assert_no_cuda_warning(): + buf.close() + _vmm_assert_released([(int(ptr), size)]) @pytest.mark.agent_authored(model="claude-fable-5-1") @@ -1834,7 +1848,7 @@ def test_vmm_buffers_alive_at_shutdown_are_freed_quietly(init_cuda): @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") +@pytest.mark.thread_unsafe(reason="records process-global warnings") 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() @@ -1842,7 +1856,6 @@ def test_vmm_failed_grow_leaves_input_intact(init_cuda): 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): @@ -1850,8 +1863,9 @@ def test_vmm_failed_grow_leaves_input_intact(init_cuda): 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() + with assert_no_cuda_warning(): + buf.close() + _vmm_assert_released([(ptr, size)]) @pytest.mark.agent_authored(model="claude-fable-5-1") @@ -1903,7 +1917,7 @@ def test_vmm_close_during_unrelated_capture_keeps_capture_valid(init_cuda, mode) @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") +@pytest.mark.thread_unsafe(reason="records process-global warnings") 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. @@ -1912,9 +1926,8 @@ def test_vmm_release_from_gc_on_capturing_stream_reports_and_unmaps(init_cuda): """ device = _vmm_device_or_skip() mr = _vmm_resource(device) - baseline = _vmm_free_memory() buf = mr.allocate(4096) - size = buf.size + ptr, size = int(buf.handle), buf.size gb = device.create_graph_builder().begin_building() try: buf.set_deallocation_stream(gb) @@ -1927,11 +1940,11 @@ def test_vmm_release_from_gc_on_capturing_stream_reports_and_unmaps(init_cuda): finally: gb.end_building() gb.close() - assert baseline - _vmm_free_memory() < size + _vmm_assert_released([(ptr, 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") +@pytest.mark.thread_unsafe(reason="records process-global warnings") 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. @@ -1945,9 +1958,8 @@ def test_vmm_release_on_default_stream_during_blocking_capture_reports_and_unmap """ device = _vmm_device_or_skip() mr = _vmm_resource(device) - baseline = _vmm_free_memory() buf = mr.allocate(4096) - size = buf.size + ptr, size = int(buf.handle), buf.size blocking = device.create_stream(options=StreamOptions(nonblocking=False)) gb = blocking.create_graph_builder().begin_building() try: @@ -1961,7 +1973,7 @@ def test_vmm_release_on_default_stream_during_blocking_capture_reports_and_unmap gb.end_building() gb.close() blocking.close() - assert baseline - _vmm_free_memory() < size + _vmm_assert_released([(ptr, size)]) @pytest.mark.agent_authored(model="claude-fable-5-1") @@ -1982,55 +1994,58 @@ def test_vmm_close_without_stream_argument_on_capturing_stream_raises(init_cuda) buf.close() -@pytest.mark.thread_unsafe(reason="cuMemGetInfo measures process-wide free memory") +@pytest.mark.thread_unsafe(reason="records process-global warnings") @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): + """Closing every buffer of a range unmaps it and frees each reservation, with no teardown failure. + + Each mode closes the buffers and then asks the driver about the exact + addresses the range used, instead of comparing a device-wide free-memory + count that other processes can move. + """ device = Device() if not device.properties.virtual_memory_management_supported: pytest.skip("Virtual memory management is not supported on this device") mr = _vmm_resource(device) - requested_size = 8 * MIB - - def allocate_and_close(warm_up=False): - if mode == "grow_moved": - # 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() - 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)) - finally: - handle_return(driver.cuMemAddressFree(decoy, buf.size)) - assert int(buffers[1].handle) != int(buf.handle) + if mode == "grow_moved": + # 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. + mr, buf = _vmm_allocate_in_free_hole(device, 2 * MIB, 4 * MIB) + 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(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(8 * MIB) + buffers = [buf] + if mode == "grow": + buffers.append(mr.modify_allocation(buf, 2 * buf.size)) + + # The reservations the range holds: one per allocate, and one more per grow. + base, size = int(buf.handle), buf.size + reservations = [(base, size)] + if len(buffers) > 1: + grown = buffers[1] + if int(grown.handle) == base: + # Grew in place: a second reservation right after the first. + reservations.append((base + size, grown.size - size)) 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 + reservations.append((int(grown.handle), grown.size)) + + # Deleters report a failed unmap, address free or release as a CUDAWarning. + with assert_no_cuda_warning(): for b in buffers: b.close() - return 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 at least one chunk per iteration; the fixed path stays near baseline. - assert baseline - free < mapped_size + _vmm_assert_released(reservations) def test_vmm_allocator_rdma_unsupported_exception(): From 312b0b4d70ca6b16c95afa4a7a9f4a18b4d8edb5 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 28 Sep 2026 17:39:32 -0700 Subject: [PATCH 09/14] fix(cuda.core): make VMM ranges immutable and give each buffer its own A range is now an immutable list of mapping handles that lives in the buffer's device pointer box. A grow copies the input's list, appends or replaces mappings, and builds a new range for its result; the input's range never changes. Mapping handles are shared between ranges, so each mapping unmaps when the last range that holds it goes. Two buffers therefore never share mutable state, which is what made concurrent grows of aliased buffers unsafe: the shared range's mapping vector was appended and iterated without a lock, and a grow that lost a race could dereference a handle another thread had emptied. With one range and one recorded stream per buffer, teardown follows the ordinary Buffer model, so the stream union, the range mutex, the base-address registry and the range header are gone. modify_allocation never returns its input any more. A request the buffer already covers returns a full alias without a driver call, so closing the result never closes the buffer passed in. It also reads the input's handle once, so a close from another thread defers the release instead of emptying what the call reads. The box behind a VMM handle is a VmmDevicePtrBox, a DevicePtrBox with the range as a member and no virtual functions. Every handle on a VirtualMemoryBuffer comes from deviceptr_create_vmm, including the size-zero buffer, which now sits on a VMM box with an empty range, so the class check in modify_allocation is what makes the downcast valid. Adds a test that grows and closes aliases of one buffer from four threads, reduced from the report on this pull request. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 119 +++++++++------- cuda_core/cuda/core/_cpp/rt/api.hpp | 43 +++--- cuda_core/cuda/core/_cpp/rt/internal.hpp | 6 + cuda_core/cuda/core/_cpp/rt/memory.cpp | 35 +++-- cuda_core/cuda/core/_cpp/rt/types.hpp | 9 +- .../cuda/core/_cpp/rt/virtual_memory.cpp | 103 +++----------- cuda_core/cuda/core/_cpp/rt/vmm_range.hpp | 38 ----- .../core/_memory/_virtual_memory_resource.pyi | 29 ++-- .../core/_memory/_virtual_memory_resource.pyx | 132 +++++++++--------- cuda_core/cuda/core/_rt.pxd | 12 +- cuda_core/cuda/core/_rt.pyx | 13 +- cuda_core/docs/source/release/1.3.0-notes.rst | 4 +- cuda_core/tests/test_memory.py | 61 +++++++- 13 files changed, 299 insertions(+), 305 deletions(-) delete mode 100644 cuda_core/cuda/core/_cpp/rt/vmm_range.hpp diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index d26a4b3c88c..ce9b0ba1be1 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -76,42 +76,57 @@ access; if access fails it unmaps and returns empty. The eight VMM entry points, ### 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 +using VmmRange = std::vector; // ascending, contiguous; sum of sizes = range total +using VmmRangeHandle = std::shared_ptr; // immutable once built +VmmRangeHandle create_vmm_range(const std::vector& mappings); +DevicePtrHandle deviceptr_create_vmm(CUdeviceptr base, const VmmRangeHandle& range); +VmmRangeHandle vmm_range(const DevicePtrHandle& h); // the range; only for handles from deviceptr_create_vmm ``` +A range is a list of mapping handles and nothing else. It is built once and never changes. A +grow reads the input's range, copies the list, appends the new mapping (or, when the buffer +moves, replaces every entry with a mapping at the new address) and wraps the copy in a new range +for the result. The input's range is untouched. Mapping handles are shared between the two +ranges, so each mapping unmaps once, when the last range that holds it goes; a chunk that only +the result holds is unmapped as soon as the result closes. + +Every buffer therefore owns exactly one range and records exactly one deallocation stream, and +two buffers never share mutable state. A full alias (a request the input already covers) shares +the input's range object, which is safe because it is immutable. This is what makes concurrent +grows of aliased buffers safe without a lock: they read a shared immutable list and write only +their own locals. Which grow extends in place and which one moves depends on what address space +the driver has free when each probe runs, exactly as for a single thread. + 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)`. +reference; grow operations call `Buffer_check_open`, copy `buf._h_ptr` into a local (to isolate +it from concurrent operations that might, e.g., call `close()`), and then call `vmm_range()` on the copy. `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 deleter 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 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 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 - the deferred-cleanup drain on the main thread, with the GIL released. +The box behind a VMM handle is a `VmmDevicePtrBox`: a `DevicePtrBox` with a `VmmRangeHandle` +member and no virtual functions, so other boxes pay nothing. Every handle on a +`VirtualMemoryBuffer` comes from `deviceptr_create_vmm`, so `modify_allocation` checks the Python +class and `vmm_range(h)` downcasts the box without a tag, the way `deallocation_stream(h)` reads +the stream. A size-zero buffer has a VMM box over an empty range. + +Deleter of a `VmmDevicePtrBox`: release the GIL; unless the +interpreter is finalizing or no stream was recorded, check the capture status of the recorded +stream and skip it 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 with the thread's capture mode switched to relaxed for +the call, so a capture on an unrelated stream is not invalidated. Then drop the range: each +mapping this buffer was the last to hold unmaps, and its reservation frees and its allocation +releases as their last references go. This is the same model as every other `Buffer`: the +recorded stream orders the release of this buffer's memory, and a caller who touches that memory +from another stream must order that work before the close. The deleter blocks on the sync, like +the pool-backed deleters block on the free, 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. @@ -147,24 +162,30 @@ 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; `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)`. - 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. + - `req <= total`: return a new `VirtualMemoryBuffer` over the same range with size + `max(req, buf.size)`; no driver call. The range already maps the request, so `cfg` is not + applied and the access of memory that is already mapped never changes. The result is a full + alias of the input when the input already covers the request; `buf` itself is never returned, + so closing the result never closes `buf`. + - `req > total`, in place: copy the input's mapping list; 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, append the + mapping to the copy, build a new range from it, and create the result's `DevicePtrHandle` on + that range at the same base, copying the input's recorded deallocation stream. 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 copied list + (shared allocation handles, their own descriptors) at `base_new + offset`, replacing the copy's + entries; create and map the new allocation; build the new range, handle (copying the input's + recorded stream), and buffer. + - Every path returns a new buffer and leaves the input open and its range unchanged. See "Why + `modify_allocation` returns a new buffer" below. + - Concurrent calls on buffers that alias one another are safe: each reads an immutable range + and writes only its own locals. Which call extends in place and which one moves depends on + what address space the driver has free when each probe runs, exactly as for a single thread. + Closing a buffer while another thread passes it to `modify_allocation` is unsupported + (concurrency.rst); the call copies the input handle first, so that misuse defers the release + rather than crashing. - `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 @@ -217,21 +238,21 @@ the whole layer; see [DESIGN.md](DESIGN.md). 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. +7. The stream a buffer recorded finishes before that buffer's share of its range is released, + except as invariant 8 states. A mapping that another buffer still holds stays mapped. 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. +9. Graph nodes and aliases keep the mappings alive. A mapping dies with the last range that holds + it, in any close order. A range is immutable once built. 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. +12. Every handle on a `VirtualMemoryBuffer` points at a VMM box whose range is the buffer's + mappings; a size-zero buffer's range is empty. 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 diff --git a/cuda_core/cuda/core/_cpp/rt/api.hpp b/cuda_core/cuda/core/_cpp/rt/api.hpp index e65617949a4..de419446b42 100644 --- a/cuda_core/cuda/core/_cpp/rt/api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/api.hpp @@ -5,6 +5,7 @@ #pragma once #include "types.hpp" +#include #include #include #include @@ -246,7 +247,8 @@ CUresult set_deallocation_stream( // 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. +// last references go. A buffer's DevicePtrHandle owns the range. Ranges are +// immutable: a grow builds a new range for its result. // ============================================================================ // Create a physical allocation via cuMemCreate. The access descriptors are @@ -279,29 +281,28 @@ VaMappingHandle create_va_mapping_handle(CUdeviceptr ptr, const MemAllocationHan 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); +// Build an immutable range from mappings in ascending, contiguous order. A +// grow builds a new range for its result and never changes the input's. +// May throw std::bad_alloc. +VmmRangeHandle create_vmm_range(const std::vector& mappings); -// Recover the range of a VMM device pointer handle; empty for any other -// handle, including a closed one. -VmmRangeHandle vmm_range(const DevicePtrHandle& h); +// The range of a device pointer handle created by deviceptr_create_vmm. Only +// for such handles: the caller (VirtualMemoryBuffer) guarantees the origin. +// Empty for an empty handle. +VmmRangeHandle vmm_range(const DevicePtrHandle& h) noexcept; -// 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; +// Range accessors: a copy of the mapping list, which a grow extends and turns +// into a new range, and the range total. Reads of an immutable range need no +// synchronization. +std::vector vmm_range_mappings(const VmmRangeHandle& range); // may throw 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. + +// Create a device pointer handle whose box owns a range (an empty range for +// a size-zero buffer). The box records no deallocation stream; set one with +// set_deallocation_stream. When the last reference is released, the recorded +// stream is synchronized (skipped with a report when that would disturb a +// capture) and the box freed, which unmaps every mapping this buffer was the +// last to hold. Returns empty handle for a null range handle. DevicePtrHandle deviceptr_create_vmm(CUdeviceptr base, const VmmRangeHandle& range); // ============================================================================ diff --git a/cuda_core/cuda/core/_cpp/rt/internal.hpp b/cuda_core/cuda/core/_cpp/rt/internal.hpp index b8f17e6087d..331e512fc77 100644 --- a/cuda_core/cuda/core/_cpp/rt/internal.hpp +++ b/cuda_core/cuda/core/_cpp/rt/internal.hpp @@ -62,6 +62,12 @@ ContextHandle deallocation_context(const DeallocationStream& stream) noexcept; // Implemented in stream.cpp bool make_deallocation_stream(const StreamHandle& h, DeallocationStream& out) noexcept; +// Implemented in virtual_memory.cpp. Synchronize the stream a VMM buffer +// recorded before its mappings are released. Skips the sync, with a report +// when a capture is the reason, if no stream was recorded, the interpreter is +// finalizing, or the sync would disturb a graph capture (VMM_DESIGN.md). +void vmm_sync_before_release(const DeallocationStream& stream) noexcept; + // Decorate a status-returning cleanup call to report whenever it fails. CUDA // calls (CUresult) are reported with the error name and description; NVRTC, // NVVM and nvJitLink calls (integer status codes) with the raw code. diff --git a/cuda_core/cuda/core/_cpp/rt/memory.cpp b/cuda_core/cuda/core/_cpp/rt/memory.cpp index eb1f36666ef..06efc3e6f7c 100644 --- a/cuda_core/cuda/core/_cpp/rt/memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/memory.cpp @@ -8,7 +8,6 @@ #include "driver_api.hpp" #include "error.hpp" #include "internal.hpp" -#include "vmm_range.hpp" #include #include #include @@ -124,6 +123,15 @@ struct DevicePtrBox { // stream tokens carry a bound context. mutable DeallocationStream deallocation; }; + +// The box behind a VirtualMemoryResource buffer: the range of mappings it +// owns (VMM_DESIGN.md). Every handle on a VirtualMemoryBuffer comes from +// deviceptr_create_vmm, so vmm_range() downcasts without a tag; there are no +// virtual functions, so other boxes pay nothing. A size-zero buffer has an +// empty range. +struct VmmDevicePtrBox : DevicePtrBox { + VmmRangeHandle range; +}; } // namespace // Recovers the owning DevicePtrBox from the aliased CUdeviceptr pointer. @@ -306,26 +314,31 @@ DevicePtrHandle deviceptr_create_mapped_graphics( // ============================================================================ DevicePtrHandle deviceptr_create_vmm(CUdeviceptr base, const VmmRangeHandle& range) { - if (!range) { + if (!range) { // a null handle; an empty range is valid 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 { + auto box = std::shared_ptr( + new VmmDevicePtrBox{{base, DeallocationStream{}}, range}, + [](VmmDevicePtrBox* b) { 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); + // Order the release on this buffer's recorded stream, then free + // the box, which drops the range: every mapping this buffer was + // the last to hold unmaps, and its reservation and allocation + // follow. + vmm_sync_before_release(b->deallocation); delete b; - range.reset(); } ); return DevicePtrHandle(box, &box->resource); } +VmmRangeHandle vmm_range(const DevicePtrHandle& h) noexcept { + // Only for handles from deviceptr_create_vmm; the VirtualMemoryBuffer + // class guarantees that for its callers. + return h ? static_cast(get_box(h))->range : VmmRangeHandle{}; +} + // ============================================================================ // MemoryResource-owned Device Pointer Handles // ============================================================================ diff --git a/cuda_core/cuda/core/_cpp/rt/types.hpp b/cuda_core/cuda/core/_cpp/rt/types.hpp index af89a75e23f..6cc061244e1 100644 --- a/cuda_core/cuda/core/_cpp/rt/types.hpp +++ b/cuda_core/cuda/core/_cpp/rt/types.hpp @@ -8,6 +8,7 @@ #include #include #include +#include // Forward declaration for NVVM - avoids nvvm.h dependency // Use void* to match cuda.bindings.cynvvm's typedef @@ -77,12 +78,14 @@ 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. +// reservation (cuMemMap), and the range of mappings a buffer owns. A range is +// the buffer's mappings in ascending, contiguous order; it is immutable once +// built, and a grow builds a new one for its result (VMM_DESIGN.md). using MemAllocationHandle = std::shared_ptr; using VaReservationHandle = std::shared_ptr; using VaMappingHandle = std::shared_ptr; -struct VmmRange; -using VmmRangeHandle = std::shared_ptr; +using VmmRange = std::vector; +using VmmRangeHandle = std::shared_ptr; using DevicePtrHandle = std::shared_ptr; diff --git a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp index 340a2f62787..bb722cf9b1b 100644 --- a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp @@ -12,10 +12,8 @@ #include "driver_api.hpp" #include "error.hpp" #include "internal.hpp" -#include "vmm_range.hpp" #include #include -#include #include #include @@ -165,33 +163,6 @@ MemAllocationHandle va_mapping_allocation(const VaMappingHandle& h) noexcept { // 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 @@ -258,71 +229,41 @@ static void sync_recorded_stream(const DeallocationStream& ds, bool& capture_ski } } -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{}; +namespace detail { +void vmm_sync_before_release(const DeallocationStream& stream) noexcept { + // Order the release on this buffer's stream, as the other DevicePtrBox + // deleters do, unless the interpreter is finalizing or the sync would + // disturb a capture. A host-located buffer recorded no stream. + if (!stream.h_stream || py_is_finalizing()) { + return; + } + bool capture_skipped = false; + sync_recorded_stream(stream, 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 buffer was unmapped without synchronizing that stream"); + } } +} // namespace detail -size_t vmm_range_count(const VmmRangeHandle& range) noexcept { - return range ? range->mappings.size() : 0; +VmmRangeHandle create_vmm_range(const std::vector& mappings) { + return std::make_shared(mappings); } -VaMappingHandle vmm_range_mapping(const VmmRangeHandle& range, size_t index) noexcept { - if (!range || index >= range->mappings.size()) { - return {}; - } - return range->mappings[index]; +std::vector vmm_range_mappings(const VmmRangeHandle& range) { + return range ? *range : VmmRange{}; } size_t vmm_range_total(const VmmRangeHandle& range) noexcept { size_t total = 0; if (range) { - for (const VaMappingHandle& m : range->mappings) { + for (const VaMappingHandle& m : *range) { 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 deleted file mode 100644 index 8c699d4397f..00000000000 --- a/cuda_core/cuda/core/_cpp/rt/vmm_range.hpp +++ /dev/null @@ -1,38 +0,0 @@ -// 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/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index cced7cf0e56..8dfe51b74bc 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -93,11 +93,11 @@ class VirtualMemoryBuffer(Buffer): 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. When the stream this close uses, given or recorded, is not + Each mapping, reservation and physical allocation goes away when the + last buffer that holds it closes. Before this buffer's share is + released, its deallocation stream is synchronized, as for every + :class:`Buffer`. Virtual memory deallocation is synchronous and cannot + be 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, @@ -173,14 +173,14 @@ class VirtualMemoryResource(MemoryResource): """ 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. + The buffer passed in stays open and usable, and is never returned. 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. Closing the returned buffer never closes ``buf``. - This method is not thread-safe with respect to two buffers that share an - address range. + Concurrent calls on buffers that alias one another are safe. Parameters ---------- @@ -200,8 +200,9 @@ class VirtualMemoryResource(MemoryResource): Returns ------- VirtualMemoryBuffer - ``buf`` itself when it already covers ``new_size``; otherwise a new - buffer of the rounded size. + A new buffer of at least ``new_size`` and at least ``buf.size`` + bytes. When ``buf`` already covers the request, the result is a + full alias of it and no driver call is made. Raises ------ diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 007cd93e00a..4b698c8633b 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -29,7 +29,6 @@ from cuda.core._rt cimport ( create_va_reservation_handle, create_vmm_range, deallocation_stream, - deviceptr_create_ref, deviceptr_create_vmm, get_last_error, get_primary_context, @@ -37,10 +36,7 @@ from cuda.core._rt cimport ( set_deallocation_stream, va_mapping_allocation, vmm_range, - vmm_range_append, - vmm_range_count, - vmm_range_mapping, - vmm_range_reserve, + vmm_range_mappings, vmm_range_total, ) from cuda.core._stream cimport Stream, Stream_accept, Stream_handle_is_default_token, Stream_is_default_token @@ -232,11 +228,11 @@ cdef class VirtualMemoryBuffer(Buffer): 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. When the stream this close uses, given or recorded, is not + Each mapping, reservation and physical allocation goes away when the + last buffer that holds it closes. Before this buffer's share is + released, its deallocation stream is synchronized, as for every + :class:`Buffer`. Virtual memory deallocation is synchronous and cannot + be 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, @@ -449,15 +445,16 @@ cdef class VirtualMemoryResource(MemoryResource): cdef MemAllocationHandle h_alloc cdef VaReservationHandle h_res cdef VaMappingHandle h_map + cdef vector[VaMappingHandle] mappings cdef VmmRangeHandle rng cdef DevicePtrHandle h_ptr if size == 0: - # Nothing to reserve or map; an empty buffer with a non-owning handle. - # 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) + # Nothing to reserve or map: a VMM box over an empty range, so the + # buffer is a VirtualMemoryBuffer like every other. 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_vmm(0, create_vmm_range(mappings)) 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) @@ -484,8 +481,8 @@ cdef class VirtualMemoryResource(MemoryResource): if not h_map: _raise_last_error() - rng = create_vmm_range(as_cu(h_res)) - vmm_range_append(rng, h_map) + mappings.push_back(h_map) + rng = create_vmm_range(mappings) h_ptr = deviceptr_create_vmm(as_cu(h_res), rng) if not h_ptr: _raise_last_error() @@ -498,14 +495,14 @@ cdef class VirtualMemoryResource(MemoryResource): """ 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. + The buffer passed in stays open and usable, and is never returned. 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. Closing the returned buffer never closes ``buf``. - This method is not thread-safe with respect to two buffers that share an - address range. + Concurrent calls on buffers that alias one another are safe. Parameters ---------- @@ -525,8 +522,9 @@ cdef class VirtualMemoryResource(MemoryResource): Returns ------- VirtualMemoryBuffer - ``buf`` itself when it already covers ``new_size``; otherwise a new - buffer of the rounded size. + A new buffer of at least ``new_size`` and at least ``buf.size`` + bytes. When ``buf`` already covers the request, the result is a + full alias of it and no driver call is made. Raises ------ @@ -545,59 +543,71 @@ cdef class VirtualMemoryResource(MemoryResource): If a driver call fails. ``buf`` is untouched when this method raises. """ cdef Buffer b + cdef DevicePtrHandle h_in + cdef size_t in_size 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 vector[VaMappingHandle] mappings + cdef size_t gran, req, total, add, addr_align, offset, i 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 VaMappingHandle h_map, m2 + cdef DevicePtrHandle h_ptr_new cdef object new_buf - if not isinstance(buf, Buffer): - raise TypeError(f"buf must be a Buffer, got {type(buf).__name__}") + # Every VirtualMemoryBuffer handle comes from deviceptr_create_vmm, so + # the class check is what makes vmm_range() below valid. + if not isinstance(buf, VirtualMemoryBuffer): + raise TypeError(f"buf must be a VirtualMemoryBuffer, got {type(buf).__name__}") b = buf Buffer_check_open(b) + # Read the input's handle and size once. The handle copy keeps the + # input's box alive for the rest of the call, so a close from another + # thread (unsupported; see concurrency.rst) defers the release instead + # of emptying what this call reads. + h_in = b._h_ptr + in_size = b._size 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 ) self._check_config(cfg) - if b._size == 0: + if in_size == 0: # 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) + self._copy_deallocation_stream((new_buf)._h_ptr, h_in) return new_buf - rng = vmm_range(b._h_ptr) - if not rng: - raise TypeError("buf was not allocated by VirtualMemoryResource.allocate") + rng = vmm_range(h_in) 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) + base = as_cu(h_in) - 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: + # The range already maps the request: a new buffer over the same + # range, at least as large as the input, with no driver call. The + # result is a distinct object, so closing it never closes ``buf``. + h_ptr_new = deviceptr_create_vmm(base, rng) + if not h_ptr_new: _raise_last_error() - self._copy_deallocation_stream(h_ptr2, b._h_ptr) - return Buffer_from_deviceptr_handle(h_ptr2, req, self, None, VirtualMemoryBuffer) + self._copy_deallocation_stream(h_ptr_new, h_in) + return Buffer_from_deviceptr_handle( + h_ptr_new, req if req > in_size else in_size, 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) + # The result's range is built from a copy of the input's mappings; the + # input's range is immutable and never changes. + mappings = vmm_range_mappings(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. @@ -616,16 +626,13 @@ cdef class VirtualMemoryResource(MemoryResource): 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: + mappings.push_back(h_map) + rng_new = create_vmm_range(mappings) + h_ptr_new = deviceptr_create_vmm(base, rng_new) + if not h_ptr_new: _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 + self._copy_deallocation_stream(h_ptr_new, h_in) + return Buffer_from_deviceptr_handle(h_ptr_new, req, self, None, VirtualMemoryBuffer) # Move: a new range that maps every existing allocation, then the new one. # The allocations are shared with the input buffer's range. @@ -635,19 +642,15 @@ cdef class VirtualMemoryResource(MemoryResource): 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) + for i in range(mappings.size()): + a = va_mapping_allocation(mappings[i]) 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 + offset += mem_allocation_size(a) + mappings[i] = m2 # the input's range still holds the old mapping with nogil: h_alloc = create_mem_allocation_handle(add, prop, descs.data(), descs.size()) if not h_alloc: @@ -656,11 +659,12 @@ cdef class VirtualMemoryResource(MemoryResource): 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) + mappings.push_back(m2) + rng_new = create_vmm_range(mappings) 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) + self._copy_deallocation_stream(h_ptr_new, h_in) 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: diff --git a/cuda_core/cuda/core/_rt.pxd b/cuda_core/cuda/core/_rt.pxd index b8b2ca5e416..e83053db14d 100644 --- a/cuda_core/cuda/core/_rt.pxd +++ b/cuda_core/cuda/core/_rt.pxd @@ -7,6 +7,7 @@ from libc.stddef cimport size_t from libc.stdint cimport intptr_t from libcpp.memory cimport shared_ptr, unique_ptr +from libcpp.vector cimport vector from cuda.bindings cimport cydriver from cuda.bindings cimport cynvrtc @@ -72,7 +73,7 @@ cdef extern from "_cpp/rt/handles.hpp" namespace "cuda_core::rt": ctypedef shared_ptr[const MemAllocationValue] MemAllocationHandle ctypedef shared_ptr[const VaReservationValue] VaReservationHandle ctypedef shared_ptr[const VaMappingValue] VaMappingHandle - ctypedef shared_ptr[VmmRange] VmmRangeHandle + ctypedef shared_ptr[const VmmRange] VmmRangeHandle # Type-erased shared owner for resources attached to graph node slots. # Typed handles above assign directly to an OpaqueHandle (shared control @@ -306,13 +307,10 @@ cdef VaMappingHandle create_va_mapping_handle( 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 VmmRangeHandle create_vmm_range(const vector[VaMappingHandle]& mappings) except+ nogil +cdef VmmRangeHandle vmm_range(const DevicePtrHandle& h) noexcept nogil +cdef vector[VaMappingHandle] vmm_range_mappings(const VmmRangeHandle& range) except+ 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 diff --git a/cuda_core/cuda/core/_rt.pyx b/cuda_core/cuda/core/_rt.pyx index 70483960a9e..0d58ecfd322 100644 --- a/cuda_core/cuda/core/_rt.pyx +++ b/cuda_core/cuda/core/_rt.pyx @@ -185,16 +185,11 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt": 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 + const vector[VaMappingHandle]& mappings) except+ nogil + VmmRangeHandle vmm_range "cuda_core::rt::vmm_range" (const DevicePtrHandle& h) noexcept nogil + vector[VaMappingHandle] vmm_range_mappings "cuda_core::rt::vmm_range_mappings" ( + const VmmRangeHandle& range) except+ 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 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 4442c239967..51b2d0a8d6e 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -172,7 +172,9 @@ Fixes and enhancements :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 + range. The result is always a new buffer, so closing it never closes the + buffer passed in; a request the buffer already covers returns a full alias + without a driver call, 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 `__); diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index bb68c17e95d..84c71c741ca 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -6,6 +6,7 @@ import multiprocessing as mp import sys import textwrap +import threading import warnings from cuda.bindings import driver @@ -1508,9 +1509,16 @@ def test_vmm_allocator_grow_allocation(handle_type): assert buffer.size == original_size assert int(buffer.handle) != 0 - # 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 + # Requests the buffer already covers return a full alias, never the buffer + # itself, so closing the result never closes the input. + same = vmm_mr.modify_allocation(grown_buffer, 4 * MIB) + smaller = vmm_mr.modify_allocation(grown_buffer, 2 * MIB) + assert same is not grown_buffer and smaller is not grown_buffer + assert same.size == smaller.size == grown_buffer.size + assert int(same.handle) == int(smaller.handle) == int(grown_buffer.handle) + same.close() + smaller.close() + _vmm_fill(grown_buffer, 0x11) # the input is still mapped # Clean up buffer.close() @@ -1612,11 +1620,11 @@ def test_vmm_unaligned_and_repeated_grow_close_cleanly(init_cuda): 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. + # A second grow of the input builds its own result. The input's range + # is immutable, so this grow cannot reuse the chunk ``second`` added. 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) + assert alias is not buf + assert alias.size == 2 * gran # Every buffer maps the first chunk. _vmm_fill(third, 0x77) @@ -1628,6 +1636,45 @@ def test_vmm_unaligned_and_repeated_grow_close_cleanly(init_cuda): second.close() +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="spawns worker threads that grow shared buffers") +def test_vmm_concurrent_grows_of_aliased_buffers_are_memory_safe(init_cuda): + """Growing and closing aliases of one buffer from several threads never crashes. + + Reduced from the reproduction on PR #2917. A range is immutable, so a grow + reads the input's range and builds a new one, and ``modify_allocation`` + never returns its input, so closing a result never closes a buffer another + thread is using. Which grow extends in place and which moves is undefined; + the process must survive and every buffer must close cleanly. + """ + device = _vmm_device_or_skip() + mr = _vmm_resource(device) + nthreads = 4 + errors = [] + + def work(buf, i, gate): + try: + gate.wait() + for j in range(6): + mr.modify_allocation(buf, (i * 8 + j + 2) * 2 * MIB).close() + except BaseException as exc: # pragma: no cover - surfaced below + errors.append(exc) + + with assert_no_cuda_warning(): + for _ in range(8): + root = mr.allocate(2 * MIB) + alias = mr.modify_allocation(root, 8 * MIB) + gate = threading.Barrier(nthreads) + threads = [threading.Thread(target=work, args=(root if i % 2 else alias, i, gate)) for i in range(nthreads)] + for t in threads: + t.start() + for t in threads: + t.join() + alias.close() + root.close() + assert errors == [] + + @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.""" From beea7662068503cd4fc6efb0544f2282a049d2b0 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 28 Sep 2026 17:47:14 -0700 Subject: [PATCH 10/14] fix(cuda.core): return an empty stream for an empty device pointer handle deallocation_stream() was the one accessor in the layer that dereferenced an empty handle instead of returning empty, as its sibling set_deallocation_stream and every as_cu() overload do. A buffer closed by one thread while another still reads its handle now gets an empty stream rather than a crash. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/api.hpp | 1 + cuda_core/cuda/core/_cpp/rt/memory.cpp | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/api.hpp b/cuda_core/cuda/core/_cpp/rt/api.hpp index 3bfff91b4df..d77e57c86e5 100644 --- a/cuda_core/cuda/core/_cpp/rt/api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/api.hpp @@ -233,6 +233,7 @@ DevicePtrHandle deviceptr_import_ipc( // Access the deallocation stream for a device pointer handle (read-only). // For non-owning handles, the stream is not used but can still be accessed. +// Returns an empty handle for an empty device pointer handle. StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept; // Set the deallocation stream for a device pointer handle. diff --git a/cuda_core/cuda/core/_cpp/rt/memory.cpp b/cuda_core/cuda/core/_cpp/rt/memory.cpp index 06efc3e6f7c..e5287a39183 100644 --- a/cuda_core/cuda/core/_cpp/rt/memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/memory.cpp @@ -146,9 +146,10 @@ static DevicePtrBox* get_box(const DevicePtrHandle& h) { ); } -// Return the stream that orders a device pointer's deallocation. +// Return the stream that orders a device pointer's deallocation; empty for an +// empty handle, as set_deallocation_stream rejects one. StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept { - return get_box(h)->deallocation.h_stream; + return h ? get_box(h)->deallocation.h_stream : StreamHandle{}; } // Replace the stream that orders a device pointer's deallocation. From 988e342e5aedff419b2da4d578bf58884e2685a2 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 28 Sep 2026 17:56:44 -0700 Subject: [PATCH 11/14] test(cuda.core): bound the joins in the concurrent VMM grow test Join each worker with the suite's sanitizer-aware timeout and assert that none is still alive before the shared buffers are closed, as the other threading tests do. Drop "undefined" from the docstring. Co-Authored-By: Claude Fable 5.1 --- cuda_core/tests/test_memory.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 84c71c741ca..e2576bd3930 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1644,8 +1644,9 @@ def test_vmm_concurrent_grows_of_aliased_buffers_are_memory_safe(init_cuda): Reduced from the reproduction on PR #2917. A range is immutable, so a grow reads the input's range and builds a new one, and ``modify_allocation`` never returns its input, so closing a result never closes a buffer another - thread is using. Which grow extends in place and which moves is undefined; - the process must survive and every buffer must close cleanly. + thread is using. Which grow extends in place and which one moves depends on + the address space the driver has free, as for one thread; the process must + survive and every buffer must close cleanly. """ device = _vmm_device_or_skip() mr = _vmm_resource(device) @@ -1669,7 +1670,8 @@ def work(buf, i, gate): for t in threads: t.start() for t in threads: - t.join() + t.join(timeout=CHILD_TIMEOUT_SEC) + assert not any(t.is_alive() for t in threads), "a grow thread did not finish" alias.close() root.close() assert errors == [] From 53479be6e8801b32ab98b353d348af4befbf8f95 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 1 Oct 2026 13:05:59 -0700 Subject: [PATCH 12/14] fix(cuda.core): keep the pending exception across a MemoryResource deleter The deleter behind Buffer.from_handle(..., mr) acquired the GIL and called deallocate() while an exception could be propagating through the caller that released the last reference. The handler that reports a failed deallocate() then cleared that exception, and the caller returned an error with no exception set, which Python reports as SystemError. report_message already saved and restored the pending exception inline. PendingExceptionGuard (py.hpp) saves the exception in flight and restores it on scope exit, dropping anything the scope itself raised. The deleter and report_message use it. The regression test releases a temporary buffer whose deallocate() fails while a TypeError propagates and expects the TypeError and a CUDAWarning. Found in review of #2917. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/memory.cpp | 4 +++ cuda_core/cuda/core/_cpp/rt/py.hpp | 39 +++++++++++++++++++++++ cuda_core/cuda/core/_cpp/rt/py_report.cpp | 12 +------ cuda_core/tests/test_memory.py | 24 ++++++++++++++ 4 files changed, 68 insertions(+), 11 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/memory.cpp b/cuda_core/cuda/core/_cpp/rt/memory.cpp index e5287a39183..7f4d26b1e25 100644 --- a/cuda_core/cuda/core/_cpp/rt/memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/memory.cpp @@ -365,6 +365,10 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* [mr, size](DevicePtrBox* b) { GILAcquireGuard gil; if (gil.acquired()) { + // The last reference may go while an exception propagates + // through the releasing caller; deallocate() must run with a + // clean error state and leave that exception in place. + PendingExceptionGuard pending; if (mr_dealloc_cb) { const DeallocationStream& stream = b->deallocation; cleanup_in_context( diff --git a/cuda_core/cuda/core/_cpp/rt/py.hpp b/cuda_core/cuda/core/_cpp/rt/py.hpp index 55212b5c42b..04c50be4a4e 100644 --- a/cuda_core/cuda/core/_cpp/rt/py.hpp +++ b/cuda_core/cuda/core/_cpp/rt/py.hpp @@ -94,6 +94,45 @@ class GILAcquireGuard { bool acquired_; }; +// Save the Python exception in flight and restore it on scope exit, dropping +// anything the scope itself raised. A deleter may run while an exception is +// propagating through the caller that released the last reference; Python +// code it calls (a MemoryResource's deallocate, the warnings machinery) must +// start with a clean error state and must not leave that exception cleared, +// or the caller returns an error with no exception set. Allocates nothing. +// Construct with the GIL held and destroy before releasing it. +class PendingExceptionGuard { +public: + PendingExceptionGuard() noexcept { +#if PY_VERSION_HEX >= 0x030C0000 + pending_ = PyErr_GetRaisedException(); +#else + PyErr_Fetch(&type_, &value_, &tb_); +#endif + } + + ~PendingExceptionGuard() { + PyErr_Clear(); +#if PY_VERSION_HEX >= 0x030C0000 + PyErr_SetRaisedException(pending_); +#else + PyErr_Restore(type_, value_, tb_); +#endif + } + + PendingExceptionGuard(const PendingExceptionGuard&) = delete; + PendingExceptionGuard& operator=(const PendingExceptionGuard&) = delete; + +private: +#if PY_VERSION_HEX >= 0x030C0000 + PyObject* pending_ = nullptr; +#else + PyObject* type_ = nullptr; + PyObject* value_ = nullptr; + PyObject* tb_ = nullptr; +#endif +}; + // as_py() - convert handle to Python wrapper object (returns new reference) namespace detail { // n.b. class lookup is not cached to avoid deadlock hazard, see DESIGN.md diff --git a/cuda_core/cuda/core/_cpp/rt/py_report.cpp b/cuda_core/cuda/core/_cpp/rt/py_report.cpp index cde00589650..92fbb53c193 100644 --- a/cuda_core/cuda/core/_cpp/rt/py_report.cpp +++ b/cuda_core/cuda/core/_cpp/rt/py_report.cpp @@ -28,22 +28,12 @@ void report_message(const char* message) noexcept { GILAcquireGuard gil; if (gil.acquired()) { // Deleters can run while a Python exception is propagating; keep it. -#if PY_VERSION_HEX >= 0x030C0000 - PyObject* pending = PyErr_GetRaisedException(); -#else - PyObject *pending_type, *pending_value, *pending_tb; - PyErr_Fetch(&pending_type, &pending_value, &pending_tb); -#endif + PendingExceptionGuard pending; if (PyErr_WarnEx(category, message, 1) != 0) { PyObject* subject = PyUnicode_FromString(message); PyErr_WriteUnraisable(subject); Py_XDECREF(subject); } -#if PY_VERSION_HEX >= 0x030C0000 - PyErr_SetRaisedException(pending); -#else - PyErr_Restore(pending_type, pending_value, pending_tb); -#endif return; } } diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index e2576bd3930..ac67d369575 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -896,6 +896,30 @@ def test_mr_deallocation_failure_warns(): buf.close() +@pytest.mark.thread_unsafe(reason="records process-global warnings and mutates the context stack") +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_mr_deallocation_during_exception_keeps_the_exception(): + """A Buffer released while an exception propagates leaves that exception in place. + + The deleter calls ``mr.deallocate`` on the Python thread. It must start + with a clean error state and restore the pending exception afterwards. + Before the fix, reporting a failed ``deallocate`` cleared the pending + exception, and the caller returned an error with no exception set, which + Python reports as ``SystemError`` (found in review of #2917). + """ + device = Device() + device.set_current() + FailingMR, _ = make_instrumented_memory_resource(deallocate_error=RuntimeError("expected deallocation failure")) + mr = FailingMR(device) + + def reject(buf): + raise TypeError("rejected") + + with pytest.warns(CUDAWarning, match="expected deallocation failure"), pytest.raises(TypeError, match="rejected"): + # The temporary Buffer is released while TypeError propagates. + reject(Buffer.from_handle(1, 1024, mr=mr)) + + @pytest.mark.thread_unsafe(reason="records process-global warnings and mutates the context stack") @pytest.mark.agent_authored(model="cursor-grok-4.5") @pytest.mark.parametrize("replace_stream", [False, True]) From 6cb1826830e6dc23dcdab30c5340da4b34a086eb Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 1 Oct 2026 13:05:59 -0700 Subject: [PATCH 13/14] docs(cuda.core): restore __weakref__ on VirtualMemoryResource and state the release wait The Python class had __weakref__ as a subclass of an extension type; the cdef class declares it, as Buffer and the pool-backed resources do. VMM_DESIGN.md compared the blocking release to pool-backed deleters, which do not block: cuMemFreeAsync is stream-ordered. The precedent is _SynchronousMemoryResource and LegacyPinnedMemoryResource, which wait in deallocate() on the same deleter path. The design doc, the class docstring, close() and the release note now say that closing a buffer waits for the work on its deallocation stream and how to control when that happens. The release note also drops a sentence about the range deleter synchronizing every recorded stream, which immutable ranges made false. The follow-up for a stream-ordered release is #2989. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 12 +++++++++--- .../core/_memory/_virtual_memory_resource.pyi | 15 +++++++++++---- .../core/_memory/_virtual_memory_resource.pyx | 16 ++++++++++++---- cuda_core/docs/source/release/1.3.0-notes.rst | 6 ++++-- 4 files changed, 36 insertions(+), 13 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index ce9b0ba1be1..35faacaa612 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -124,9 +124,15 @@ the call, so a capture on an unrelated stream is not invalidated. Then drop the mapping this buffer was the last to hold unmaps, and its reservation frees and its allocation releases as their last references go. This is the same model as every other `Buffer`: the recorded stream orders the release of this buffer's memory, and a caller who touches that memory -from another stream must order that work before the close. The deleter blocks on the sync, like -the pool-backed deleters block on the free, and it may run inside the deferred-cleanup drain on -the main thread, with the GIL released. +from another stream must order that work before the close. The deleter blocks until that work +completes. `cuMemUnmap` is not stream-ordered, so waiting is the only way to honor the order; +`_SynchronousMemoryResource` and `LegacyPinnedMemoryResource` wait the same way in +`deallocate()`, while pool-backed buffers never block because `cuMemFreeAsync` is stream-ordered. +The wait happens wherever the last reference goes: an explicit `close()`, a garbage collection, +or the deferred-cleanup drain on the main thread, with the GIL released. A caller who needs to +control when it happens closes the buffer explicitly or records an idle deallocation stream. +Emulating stream order with a host callback that hands the release to the deferred-cleanup +queue is planned as a follow-up for every synchronous release in cuda.core. 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. diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index 8dfe51b74bc..2deaca761e0 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -94,10 +94,11 @@ class VirtualMemoryBuffer(Buffer): """Release this buffer's share of its address range. Each mapping, reservation and physical allocation goes away when the - last buffer that holds it closes. Before this buffer's share is - released, its deallocation stream is synchronized, as for every - :class:`Buffer`. Virtual memory deallocation is synchronous and cannot - be captured. When the stream this close uses, given or recorded, is not + last buffer that holds it closes. Virtual memory is unmapped + synchronously, so this call waits for the work queued on the + deallocation stream before it releases this buffer's share; a buffer + released by the garbage collector waits at that point instead. The + unmap cannot be 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, @@ -134,6 +135,12 @@ class VirtualMemoryResource(MemoryResource): 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. + + Virtual memory is unmapped synchronously, so closing a buffer waits for + the work queued on its deallocation stream. A buffer released by the + garbage collector waits at that point instead. To control when the wait + happens, close the buffer explicitly or record an idle stream with + :meth:`Buffer.set_deallocation_stream`. """ device: object config: object diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 4b698c8633b..def6ad9ee57 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -229,10 +229,11 @@ cdef class VirtualMemoryBuffer(Buffer): """Release this buffer's share of its address range. Each mapping, reservation and physical allocation goes away when the - last buffer that holds it closes. Before this buffer's share is - released, its deallocation stream is synchronized, as for every - :class:`Buffer`. Virtual memory deallocation is synchronous and cannot - be captured. When the stream this close uses, given or recorded, is not + last buffer that holds it closes. Virtual memory is unmapped + synchronously, so this call waits for the work queued on the + deallocation stream before it releases this buffer's share; a buffer + released by the garbage collector waits at that point instead. The + unmap cannot be 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, @@ -288,11 +289,18 @@ cdef class VirtualMemoryResource(MemoryResource): 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. + + Virtual memory is unmapped synchronously, so closing a buffer waits for + the work queued on its deallocation stream. A buffer released by the + garbage collector waits at that point instead. To control when the wait + happens, close the buffer explicitly or record an idle stream with + :meth:`Buffer.set_deallocation_stream`. """ cdef: readonly object device readonly object config + object __weakref__ def __init__(self, device_id: Device | int, config: VirtualMemoryResourceOptions | None = None) -> None: self.device = Device(device_id) 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 bd3b229497c..d3abd683882 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -210,8 +210,10 @@ Fixes and enhancements 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 + to ``allocate()``, and closing it waits for the work queued on that stream + before it unmaps, because ``cuMemUnmap`` is not stream-ordered; close the + buffer explicitly, or record an idle deallocation stream, to control when + that wait happens. ``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. From 818a4406ae16fbae6ef5ce4987403361f5ac3c9f Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 1 Oct 2026 13:28:06 -0700 Subject: [PATCH 14/14] fix(cuda.core): query capture state with cuStreamIsCapturing in the VMM deleter cuStreamGetCaptureInfo has two slots in the cuda-bindings loader (v2 and v3) and a different arity per CUDA major, which needed a build-major fence and broke the driver-table test that builds a fake table from the first slot per name. cuStreamIsCapturing is the query that cuStreamGetCaptureInfo makes first: same status, including the implicit capture error for the legacy stream, one slot, one signature. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 5 +++-- cuda_core/cuda/core/_cpp/rt/driver_api.hpp | 3 +-- cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp | 9 +++------ 3 files changed, 7 insertions(+), 10 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index 35faacaa612..c2001b5da19 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -70,8 +70,9 @@ size_t mem_allocation_size(...) noexcept; size_t va_reservation_size(...) noexc 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. +access; if access fails it unmaps and returns empty. The seven VMM entry points, plus +`cuStreamSynchronize`, `cuStreamIsCapturing` and `cuThreadExchangeStreamCaptureMode`, join the +`driver_api` function table. ### The range and the device pointer diff --git a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp index 4c61b8386de..9efc1186c2c 100644 --- a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp @@ -137,8 +137,7 @@ namespace cuda_core::rt { X(cuMemSetAccess, 10020) \ /* cuda-bindings requests 7000 (PTDS) or 2000 (legacy) */ \ X(cuStreamSynchronize, 7000) \ - /* cuda.h maps to _v3 (12030) on the CUDA 13 header, _v2 (11030) on CUDA 12 */ \ - X(cuStreamGetCaptureInfo, 12030) \ + X(cuStreamIsCapturing, 10000) \ X(cuThreadExchangeStreamCaptureMode, 10010) #define CUDA_CORE_DECLARE_DRIVER_FN(name, introduced) extern decltype(&name) p_##name; diff --git a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp index c02011d73eb..8a946180794 100644 --- a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp @@ -167,14 +167,11 @@ MemAllocationHandle va_mapping_allocation(const VaMappingHandle& h) noexcept { // 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. +// such a capture; cuStreamIsCapturing does not. It is the query +// cuStreamGetCaptureInfo makes first, with one signature on every CUDA major. static bool sync_would_disturb_capture(CUstream stream) noexcept { CUstreamCaptureStatus status = CU_STREAM_CAPTURE_STATUS_NONE; -#if CUDA_CORE_BUILD_MAJOR >= 13 - const CUresult result = DRIVER_CALL(cuStreamGetCaptureInfo, stream, &status, nullptr, nullptr, nullptr, nullptr, nullptr); -#else - const CUresult result = DRIVER_CALL(cuStreamGetCaptureInfo, stream, &status, nullptr, nullptr, nullptr, nullptr); -#endif + const CUresult result = DRIVER_CALL(cuStreamIsCapturing, stream, &status); if (result == CUDA_ERROR_STREAM_CAPTURE_IMPLICIT) { return true; }