From 7f20e80c6f536e5ca7b38e3495dc6bbb17d7ab57 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 15 Oct 2026 16:08:50 -0700 Subject: [PATCH 1/3] reject invalid cuda stream --- cpp/include/rapidsmpf/memory/buffer.hpp | 1 + .../rapidsmpf/memory/buffer_resource.hpp | 16 ++++++------ cpp/include/rapidsmpf/memory/host_buffer.hpp | 14 +++++++++- cpp/src/memory/buffer.cpp | 5 ++++ cpp/src/memory/buffer_resource.cpp | 20 +++++++------- cpp/src/memory/host_buffer.cpp | 26 +++++++++++++++++-- cpp/tests/test_buffer.cpp | 15 +++++++++++ cpp/tests/test_buffer_resource.cpp | 14 ++++++++++ cpp/tests/test_host_buffer.cpp | 19 ++++++++++++++ 9 files changed, 110 insertions(+), 20 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/buffer.hpp b/cpp/include/rapidsmpf/memory/buffer.hpp index 42b996669..ecb97b834 100644 --- a/cpp/include/rapidsmpf/memory/buffer.hpp +++ b/cpp/include/rapidsmpf/memory/buffer.hpp @@ -254,6 +254,7 @@ class Buffer { * @param new_stream The new CUDA stream. * * @throws std::logic_error If the buffer is locked. + * @throws std::invalid_argument If @p new_stream is `cuda::invalid_stream`. * * @code{.cpp} * // Example: merge buffers from different streams onto a single stream. diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index d49b1f5f2..41bf70989 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -483,6 +483,7 @@ class BufferResource : public std::enable_shared_from_this { * * @throws std::invalid_argument if the memory type does not match the reservation. * @throws rapidsmpf::reservation_error if `size` exceeds the size of the reservation. + * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`. */ std::unique_ptr make_buffer( std::size_t size, cuda::stream_ref stream, MemoryReservation& reservation @@ -497,6 +498,8 @@ class BufferResource : public std::enable_shared_from_this { * @param stream CUDA stream to use for device allocations. * @param reservation The memory reservation to consume for the allocation. * @return A unique pointer to the allocated Buffer. + * + * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`. */ std::unique_ptr make_buffer( cuda::stream_ref stream, MemoryReservation&& reservation @@ -512,13 +515,10 @@ class BufferResource : public std::enable_shared_from_this { * the Buffer is created with `MemoryType::PINNED_HOST`; otherwise it is * created with `MemoryType::DEVICE`. * - * If @p stream differs from the device buffer's current stream: - * - @p stream is synchronized with the device buffer's current stream, and - * - the device buffer's current stream is updated to @p stream. - * - * If @p data is empty and has no stream (e.g. a default-constructed - * `rmm::device_buffer`), no synchronization is performed and its stream is set to - * @p stream. + * @p stream is ordered after any work pending on the device buffer's current + * stream, and the device buffer's stream is updated to @p stream. If @p data is + * empty (e.g. a default-constructed `rmm::device_buffer`, which has no stream), it + * owns no memory, so no synchronization is performed. * * @param data Unique pointer to the device or pinned host buffer. * @param stream CUDA stream associated with the new Buffer. Use or synchronize with @@ -528,7 +528,7 @@ class BufferResource : public std::enable_shared_from_this { * @return Unique pointer to the resulting Buffer. * * @throws std::invalid_argument If @p spill_token is set and @p data is not - * host-accessible. + * host-accessible, or if @p stream is `cuda::invalid_stream`. */ std::unique_ptr move( std::unique_ptr data, diff --git a/cpp/include/rapidsmpf/memory/host_buffer.hpp b/cpp/include/rapidsmpf/memory/host_buffer.hpp index f83b090f5..16ff10a35 100644 --- a/cpp/include/rapidsmpf/memory/host_buffer.hpp +++ b/cpp/include/rapidsmpf/memory/host_buffer.hpp @@ -42,6 +42,8 @@ class HostBuffer { * so the buffer shares ownership of the resource (e.g. bumps the refcount * when constructed from a shared-ownership resource); an implicit * conversion from `rmm::host_async_resource_ref` is also supported. + * + * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`. */ HostBuffer( std::size_t size, @@ -155,6 +157,8 @@ class HostBuffer { * @param mr Host memory resource used to allocate the buffer. * * @return A new `HostBuffer` containing a copy of `data`. + * + * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`. */ static HostBuffer from_uint8_vector( std::vector const& data, @@ -173,6 +177,8 @@ class HostBuffer { * @param stream CUDA stream to associate with this buffer. * * @return A new `HostBuffer` owning the vector's memory. + * + * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`. */ static HostBuffer from_owned_vector( std::vector&& data, cuda::stream_ref stream @@ -185,6 +191,9 @@ class HostBuffer { * the device buffer contains host-accessible memory (e.g., pinned host memory * allocated via a managed or pinned memory resource). * + * @p stream is ordered after any work pending on the device buffer's stream, and + * the memory is deallocated on the `HostBuffer`'s stream. + * * @warning The caller is responsible for ensuring the device buffer's memory is * host-accessible. Using this with non-host-accessible device memory will result * in a std::invalid_argument exception. @@ -194,7 +203,8 @@ class HostBuffer { * * @return A new `HostBuffer` owning the device buffer's memory. * - * @throws std::invalid_argument if @p pinned_host_buffer is null. + * @throws std::invalid_argument if @p pinned_host_buffer is null or @p stream is + * `cuda::invalid_stream`. * @throws std::logic_error if @p pinned_host_buffer is not of pinned memory host * memory type (see warning for details). * @@ -214,6 +224,8 @@ class HostBuffer { * @param deallocate_fn Callable invoked with the current stream to release the * underlying memory. It captures all resources needed for deallocation (e.g. * memory resource, raw pointer, size). + * + * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`. */ HostBuffer( std::span span, diff --git a/cpp/src/memory/buffer.cpp b/cpp/src/memory/buffer.cpp index 37d9d09b6..cfde6b9cc 100644 --- a/cpp/src/memory/buffer.cpp +++ b/cpp/src/memory/buffer.cpp @@ -146,6 +146,11 @@ Buffer::DiskBufferT Buffer::release_disk_buffer() { void Buffer::rebind_stream(cuda::stream_ref new_stream) { throw_if_locked(); + RAPIDSMPF_EXPECTS( + new_stream != cuda::invalid_stream, + "new_stream must be valid", + std::invalid_argument + ); if (new_stream.get() == stream_.get()) { return; } diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 42bc75a08..ba071dfe8 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -276,6 +276,9 @@ std::size_t BufferResource::release(MemoryReservation& reservation, std::size_t std::unique_ptr BufferResource::make_buffer( std::size_t size, cuda::stream_ref stream, MemoryReservation& reservation ) { + RAPIDSMPF_EXPECTS( + stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument + ); auto const mem_type = reservation.mem_type_; // A disk buffer is created with mkstemp, which never touches the stream, so a // stream-ordered timing would measure stream backlog rather than the creation. @@ -328,21 +331,20 @@ std::unique_ptr BufferResource::move( cuda::stream_ref stream, std::shared_ptr spill_token ) { + RAPIDSMPF_EXPECTS( + stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument + ); RAPIDSMPF_EXPECTS( spill_token == nullptr || is_host_accessible(data->memory_resource()), "a spill token cannot be attached to device memory", std::invalid_argument ); - cuda::stream_ref upstream = data->stream(); - if (upstream.get() != stream.get()) { - // Default-constructed RMM buffers have no stream or work to join. - bool const empty_without_stream = - data->capacity() == 0 && upstream == cuda::invalid_stream; - if (!empty_without_stream) { - cuda_stream_join(stream, upstream); - } - data->set_stream(stream); + // An empty buffer owns no memory, so there is no work to join. This also covers + // default-constructed RMM buffers, which have no stream. + if (data->capacity() > 0) { + cuda_stream_join(stream, data->stream()); } + data->set_stream(stream); if (is_host_accessible(data->memory_resource())) { auto pinned_host_buffer = std::make_unique( diff --git a/cpp/src/memory/host_buffer.cpp b/cpp/src/memory/host_buffer.cpp index e908d32c9..4a31494be 100644 --- a/cpp/src/memory/host_buffer.cpp +++ b/cpp/src/memory/host_buffer.cpp @@ -7,6 +7,7 @@ #include +#include #include #include #include @@ -20,6 +21,9 @@ HostBuffer::HostBuffer( cuda::mr::any_resource mr ) : stream_{stream} { + RAPIDSMPF_EXPECTS( + stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument + ); if (size > 0) { auto* ptr = static_cast( mr.allocate(stream_, size, alignof(::cuda::std::max_align_t)) @@ -36,7 +40,11 @@ HostBuffer::HostBuffer( cuda::stream_ref stream, std::function deallocate_fn ) - : stream_{stream}, span_{span}, deallocate_fn_{std::move(deallocate_fn)} {} + : stream_{stream}, span_{span}, deallocate_fn_{std::move(deallocate_fn)} { + RAPIDSMPF_EXPECTS( + stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument + ); +} void HostBuffer::deallocate_async() noexcept { if (!span_.empty() && deallocate_fn_) { @@ -136,6 +144,9 @@ HostBuffer HostBuffer::from_owned_vector( HostBuffer HostBuffer::from_rmm_device_buffer( std::unique_ptr pinned_host_buffer, cuda::stream_ref stream ) { + RAPIDSMPF_EXPECTS( + stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument + ); RAPIDSMPF_EXPECTS( pinned_host_buffer != nullptr, "pinned_host_buffer must not be null", @@ -148,6 +159,14 @@ HostBuffer HostBuffer::from_rmm_device_buffer( std::logic_error ); + // Order `stream` after pending work on the buffer's stream. An empty buffer owns + // no memory, so there is no work to join. This also covers default-constructed + // RMM buffers, which have no stream. + if (pinned_host_buffer->capacity() > 0) { + cuda_stream_join(stream, pinned_host_buffer->stream()); + } + pinned_host_buffer->set_stream(stream); + // Wrap in shared_ptr so the lambda is copyable (required by std::function). auto shared_db = std::make_shared(std::move(*pinned_host_buffer)); std::span span{ @@ -157,7 +176,10 @@ HostBuffer HostBuffer::from_rmm_device_buffer( return HostBuffer{ std::move(span), stream, - [shared_db_ = std::move(shared_db)](cuda::stream_ref) mutable { + // Deallocate on the HostBuffer's current stream, which `set_stream()` may + // have changed since construction. + [shared_db_ = std::move(shared_db)](cuda::stream_ref s) mutable { + shared_db_->set_stream(s); shared_db_.reset(); } }; diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index c1959a2a0..8c8ad55ee 100644 --- a/cpp/tests/test_buffer.cpp +++ b/cpp/tests/test_buffer.cpp @@ -238,3 +238,18 @@ TEST_P(BufferRebindStreamTest, ThrowsWhenLocked) { EXPECT_NO_THROW(buffer->rebind_stream(stream2)); EXPECT_EQ(buffer->stream().get(), stream2.get()); } + +TEST(BufferRebindStream, ThrowsOnInvalidStream) { + StreamPool stream_pool{1}; + auto stream = stream_pool.get_stream(); + auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); + + auto [reserve, _] = br->reserve(MemoryType::DEVICE, 1_KiB, AllowOverbooking::YES); + auto buffer = br->make_buffer(1_KiB, stream, reserve); + + EXPECT_THROW( + buffer->rebind_stream(cuda::stream_ref{cuda::invalid_stream}), + std::invalid_argument + ); + EXPECT_EQ(buffer->stream().get(), stream.get()); +} diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index dca11f142..a46521710 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -79,6 +79,20 @@ TEST(BufferResource, MoveDefaultConstructedDeviceBuffer) { EXPECT_NO_THROW(buffer->latest_write_event().host_wait()); } +TEST(BufferResource, InvalidStreamThrows) { + cuda::stream_ref const invalid{cuda::invalid_stream}; + auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); + + auto [reservation, _] = br->reserve(MemoryType::DEVICE, 0, AllowOverbooking::YES); + EXPECT_THROW( + std::ignore = br->make_buffer(0, invalid, reservation), std::invalid_argument + ); + EXPECT_THROW( + std::ignore = br->move(std::make_unique(), invalid), + std::invalid_argument + ); +} + TEST(BufferResource, ReservationOverbooking) { // Create a buffer resource that always reports 10 KiB of available device memory. auto br = BufferResource::create( diff --git a/cpp/tests/test_host_buffer.cpp b/cpp/tests/test_host_buffer.cpp index e86ba454c..72deeae78 100644 --- a/cpp/tests/test_host_buffer.cpp +++ b/cpp/tests/test_host_buffer.cpp @@ -269,6 +269,25 @@ TEST(PinnedResource, transient_mr) { EXPECT_NO_THROW(test_buffer(std::move(buffer), source_data)); } +TEST(HostBuffer, invalid_stream_throws) { + cuda::stream_ref const invalid{cuda::invalid_stream}; + std::vector const data(8); + + EXPECT_THROW( + std::ignore = rapidsmpf::HostBuffer::from_owned_vector( + std::vector(data), invalid + ), + std::invalid_argument + ); + auto br = + rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); + EXPECT_THROW( + std::ignore = + rapidsmpf::HostBuffer::from_uint8_vector(data, invalid, br->host_mr()), + std::invalid_argument + ); +} + namespace { /// Discover the actual pool size the driver creates when a small max is requested. From b9fcf58d8aeb0a123e9c3a85329b18a17a8e95f4 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Fri, 16 Oct 2026 09:28:57 -0700 Subject: [PATCH 2/3] addressing PR comments --- cpp/include/rapidsmpf/memory/host_buffer.hpp | 2 ++ cpp/src/memory/host_buffer.cpp | 3 +++ cpp/tests/test_host_buffer.cpp | 18 ++++++++++++++++++ 3 files changed, 23 insertions(+) diff --git a/cpp/include/rapidsmpf/memory/host_buffer.hpp b/cpp/include/rapidsmpf/memory/host_buffer.hpp index 16ff10a35..5c398ebca 100644 --- a/cpp/include/rapidsmpf/memory/host_buffer.hpp +++ b/cpp/include/rapidsmpf/memory/host_buffer.hpp @@ -143,6 +143,8 @@ class HostBuffer { * or is otherwise properly synchronized before switching streams. * * @param stream The new CUDA stream. + * + * @throws std::invalid_argument if @p stream is `cuda::invalid_stream`. */ void set_stream(cuda::stream_ref stream); diff --git a/cpp/src/memory/host_buffer.cpp b/cpp/src/memory/host_buffer.cpp index 4a31494be..5cb83ff9e 100644 --- a/cpp/src/memory/host_buffer.cpp +++ b/cpp/src/memory/host_buffer.cpp @@ -98,6 +98,9 @@ std::byte const* HostBuffer::data() const noexcept { } void HostBuffer::set_stream(cuda::stream_ref new_stream) { + RAPIDSMPF_EXPECTS( + new_stream != cuda::invalid_stream, "stream must be valid", std::invalid_argument + ); stream_ = new_stream; } diff --git a/cpp/tests/test_host_buffer.cpp b/cpp/tests/test_host_buffer.cpp index 72deeae78..f2d03d76e 100644 --- a/cpp/tests/test_host_buffer.cpp +++ b/cpp/tests/test_host_buffer.cpp @@ -270,6 +270,7 @@ TEST(PinnedResource, transient_mr) { } TEST(HostBuffer, invalid_stream_throws) { + cuda::stream_ref const valid{cudaStreamLegacy}; cuda::stream_ref const invalid{cuda::invalid_stream}; std::vector const data(8); @@ -286,6 +287,23 @@ TEST(HostBuffer, invalid_stream_throws) { rapidsmpf::HostBuffer::from_uint8_vector(data, invalid, br->host_mr()), std::invalid_argument ); + EXPECT_THROW( + rapidsmpf::HostBuffer(data.size(), invalid, br->host_mr()), std::invalid_argument + ); + + auto buffer = + rapidsmpf::HostBuffer::from_owned_vector(std::vector(data), valid); + EXPECT_THROW(buffer.set_stream(invalid), std::invalid_argument); + + if (!rapidsmpf::is_pinned_memory_resources_supported()) { + GTEST_SKIP() << "PinnedMemoryResource is not supported"; + } + EXPECT_THROW( + std::ignore = rapidsmpf::HostBuffer::from_rmm_device_buffer( + std::make_unique(8, valid, br->pinned_mr()), invalid + ), + std::invalid_argument + ); } namespace { From 86798db660152caea3d9c1b55845ec1b3c8764f1 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Fri, 16 Oct 2026 13:04:26 -0700 Subject: [PATCH 3/3] addressing PR comments --- cpp/tests/test_host_buffer.cpp | 58 ++++++++++++++++++++++++++++------ 1 file changed, 48 insertions(+), 10 deletions(-) diff --git a/cpp/tests/test_host_buffer.cpp b/cpp/tests/test_host_buffer.cpp index f2d03d76e..cfee2a820 100644 --- a/cpp/tests/test_host_buffer.cpp +++ b/cpp/tests/test_host_buffer.cpp @@ -15,12 +15,14 @@ #include +#include #include #include #include #include #include +#include #include #include #include @@ -269,6 +271,39 @@ TEST(PinnedResource, transient_mr) { EXPECT_NO_THROW(test_buffer(std::move(buffer), source_data)); } +TEST(PinnedResource, from_rmm_device_buffer_stream_ordering) { + if (!rapidsmpf::is_pinned_memory_resources_supported()) { + GTEST_SKIP() << "PinnedMemoryResource is not supported"; + } + auto br = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + ); + auto mr = br->try_pinned_mr(); + rmm::cuda_stream stream_a; + rmm::cuda_stream stream_b; + rmm::cuda_stream stream_c; + constexpr std::size_t size = 1 << 20; + + // Fill the pinned buffer asynchronously on stream A, without synchronizing. + auto pinned_host_buffer = std::make_unique(size, stream_a, *mr); + RAPIDSMPF_CUDA_TRY( + cudaMemsetAsync(pinned_host_buffer->data(), 0xAB, size, stream_a.value()) + ); + + // The new stream B must be ordered after the fill on A. + auto buffer = rapidsmpf::HostBuffer::from_rmm_device_buffer( + std::move(pinned_host_buffer), stream_b + ); + EXPECT_EQ(buffer.stream().get(), stream_b.value()); + EXPECT_EQ(buffer.copy_to_uint8_vector(), std::vector(size, 0xAB)); + + buffer.set_stream(stream_c); // Switch to stream C; the memory must be freed on C. + EXPECT_EQ(buffer.stream().get(), stream_c.value()); + buffer.deallocate_async(); + EXPECT_TRUE(buffer.empty()); + stream_c.synchronize(); +} + TEST(HostBuffer, invalid_stream_throws) { cuda::stream_ref const valid{cudaStreamLegacy}; cuda::stream_ref const invalid{cuda::invalid_stream}; @@ -280,8 +315,12 @@ TEST(HostBuffer, invalid_stream_throws) { ), std::invalid_argument ); - auto br = - rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); + auto br = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), + rapidsmpf::is_pinned_memory_resources_supported() + ? std::make_optional() + : std::nullopt + ); EXPECT_THROW( std::ignore = rapidsmpf::HostBuffer::from_uint8_vector(data, invalid, br->host_mr()), @@ -295,15 +334,14 @@ TEST(HostBuffer, invalid_stream_throws) { rapidsmpf::HostBuffer::from_owned_vector(std::vector(data), valid); EXPECT_THROW(buffer.set_stream(invalid), std::invalid_argument); - if (!rapidsmpf::is_pinned_memory_resources_supported()) { - GTEST_SKIP() << "PinnedMemoryResource is not supported"; + if (auto pinned_mr = br->try_pinned_mr()) { + EXPECT_THROW( + std::ignore = rapidsmpf::HostBuffer::from_rmm_device_buffer( + std::make_unique(8, valid, *pinned_mr), invalid + ), + std::invalid_argument + ); } - EXPECT_THROW( - std::ignore = rapidsmpf::HostBuffer::from_rmm_device_buffer( - std::make_unique(8, valid, br->pinned_mr()), invalid - ), - std::invalid_argument - ); } namespace {