EPIC: VirtualMemoryResource ownership redesign

[Andy-Jost]

This epic collects the open VirtualMemoryResource bugs, states the fix, and asks contributors to hold new VMM pull requests until the fix lands.

What is wrong

The twelve issues and seven pull requests listed below report about twenty bugs in VirtualMemoryResource. All of them live in one file of about 640 lines. Most of the serious ones have one cause.

When a buffer is freed, the resource knows only the pointer and size it recorded at allocation. From those two values, the code must undo every driver call that built the buffer, in the right order. It does this by hand, with eight calls spread over three code paths. That works for a plain allocation. It fails after a grow.

A grow extends an existing buffer. If the address range right after the buffer is free, the resource reserves it and maps new memory there (the fast path). If that range is taken, the resource reserves a larger range elsewhere and moves the buffer (the slow path). Either way, a grown buffer owns two address reservations, two physical allocations, and two mappings. The driver frees a reservation only when the pointer and size match one reservation exactly (cuMemAddressFree). So no single deallocate(pointer, size) call can free a grown buffer, whatever size it passes.

The pool-backed resources do not have this problem. Each of their buffers carries a C++ handle that knows how to free itself and holds the handles it depends on.

The fix

The implementation plan is posted in this comment; it will land in the PR as cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md.

Move VirtualMemoryResource onto the same handle layer (_rt). Each physical allocation, address reservation, and mapping gets its own std::shared_ptr handle, and a Buffer owns its mappings. Teardown order then follows from ownership instead of from hand-paired driver calls. The redesign also decides how the free is ordered on the stream (#2886), whether a subclass's deallocate() still runs (#2615), and whether a grown buffer keeps its pointer and identity.

This fixes #2887, the unaligned size recorded on grown buffers, the second grow that fails, the slow-path failure that cannot be undone, the stream argument that allocate() drops, and the fast path that never runs (#2388 item 2). Owner: @Andy-Jost. Milestone: cuda.core 1.3.0.

Two smaller groups of work go with it:

1. Fixes that were in flight when the epic opened: cuda.core: close the old buffer when a VMM grow moves the mapping #2880 (spurious warning when a slow-path grow moves the buffer), fix(cuda.core): order VMM unmaps on the stream #2889 (stream sync before unmap), and the failed-reservation part of Avoid freeing failed VMM grow reservations #2237. The redesign PR fix(cuda.core): move VirtualMemoryResource onto the _rt handle layer #2917 replaces the module they edit and carries their fixes, so they are superseded and close when fix(cuda.core): move VirtualMemoryResource onto the _rt handle layer #2917 merges. Fix VMM handle leaks in virtual memory paths #2235, which fixes the physical-allocation leak, has merged.
2. Option validation, as small follow-up PRs after the redesign: addr_align, shrink requests, MANAGED, host location defaults ([BUG]: VirtualMemoryResource host_numa allocations always fail #2694), the Windows default handle type, an assert that guards user input, size 0, device_id and is_device_accessible for host-located resources, and the config= argument that persists (with Review test_memory.py::test_vmm_allocator_policy_configuration xfail #1300).

Contributors

Thanks to @fallintoplace and @aryanputta for the reports and fixes. #2344 and #2235 found and closed the largest leak, and #2886 found the missing stream sync.

Please hold new VMM pull requests until the redesign lands. For the open PRs, the table below states what happens to each. Changes that try to free a grown buffer through one deallocate(pointer, size) call cannot be merged, for the driver reason above (#2887, #2890). Held PRs stay open. Discussion of the plan belongs on this epic.

On tests: we do not merge tests that monkeypatch driver.* entry points or hand fake Buffer objects to the resource (see the review on #2235). A VMM test allocates against the real driver and checks what it can observe: the change in free memory that cuMemGetInfo reports, the buffer contents after a grow, or the absence of a CUDAWarning. To force the slow path, reserve a decoy range right after the buffer, as the #2917 tests do.

Issues

#	Issue	Fixed by
#2344	leak after allocate and grow	#2235 (merged)
#2877	spurious warning on the slow grow path	#2917 (the redesign; supersedes #2880)
#2886	unmap without a stream sync	#2917 (supersedes #2889)
#2345	a failed adjacent reservation gets freed	#2917 (supersedes #2237; no reservation is freed by hand)
#2887	a fast-path grow leaves its extension unfreeable	the redesign
#2388	four defects	items 1, 2 and 3 by #2917 (no rollback remap, live fast path, no finalizer warning); item 4 by #2418
#2907	grown buffers record an unaligned size	the redesign
#2908	a second grow of a grown buffer fails	the redesign
#2694	host_numa allocations fail	option validation
#2909	modify_allocation(config=) persists on the resource	option validation, with #1300
#2910	options and location validation	option validation
#1300	xfail review of the policy test	option validation, same question as the config= item

#2882 and #2884 were duplicates of #2344 and #2235 and are closed. The feature requests #2057 (multicast objects) and #2358 (logical endpoints) are not part of this epic; they depend on the redesign, because it changes how VMM buffers own their mappings.

Pull requests

PR	Status
#2235	Merged. Every open VMM PR needs a rebase on main to pick up its Transaction.on_exit changes
#2880	Superseded by #2917, which rewrites the module and fixes #2877 and #2388 items 1 and 3 itself; closes when #2917 merges
#2889	Superseded by #2917, which fixes #2886; closes when #2917 merges
#2237	Superseded by #2917, which fixes #2345 and makes the fast path live; closes when #2917 merges
#2407	Closing: a subset of #2237, and the fast path must stay off until #2887 is fixed
#2440	Closing: #2880 removes the code it edits
#2890	Closing: one call cannot free two reservations; the fix is the redesign

The tables are updated as PRs merge or close.