diff --git a/cuda_core/cuda/core/_memory/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index 4d8bd657968..5bf6511fa79 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -127,6 +127,41 @@ class Buffer: stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional The stream object to use for asynchronous deallocation. If None, the deallocation stream stored in the handle is used. + + See Also + -------- + set_deallocation_stream + Change the deallocation stream without closing the buffer. + """ + + def set_deallocation_stream(self, stream: Stream | GraphBuilder) -> None: + """Change the stream that orders this buffer's eventual deallocation. + + The buffer remains open and usable. A later :meth:`close` without a + stream, garbage collection, or release of the final retained device + pointer handle uses the replacement stream. + + This method does not synchronize streams or establish dependencies. + The caller must ensure that allocation and all accesses are ordered + before the deallocation on ``stream``. + + Parameters + ---------- + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder` + The stream to use for eventual asynchronous deallocation. + + Raises + ------ + RuntimeError + If the buffer is already closed, or if a default-stream token + cannot be bound because no CUDA context is current. + TypeError + If ``stream`` is ``None`` or is not an accepted stream object. + + Notes + ----- + Synchronizing concurrent mutation and destruction of the same buffer + is the caller's responsibility. """ def __enter__(self): diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 2f9d3bce218..552bb0bcc8d 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -252,7 +252,7 @@ cdef class Buffer: # The parent process's stream is not portable across processes, so the # pickle path cannot thread an explicit stream through. Seed the # imported buffer's deallocation with the current context's default - # stream; the receiver can override via buffer.close(stream). + # stream; the receiver can override it before or during close. return Buffer.from_ipc_descriptor(mr, ipc_descriptor, stream=default_stream()) def __reduce__(self) -> tuple[object, ...]: @@ -347,9 +347,45 @@ cdef class Buffer: stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional The stream object to use for asynchronous deallocation. If None, the deallocation stream stored in the handle is used. + + See Also + -------- + set_deallocation_stream + Change the deallocation stream without closing the buffer. """ Buffer_close(self, stream) + def set_deallocation_stream(self, stream: Stream | GraphBuilder) -> None: + """Change the stream that orders this buffer's eventual deallocation. + + The buffer remains open and usable. A later :meth:`close` without a + stream, garbage collection, or release of the final retained device + pointer handle uses the replacement stream. + + This method does not synchronize streams or establish dependencies. + The caller must ensure that allocation and all accesses are ordered + before the deallocation on ``stream``. + + Parameters + ---------- + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder` + The stream to use for eventual asynchronous deallocation. + + Raises + ------ + RuntimeError + If the buffer is already closed, or if a default-stream token + cannot be bound because no CUDA context is current. + TypeError + If ``stream`` is ``None`` or is not an accepted stream object. + + Notes + ----- + Synchronizing concurrent mutation and destruction of the same buffer + is the caller's responsibility. + """ + Buffer_set_deallocation_stream(self, stream) + def __enter__(self): return self @@ -707,15 +743,21 @@ cdef tuple Buffer_coerce_batch(object buffers, str what, str single_hint): return tuple(out) +cdef inline void Buffer_set_deallocation_stream(Buffer self, object stream): + """Validate and replace a live buffer's deallocation recipe.""" + if not self._h_ptr: + raise RuntimeError("Cannot set the deallocation stream on a closed Buffer") + cdef Stream s = Stream_accept(stream) + _apply_deallocation_stream(self._h_ptr, s._h_stream) + + cdef inline void Buffer_close(Buffer self, object stream): """Close a buffer, freeing its memory.""" - cdef Stream s if not self._h_ptr: return # Update deallocation stream if provided if stream is not None: - s = Stream_accept(stream) - _apply_deallocation_stream(self._h_ptr, s._h_stream) + Buffer_set_deallocation_stream(self, stream) # Reset handle - RAII deleter will free the memory (and release owner ref in C++) self._h_ptr.reset() self._size = 0 diff --git a/cuda_core/docs/source/api.rst b/cuda_core/docs/source/api.rst index 38ff695ace5..5ee34d34f54 100644 --- a/cuda_core/docs/source/api.rst +++ b/cuda_core/docs/source/api.rst @@ -77,6 +77,14 @@ Memory management ManagedMemoryResourceOptions VirtualMemoryResourceOptions +A :class:`Buffer` records the stream that will order its eventual deallocation. +Use :meth:`Buffer.set_deallocation_stream` to replace that stream without +closing the buffer. Changing the recorded stream does not synchronize streams; +the caller must order allocation and every access before the deallocation, +using events or other CUDA synchronization mechanisms as needed. See +:cuda-core-example:`buffer_deallocation_stream.py ` +for a complete example. + CUDA compilation toolchain -------------------------- diff --git a/cuda_core/docs/source/examples.rst b/cuda_core/docs/source/examples.rst index cf13961c6dc..f5ae0c10300 100644 --- a/cuda_core/docs/source/examples.rst +++ b/cuda_core/docs/source/examples.rst @@ -39,6 +39,13 @@ Linking and graphs - :cuda-core-example:`cuda_graphs.py ` captures and replays a multi-kernel CUDA graph to reduce launch overhead. +Memory management +----------------- + +- :cuda-core-example:`buffer_deallocation_stream.py ` + transfers a buffer between streams and safely changes the stream that orders + its deallocation. + Interoperability and memory access ---------------------------------- diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index 9e96d54b3d2..4e6f29ba9a6 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -35,6 +35,10 @@ Fixes and enhancements free recipe can pin that context. (`#2497 `__) +- Added :meth:`Buffer.set_deallocation_stream` to change the stream that orders + a buffer's eventual deallocation without closing the buffer. + (`#2600 `__) + - Explicit calls to ``deallocate()`` on pool-backed memory resources and :class:`GraphMemoryResource` now propagate errors from the underlying CUDA free operation. Previously, these errors could be suppressed. Automatic diff --git a/cuda_core/examples/buffer_deallocation_stream.py b/cuda_core/examples/buffer_deallocation_stream.py new file mode 100644 index 00000000000..49579040590 --- /dev/null +++ b/cuda_core/examples/buffer_deallocation_stream.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# ################################################################################ +# +# This example transfers a buffer from a producer stream to a consumer stream. +# An event orders the consumer after the producer. The buffer then records the +# consumer stream for its eventual deallocation. +# +# ################################################################################ + +# /// script +# dependencies = ["cuda_bindings", "cuda_core"] +# /// + +import ctypes + +from cuda.core import Device, LegacyPinnedMemoryResource + + +def produce_data(device, stream, size, value): + """Allocate and fill a buffer on the producer stream.""" + buffer = device.allocate(size, stream=stream) + buffer.fill(value, stream=stream) + ready = stream.record() + return buffer, ready + + +def consume_data(buffer, ready, output, stream): + """Submit consumer work and transfer the deallocation stream.""" + stream.wait(ready) + buffer.set_deallocation_stream(stream) + buffer.copy_to(output, stream=stream) + + +def main(): + device = Device() + device.set_current() + producer_stream = device.create_stream() + consumer_stream = device.create_stream() + pinned_mr = LegacyPinnedMemoryResource() + + size = 4096 + value = 42 + buffer = None + ready = None + output = None + + try: + output = pinned_mr.allocate(size) + buffer, ready = produce_data(device, producer_stream, size, value) + consume_data(buffer, ready, output, consumer_stream) + + # No stream argument is needed. The buffer now records consumer_stream. + # The free operation runs after the copy on that stream. + buffer.close() + buffer = None + consumer_stream.sync() + + result = ctypes.string_at(int(output.handle), output.size) + assert result == bytes([value]) * size + print("Buffer deallocation stream transfer completed.") + finally: + if buffer is not None: + buffer.close() + if output is not None: + output.close() + if ready is not None: + ready.close() + consumer_stream.close() + producer_stream.close() + + +if __name__ == "__main__": + main() diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 4e781434fb8..3fdbe98885f 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -516,6 +516,62 @@ def test_mr_deallocate_receives_stream(): assert telemetry["deallocations"][-1]["stream"].handle == stream.handle +@pytest.mark.agent_authored(model="gpt-5.6") +@pytest.mark.parametrize( + ("configuration", "destruction"), + [ + ("initialization", "close"), + ("initialization", "gc"), + ("setter", "close"), + ("setter", "gc"), + ("close", "close"), + ], +) +def test_buffer_deallocation_stream_configuration_paths(configuration, destruction): + """Creation, mutation, and close overrides use the requested stream.""" + import gc + + device = Device() + device.set_current() + initial_stream = device.create_stream() + target_stream = device.create_stream() + CapturingMR, telemetry = make_instrumented_memory_resource(record_streams=True) + mr = CapturingMR(device) + + stream = target_stream if configuration == "initialization" else initial_stream + buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream) + if configuration == "setter": + handle = buf.handle + buf.set_deallocation_stream(target_stream) + assert buf.handle == handle + assert buf.size == 1024 + + if destruction == "close": + buf.close(stream=target_stream if configuration == "close" else None) + else: + del buf + gc.collect() + + assert len(telemetry["deallocations"]) == 1 + assert telemetry["deallocations"][0]["stream"].handle == target_stream.handle + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_set_deallocation_stream_rejects_none_and_closed_buffer(): + device = Device() + device.set_current() + stream = device.create_stream() + mr = StubMemoryResource(device) + buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream) + + with pytest.raises(TypeError, match="stream is required"): + buf.set_deallocation_stream(None) + + buf.close() + with pytest.raises(RuntimeError, match="closed Buffer"): + buf.set_deallocation_stream(stream) + + @pytest.mark.parametrize("buffer_type", [Buffer, ManagedBuffer]) def test_from_handle_mr_records_default_stream(buffer_type): """When a Buffer/ManagedBuffer is minted via :meth:`from_handle` with ``mr``