diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index 2c4f3f6867e..7d05fc87b09 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/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,31 +402,20 @@ 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])) @@ -429,36 +423,36 @@ def _remap_old() -> None: # 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) 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 413465d8b3c..fbb9c128260 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/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 `__) + +- 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 `__) + - ``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 diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index b61a2d5b2f3..d4a4f2e6f7f 100644 --- a/cuda_core/tests/test_memory.py +++ b/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():