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.