cuda.core: release VMM allocation handles once they are mapped

cuda_core/cuda/core/_memory/_virtual_memory_resource.py

@@ -328,31 +328,34 @@ def _grow_allocation_fast_path(
             Buffer: The same buffer object with its size updated to `new_size`.
         """
         with Transaction() as trans:
-            # Create new physical memory for the additional size
+            # The caller reserved the extension range; free it unless the grow commits.
             trans.append(
                 lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemAddressFree(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)
-            # Register undo for creation
-            trans.append(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.append(
-                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))
+            try:
+                # 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.append(
+                    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()
+                # All succeeded, cancel undo actions
+                trans.commit()
+            finally:
+                # The mapping holds its own reference; drop the one from cuMemCreate
+                # on success and failure alike (#2882).
+                raise_if_driver_error(driver.cuMemRelease(new_handle)[0])
 
         # Update the buffer size (pointer stays the same). `Buffer.size` has
         # no public setter, so this reaches into the private attribute.
@@ -372,10 +375,11 @@ def _grow_allocation_slow_path(
         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.
+        This function reserves a new, larger virtual address (VA) range, maps the old
+        physical memory to the beginning of the new VA range as a second mapping, creates
+        and maps new physical memory for the additional size, sets access permissions, and
+        then closes the old buffer, which releases the old VA range through
+        :meth:`deallocate`.
 
         Args:
             buf (Buffer): The buffer to grow.
@@ -386,8 +390,9 @@ def _grow_allocation_slow_path(
             addr_align (int): The required address alignment for the new VA range.
 
         Returns:
-            Buffer: The buffer object updated with the new pointer and size.
+            Buffer: A new buffer for the new VA range. ``buf`` is closed.
         """
+        aligned_prev_size = total_aligned_size - aligned_additional_size
         with Transaction() as trans:
             # Reserve a completely new, larger VA range
             res, new_ptr = driver.cuMemAddressReserve(total_aligned_size, addr_align, 0, 0)
@@ -397,68 +402,57 @@ def _grow_allocation_slow_path(
                 lambda np=new_ptr, s=total_aligned_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0])
             )
 
-            # Get the old allocation handle for remapping
+            # Retain the old allocation handle to map it a second time. The old
+            # mapping keeps the memory alive, so the retained reference is dropped
+            # again right after the new mapping is attempted (#2882).
             result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle)
             raise_if_driver_error(result)
-            # Register undo for old_handle
-            trans.append(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.append(_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)
+            try:
+                # Map the old physical memory to the new VA range (aligned previous size).
+                # The old VA range stays mapped too (virtual aliasing), so the old buffer
+                # is untouched if anything below fails and is released as a whole by
+                # buf.close() once the new mapping is complete.
+                (res,) = driver.cuMemMap(int(new_ptr), aligned_prev_size, 0, old_handle, 0)
+                raise_if_driver_error(res)
+            finally:
+                raise_if_driver_error(driver.cuMemRelease(old_handle)[0])
 
             # Register undo for mapping
             trans.append(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)
+            try:
+                # 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 new physical memory
-            trans.append(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.append(
-                lambda base=int(new_ptr), offs=aligned_prev_size, s=aligned_additional_size: raise_if_driver_error(
-                    driver.cuMemUnmap(base + offs, s)[0]
+                # Register undo for mapping
+                trans.append(
+                    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()
+                # 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)
 
-        # Free the old VA range (aligned previous size)
-        (res2,) = driver.cuMemAddressFree(int(buf.handle), aligned_prev_size)
-        raise_if_driver_error(res2)
+                # All succeeded, cancel undo actions
+                trans.commit()
+            finally:
+                # The mapping holds its own reference; drop the one from cuMemCreate
+                # on success and failure alike (#2882).
+                raise_if_driver_error(driver.cuMemRelease(new_handle)[0])
 
-        # Invalidate the old buffer so its destructor won't try to free again
-        buf._clear()
+        # Release the old VA range through the resource: closing the buffer runs
+        # deallocate(), which unmaps the old range, frees its reservation and
+        # releases the handle reference it retains. The physical memory stays
+        # alive through the new mapping.
+        buf.close()
 
         # Return a new Buffer for the new mapping
         return Buffer.from_handle(ptr=new_ptr, size=new_size, mr=self)
@@ -549,32 +543,37 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) ->
         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)
-            # Register undo for physical memory
-            trans.append(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.append(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)
-            trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0]))
-            raise_if_driver_error(res)
+        # ---- Create physical memory ----
+        res, handle = driver.cuMemCreate(aligned_size, prop, 0)
+        raise_if_driver_error(res)
+        try:
+            with Transaction() as trans:
+                # ---- 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.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemAddressFree(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))
+                # ---- Map physical memory into VA ----
+                (res,) = driver.cuMemMap(ptr, aligned_size, 0, handle, 0)
+                trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0]))
                 raise_if_driver_error(res)
 
-            trans.commit()
+                # ---- 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()
+        finally:
+            # The mapping holds its own reference to the physical allocation, so the
+            # reference returned by cuMemCreate is dropped here whether or not the
+            # mapping succeeded. Keeping it made deallocate() unable to free the
+            # memory: the driver frees an allocation only once every mapping is
+            # unmapped and every handle reference is released (#2882).
+            raise_if_driver_error(driver.cuMemRelease(handle)[0])
 
         # Done — return a Buffer that tracks this VA range
         buf = Buffer.from_handle(ptr=ptr, size=aligned_size, mr=self)

cuda_core/docs/source/release/1.3.0-notes.rst

