diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index ea1e2455c6f..5b693833a5c 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -500,8 +500,7 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> size : int The size in bytes of the buffer to allocate. stream : Stream, optional - Keyword-only. Unused because virtual memory operations are - synchronous. + Keyword-only. Validated when provided but otherwise unused. Returns ------- @@ -587,16 +586,17 @@ def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | Grap size : int The size in bytes of the memory to deallocate. stream : Stream, optional - Keyword-only. Unused because virtual memory operations are - synchronous. + Keyword-only. If provided for a device-located resource, + ``stream.sync()`` is called before the virtual memory is unmapped. """ ptr = 0 if ptr is None else int(ptr) - if stream is not None: + if stream is not None and self.is_device_accessible: from cuda.core._stream import Stream_accept - Stream_accept(stream) - # The mapping owns the allocation; unmapping frees its backing memory when no external references remain. + Stream_accept(stream).sync() + result, handle = driver.cuMemRetainAllocationHandle(ptr) + raise_if_driver_error(result) (result,) = driver.cuMemUnmap(ptr, size) raise_if_driver_error(result) (result,) = driver.cuMemAddressFree(ptr, size) 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 31598150775..d795b8253dc 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -34,12 +34,10 @@ New features Fixes and enhancements ---------------------- -- A :class:`Program` that is collected as part of a reference cycle no longer - raises ``AttributeError`` from its destructor, and the source file it wrote - for a ``debug`` or ``lineinfo`` build is removed. The cyclic collector - clears object attributes before the destructor runs, and the destructor - read the program options to locate that file. - (`#2876 `__) +- :meth:`VirtualMemoryResource.deallocate` now synchronizes a supplied stream + before unmapping device-located virtual memory. Host-located resources skip + this synchronization because they do not retain a device context. + (`#2889 `__) - ``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: diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index c20e263262b..eb0680e4ef0 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1538,6 +1538,21 @@ def allocate_and_close(): assert baseline - free < aligned_size +@pytest.mark.parametrize("location_type", ["device", "host"]) +def test_vmm_deallocate_with_stream_has_no_warning(init_cuda, location_type): + """VMM buffers close cleanly with an explicit stream for both locations.""" + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + device.set_current() + stream = device.create_stream() + mr = VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type)) + + with assert_no_cuda_warning(): + buffer = mr.allocate(2 * 1024 * 1024) + buffer.close(stream) + + def test_vmm_allocator_rdma_unsupported_exception(): """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it.