diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index ea1e2455c6f..b0d50e628fd 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -371,10 +371,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. @@ -385,8 +386,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) @@ -396,28 +398,17 @@ 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 mappings + # keep the memory alive, so the retained reference is dropped again on + # either outcome. result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle) raise_if_driver_error(result) trans.on_exit(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.on_failure(_remap_old) - - # Remap the old physical memory to the new VA range (aligned previous size) + # 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) @@ -449,12 +440,10 @@ def _remap_old() -> None: # All succeeded, cancel undo actions trans.commit() - # Free the old VA range (aligned previous size) - (res2,) = driver.cuMemAddressFree(int(buf.handle), aligned_prev_size) - raise_if_driver_error(res2) - - # 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 and frees its reservation. 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) 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 5a062932ccd..13c1c641356 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -41,6 +41,14 @@ Fixes and enhancements read the program options to locate that file. (`#2876 `__) +- 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 c20e263262b..b4e5a4b9ebc 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1538,6 +1538,48 @@ def allocate_and_close(): assert baseline - free < aligned_size +@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(): """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it.