@@ -34,6 +34,21 @@ New features
 Fixes and enhancements
 ----------------------
 
+- :class:`~_memory.VirtualMemoryResource` now releases the physical allocation
+  handle once it is mapped, so closing a buffer returns its memory to the
+  device. Previously every allocation, and every chunk added by
+  :meth:`~_memory.VirtualMemoryResource.modify_allocation`, stayed allocated
+  until the process exited.
+  (`#2882 <https://github.com/NVIDIA/cuda-python/issues/2882>`__)
+
+- Growing a :class:`~_memory.VirtualMemoryResource` allocation with
+  :meth:`~_memory.VirtualMemoryResource.modify_allocation` when the address
+  range cannot be extended in place no longer reports a spurious
+  :class:`CUDAWarning`. The old buffer is now closed through the resource's
+  :meth:`~_memory.VirtualMemoryResource.deallocate` instead of having its
+  range freed by hand and then released a second time.
+  (`#2877 <https://github.com/NVIDIA/cuda-python/issues/2877>`__)
+
 - ``Graph.__getitem__`` now declares an overload for each node type that has
   an executable view, so type checkers and editors see the precise view type:
   indexing with a :class:`~graph.KernelNode` yields an

cuda_core/tests/test_memory.py

@@ -1465,14 +1465,16 @@ 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,)
 
+    # Runs on success as well: the mapping keeps the memory alive, so the
+    # reference returned by cuMemCreate must be dropped once it is mapped (#2882).
+    def fake_release(handle):
+        calls.append(("release", handle))
+        return (SUCCESS,)
+
     monkeypatch.setattr(driver, "cuMemCreate", fake_create)
     monkeypatch.setattr(driver, "cuMemMap", fake_map)
     monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access)
@@ -1498,11 +1500,100 @@ def __init__(self, size):
     assert result is buf
     assert buf._size == new_size
 
-    # Successful commit: create, map, set access, and no rollback calls.
-    assert [c[0] for c in calls] == ["create", "map", "set_access"]
+    # Successful commit: create, map, set access, release the creation
+    # reference, and no rollback calls.
+    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")
+@pytest.mark.thread_unsafe(reason="measures device-wide free memory")
+def test_vmm_allocator_close_returns_physical_memory(init_cuda):
+    """Closing VMM buffers returns their physical memory to the device (#2882).
+
+    allocate() and both grow paths kept the reference that cuMemCreate returns,
+    so the memory stayed allocated after the range was unmapped and freed.
+    """
+    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"),
+    )
+    chunk = 32 * 1024 * 1024
+    rounds = 8
+
+    def free_memory():
+        device.sync()
+        return handle_return(driver.cuMemGetInfo())[0]
+
+    # Plain allocate and close: the leak was one chunk per round.
+    before = free_memory()
+    for _ in range(rounds):
+        vmm_mr.allocate(chunk).close()
+    retained = before - free_memory()
+    assert retained < rounds * chunk // 2, f"{retained >> 20} MiB still allocated after close"
+
+    # Grow through the slow path (forced by a decoy reservation), then close:
+    # the leak was two chunks per round.
+    before = free_memory()
+    for _ in range(rounds):
+        buf = vmm_mr.allocate(chunk)
+        decoy = handle_return(driver.cuMemAddressReserve(chunk, 0, int(buf.handle) + buf.size, 0))
+        try:
+            grown = vmm_mr.modify_allocation(buf, 2 * chunk)
+        finally:
+            handle_return(driver.cuMemAddressFree(decoy, chunk))
+        grown.close()
+    retained = before - free_memory()
+    assert retained < rounds * chunk, f"{retained >> 20} MiB still allocated after grow and close"
+
+
+@pytest.mark.agent_authored(model="claude-fable-5-1")
+@pytest.mark.thread_unsafe(reason="warning capture is process-global")
+def test_vmm_allocator_grow_allocation_slow_path_closes_old_buffer(init_cuda):
+    """The slow grow path closes the old buffer without a spurious CUDAWarning (#2877).
+
+    It used to free the old VA range by hand and then reset the old buffer's
+    handle, whose deleter called deallocate() on the freed range a second time.
+    """
+    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"),
+    )
+    buf = vmm_mr.allocate(2 * 1024 * 1024)
+    old_ptr, old_size = int(buf.handle), buf.size
+    handle_return(driver.cuMemsetD8(old_ptr, 7, old_size))
+    # The memset is asynchronous and the grow unmaps the old range, so let it
+    # finish first; on WDDM the batched kernel otherwise faults after the unmap.
+    device.sync()
+
+    # Occupy the address range right after buf so the adjacent reservation cannot
+    # be honored and modify_allocation has to take the slow path.
+    decoy = handle_return(driver.cuMemAddressReserve(old_size, 0, old_ptr + old_size, 0))
+    try:
+        with assert_no_cuda_warning():
+            grown = vmm_mr.modify_allocation(buf, 2 * old_size)
+    finally:
+        handle_return(driver.cuMemAddressFree(decoy, old_size))
+
+    assert buf.is_closed
+    assert int(grown.handle) != old_ptr
+    assert grown.size == 2 * old_size
+    # The old contents are reachable through the new mapping.
+    host = (ctypes.c_ubyte * old_size)()
+    handle_return(driver.cuMemcpyDtoH(ctypes.addressof(host), int(grown.handle), old_size))
+    assert bytes(host) == bytes([7]) * old_size
+    grown.close()
 
 
 def test_vmm_allocator_rdma_unsupported_exception():