Avoid freeing failed VMM grow reservations

[fallintoplace]

Summary

• Split the VMM grow adjacent-reservation fallback into failed-reservation and noncontiguous-reservation cases.
• Only call cuMemAddressFree when cuMemAddressReserve succeeded and returned a noncontiguous address.
• Add regression coverage for both fallback paths.

Why

The previous grow path called cuMemAddressFree(new_ptr, ...) even when the adjacent cuMemAddressReserve call failed with a recoverable error. In that case the returned pointer is not an owned reservation, so cleanup can turn a recoverable fallback into another driver error.

Checks

• python3 -m py_compile cuda_core/cuda/core/_memory/_virtual_memory_resource.py cuda_core/tests/test_memory.py
• git diff --check -- cuda_core/cuda/core/_memory/_virtual_memory_resource.py cuda_core/tests/test_memory.py

Not run locally:

• uv run --no-project --with pytest --with cuda-bindings --with cuda-pathfinder --with numpy python -m pytest cuda_core/tests/test_memory.py -k 'vmm_allocator_grow_allocation_does_not_free_failed_adjacent_reservation or vmm_allocator_grow_allocation_frees_noncontiguous_adjacent_reservation'

The isolated test environment on this macOS arm64 host cannot resolve cuda-bindings, because published wheels are available for Linux and Windows but not this platform.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[leofang]

xref: #2235 (comment)

[aryanputta]

Hello @fallintoplace, a heads-up on an interaction between this PR and #2407, since we are both editing this guard.

At head 9bf37ed the split reads:

expected_ptr = int(buf.handle) + aligned_prev_size
if res != driver.CUresult.CUDA_SUCCESS:
    ...
    return self._grow_allocation_slow_path(...)

if new_ptr != expected_ptr:
    (res2,) = driver.cuMemAddressFree(new_ptr, aligned_additional_size)
    ...

new_ptr comes back from cuMemAddressReserve as a CUdeviceptr, and CUdeviceptr defines __int__ but no __eq__ / __richcmp__, so CUdeviceptr(x) != x is always true. new_ptr != expected_ptr therefore always takes the fallback, even when the driver granted the exact contiguous address that was requested. That is defect 2 of #2388, and this restructure carries it into the new shape: the second branch now also calls cuMemAddressFree on a reservation the driver granted correctly, then falls through to the slow path every time.

int(new_ptr) != expected_ptr fixes it. That one-line change is what #2407 does against the pre-split code, so whichever of these lands second will conflict on this line. Happy to rebase #2407 onto this restructure if you would rather land yours first, or you can fold the int() in here and I will close #2407 as redundant. Either way works, I just want to make sure the fast path does not stay unreachable through the refactor.

One note on the new tests: _make_mock_vmm_resource and fake_addr_reserve return stale_ptr as a plain int, so new_ptr != expected_ptr behaves like a normal int comparison under the mock and the real CUdeviceptr behavior is not exercised. Returning driver.CUdeviceptr(...) from the fake would surface it.

[fallintoplace]

Thank you @aryanputta. Fixed in 08583ca by comparing int(new_ptr) with the expected address. I also updated the reservation mocks to return CUdeviceptr and added coverage confirming that an exactly contiguous reservation reaches the fast path without being freed.

[aryanputta]

Thank you, that was quick. I checked 08583ca and it covers everything: int(new_ptr) != expected_ptr in the guard, all three reservation fakes now returning driver.CUdeviceptr(...) so the mocks exercise the real comparison, and test_vmm_allocator_grow_allocation_uses_contiguous_adjacent_reservation covering the dispatch itself rather than only the two branches.

That makes #2407 redundant on both the fix and the coverage. I will close it once this lands rather than right now, only so defect 2 of #2388 does not sit without an open PR while this one waits on vetter validation. Nothing for you to do about that, and I am not planning to touch that guard again.

One piece of #2407 is not duplicated here: a release note bullet, since this changes observable behavior of a stable API. Feel free to lift it if you agree it belongs:

- Growing a :class:`VirtualMemoryResource` allocation now preserves the buffer's
  base pointer when the driver can extend the reservation contiguously.
  Previously an address comparison never matched, so every grow re-reserved and
  remapped the entire range.
  (`#2237 <https://github.com/NVIDIA/cuda-python/pull/2237>`__)

It goes under "Fixes and enhancements" in cuda_core/docs/source/release/1.2.0-notes.rst.

[Andy-Jost]

Thanks for the PR. VMM work is consolidated under #2906, and this PR is on the list to land for cuda.core 1.3.0. Before review:

The failed-reservation half is correct and can land on its own. Keep the fast path off until #2887 is fixed: replace the second if (int(new_ptr) != expected_ptr) with an unconditional cuMemAddressFree of the adjacent reservation followed by the slow-path call, with the comment # TODO(#2887): the fast path leaks its extension; take the slow path until the buffer owns its mappings. That keeps today's behavior and stops freeing a reservation that was never granted (#2345).

Also: move the release note to 1.3.0-notes.rst and reword it for #2345; rebase on main (#2235 has merged); replace the monkeypatched tests with one that runs against the driver: allocate, grow with a decoy reservation placed right after the buffer, and assert the grow succeeds with no CUDAWarning. If you prefer, we can split the PR.

Please keep discussion of the overall plan on #2906.