From 97340589ddf7c9fa6390d4497b767712ec5f5cb7 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 13 Aug 2026 06:19:48 +0000 Subject: [PATCH 01/32] Make the CUDA stream pool per-thread and per-device The global pool handed overlapping streams to concurrent threads, creating false dependencies between unrelated work. Each thread now owns a pool per device that grows on demand up to LIBCUDF_STREAM_POOL_SIZE, and pools are recycled through a free list when a thread exits so thread churn does not accumulate streams. --- .../cudf/detail/utilities/stream_pool.hpp | 35 ++-- cpp/src/utilities/stream_pool.cpp | 176 +++++++++++++++--- cpp/tests/CMakeLists.txt | 1 + .../utilities_tests/stream_pool_tests.cpp | 165 ++++++++++++++++ 4 files changed, 336 insertions(+), 41 deletions(-) create mode 100644 cpp/tests/utilities_tests/stream_pool_tests.cpp diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index d68527d4c6ef..4dcadd05eedd 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -50,9 +50,13 @@ class cuda_stream_pool { /** * @brief Get a set of `cuda_stream_view` objects from the pool. * - * An attempt is made to ensure that the returned vector does not contain duplicate - * streams, but this cannot be guaranteed if `count` is greater than the value returned by - * `get_stream_pool_size()`. + * The returned streams are distinct unless `count` is greater than the value returned by + * `get_stream_pool_size()`, in which case streams are repeated. + * + * Consecutive calls are served from different streams where the pool is large enough, so a + * nested call generally does not return streams that its caller is already using. This is not + * guaranteed: once the pool has reached its maximum size the assignment wraps around, so a + * request for more than half the pool can overlap with the streams the caller holds. * * This function is thread safe with respect to other calls to the same function. * @@ -62,11 +66,11 @@ class cuda_stream_pool { virtual std::vector get_streams(std::size_t count) = 0; /** - * @brief Get the number of unique stream objects in the pool. + * @brief Get the maximum number of unique stream objects the pool can provide. * * This function is thread safe with respect to other calls to the same function. * - * @return the number of stream objects in the pool + * @return the maximum number of stream objects in the pool */ [[nodiscard]] virtual std::size_t get_stream_pool_size() const = 0; @@ -80,7 +84,14 @@ class cuda_stream_pool { cuda_stream_pool* create_global_cuda_stream_pool(); /** - * @brief Get the global stream pool. + * @brief Get the calling thread's stream pool for the current device. + * + * Each thread has its own pool for each device it uses, so concurrent threads are handed distinct + * streams. Pools are created empty and grow on demand up to a maximum that can be configured with + * the `LIBCUDF_STREAM_POOL_SIZE` environment variable. + * + * The returned streams may be used from any thread, but must not be used after the thread that + * obtained them has exited: a pool is recycled for reuse by another thread at that point. */ cuda_stream_pool& global_cuda_stream_pool(); @@ -88,10 +99,12 @@ cuda_stream_pool& global_cuda_stream_pool(); * @brief Acquire a set of `cuda_stream_view` objects and synchronize them to an event on another * stream. * - * By default an underlying `rmm::cuda_stream_pool` is used to obtain the streams. The only other - * implementation at present is a debugging version that always returns the stream returned by - * `cudf::get_default_stream()`. To use this debugging version, set the environment variable - * `LIBCUDF_USE_DEBUG_STREAM_POOL`. + * By default the calling thread's stream pool is used to obtain the streams, so streams are not + * shared with concurrently forking threads. The only other implementation at present is a debugging + * version that always returns the stream returned by `cudf::get_default_stream()`. To use this + * debugging version, set the environment variable `LIBCUDF_USE_DEBUG_STREAM_POOL`. + * + * The returned streams must not be used after the calling thread has exited. * * Example usage: * @code{.cpp} diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 6c65a4067347..e9cbfc3b63d3 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -3,26 +3,25 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include #include #include #include #include -#include +#include #include #include #include +#include #include namespace cudf::detail { -// TODO: what is a good number here. what's the penalty for making it larger? -// Dave Baranec rule of thumb was max_streams_needed * num_concurrent_threads, -// where num_concurrent_threads was estimated to be 4. so using 32 will allow -// for 8 streams per thread, which should be plenty (decoding will be up to 4 -// kernels when delta_byte_array decoding is added). rmm::cuda_stream_pool -// defaults to 16. +// Maximum number of streams a single thread's pool will create. Sized to cover the largest +// number of streams requested by a single `fork_streams` call in libcudf, which is the number +// of distinct parquet decode kernels (see `decode_kernel_mask`). std::size_t constexpr STREAM_POOL_SIZE = 32; // FIXME: "borrowed" from rmm...remove when this stream pool is moved there @@ -81,6 +80,16 @@ rmm::cuda_device_id get_current_cuda_device() return rmm::cuda_device_id{device_id}; } +/** + * @brief Returns the maximum number of streams a single thread's pool will hold. + */ +std::size_t stream_pool_size() +{ + static std::size_t const size = + std::max(1, getenv_or("LIBCUDF_STREAM_POOL_SIZE", STREAM_POOL_SIZE)); + return size; +} + /** * @brief Returns a cudaEvent_t for the current thread. * @@ -102,29 +111,56 @@ cudaEvent_t event_for_thread() } // namespace /** - * @brief Implementation of `cuda_stream_pool` that wraps an `rmm::cuda_stram_pool`. + * @brief Implementation of `cuda_stream_pool` that creates streams on demand. + * + * Instances are owned by a single thread at a time, so no synchronization is needed. The pool + * never shrinks; it grows to the largest number of streams requested so far, up to + * `stream_pool_size()`. */ -class rmm_cuda_stream_pool : public cuda_stream_pool { - rmm::cuda_stream_pool _pool; +class growing_cuda_stream_pool : public cuda_stream_pool { + std::vector _streams; + std::size_t _next_stream{0}; + + /** + * @brief Creates streams until the pool can serve `count` streams with room to spare. + * + * Twice the requested count is created so that consecutive requests are served from different + * streams. Nested requests, such as decompression forking streams while its caller is using + * forked streams of its own, then avoid colliding with the streams they are nested inside, + * except where the rotation wraps around a pool that has reached `stream_pool_size()`. + */ + void grow_to(std::size_t count) + { + auto const target = std::min(2 * count, stream_pool_size()); + while (_streams.size() < target) { + _streams.emplace_back(rmm::cuda_stream::flags::non_blocking); + } + } public: - rmm_cuda_stream_pool() : _pool{STREAM_POOL_SIZE, rmm::cuda_stream::flags::non_blocking} {} - rmm::cuda_stream_view get_stream() override { return _pool.get_stream(); } + rmm::cuda_stream_view get_stream() override { return get_streams(1).front(); } + rmm::cuda_stream_view get_stream(stream_id_type stream_id) override { - return _pool.get_stream(stream_id); + // The id maps to the same stream on every call: growing for `stream_id` leaves the pool either + // larger than `stream_id` or at exactly `stream_pool_size()`, so the modulus below is fixed. + grow_to(stream_id + 1); + return _streams[stream_id % _streams.size()].view(); } std::vector get_streams(std::size_t count) override { - auto streams = std::vector(); - for (uint32_t i = 0; i < count; i++) { - streams.emplace_back(_pool.get_stream()); + grow_to(count); + auto const first = std::exchange(_next_stream, _next_stream + count); + auto streams = std::vector(); + streams.reserve(count); + for (std::size_t i = 0; i < count; i++) { + streams.emplace_back(_streams[(first + i) % _streams.size()].view()); } return streams; } - [[nodiscard]] std::size_t get_stream_pool_size() const override { return STREAM_POOL_SIZE; } + [[nodiscard]] std::size_t get_stream_pool_size() const override { return stream_pool_size(); } }; /** @@ -149,26 +185,106 @@ class debug_cuda_stream_pool : public cuda_stream_pool { cuda_stream_pool* create_global_cuda_stream_pool() { if (getenv("LIBCUDF_USE_DEBUG_STREAM_POOL")) return new debug_cuda_stream_pool(); - return new rmm_cuda_stream_pool(); + return new growing_cuda_stream_pool(); } +namespace { + /** - * @brief Returns a reference to the global stream pool for the current device. - * @return `cuda_stream_pool` valid on the current device. + * @brief Free lists of pools that are not currently owned by any thread, one list per device. + * + * Pools are recycled instead of destroyed so that applications which create and destroy many + * threads do not accumulate streams. The registry is intentionally leaked so that its lifetime + * covers the `thread_local` destructors that push pools back into it. + * + * Pools cannot move between devices; a stream is bound to the device that was current when the + * stream was created, so each device has its own list. */ -cuda_stream_pool& global_cuda_stream_pool() +class stream_pool_registry { + std::mutex _mutex; + std::vector> _free_pools; + + public: + stream_pool_registry() : _free_pools(get_num_cuda_devices()) {} + + /** + * @brief Takes a pool for `device_id`, reusing a retired one if there is one available. + */ + cuda_stream_pool* acquire(rmm::cuda_device_id device_id) + { + { + std::lock_guard const lock(_mutex); + auto& free_pools = _free_pools[device_id.value()]; + if (not free_pools.empty()) { + auto* pool = free_pools.back(); + free_pools.pop_back(); + return pool; + } + } + return create_global_cuda_stream_pool(); + } + + /** + * @brief Returns a pool so that another thread can reuse it. + * + * Called from a `thread_local` destructor, so it must not call into CUDA; destroying streams + * here would race with driver teardown when the main thread exits. + */ + void release(rmm::cuda_device_id device_id, cuda_stream_pool* pool) noexcept + { + std::lock_guard const lock(_mutex); + _free_pools[device_id.value()].push_back(pool); + } +}; + +stream_pool_registry& pool_registry() { - // using bare pointers here to deliberately allow them to leak. otherwise we wind up with - // seg faults trying to destroy stream objects after the context has shut down. - static std::vector pools(get_num_cuda_devices()); - static std::mutex mutex; - auto const device_id = get_current_cuda_device(); + static auto* registry = new stream_pool_registry(); + return *registry; +} + +/** + * @brief Owns the calling thread's pool for each device, and retires them when the thread exits. + */ +class thread_stream_pools { + std::vector _pools; + + public: + thread_stream_pools() : _pools(get_num_cuda_devices(), nullptr) {} + + ~thread_stream_pools() + { + for (rmm::cuda_device_id::value_type device = 0; std::cmp_less(device, _pools.size()); + device++) { + if (_pools[device] != nullptr) { + pool_registry().release(rmm::cuda_device_id{device}, _pools[device]); + } + } + } + + thread_stream_pools(thread_stream_pools const&) = delete; + thread_stream_pools& operator=(thread_stream_pools const&) = delete; + thread_stream_pools(thread_stream_pools&&) = delete; + thread_stream_pools& operator=(thread_stream_pools&&) = delete; - std::lock_guard const lock(mutex); - if (pools[device_id.value()] == nullptr) { - pools[device_id.value()] = create_global_cuda_stream_pool(); + cuda_stream_pool& pool_for(rmm::cuda_device_id device_id) + { + auto*& pool = _pools[device_id.value()]; + if (pool == nullptr) { pool = pool_registry().acquire(device_id); } + return *pool; } - return *pools[device_id.value()]; +}; + +} // namespace + +/** + * @brief Returns a reference to the calling thread's stream pool for the current device. + * @return `cuda_stream_pool` owned by the current thread and valid on the current device. + */ +cuda_stream_pool& global_cuda_stream_pool() +{ + thread_local thread_stream_pools pools; + return pools.pool_for(get_current_cuda_device()); } std::vector fork_streams(rmm::cuda_stream_view stream, std::size_t count) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 06c0462ebed9..6b27b65b18a2 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -452,6 +452,7 @@ ConfigureTest( utilities_tests/lists_column_wrapper_tests.cpp utilities_tests/memory_resource_tests.cpp utilities_tests/pinned_memory_tests.cpp + utilities_tests/stream_pool_tests.cpp utilities_tests/type_check_tests.cpp utilities_tests/type_list_tests.cpp ) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp new file mode 100644 index 000000000000..f2a587c6dd3e --- /dev/null +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -0,0 +1,165 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +class StreamPoolTest : public cudf::test::BaseFixture {}; + +namespace { + +std::vector fork_and_collect(std::size_t count) +{ + auto const streams = cudf::detail::fork_streams(cudf::get_default_stream(), count); + auto values = std::vector{}; + std::transform(streams.begin(), streams.end(), std::back_inserter(values), [](auto stream) { + return stream.value(); + }); + cudf::detail::join_streams(streams, cudf::get_default_stream()); + return values; +} + +} // namespace + +TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) +{ + auto constexpr num_streams = 8; + auto constexpr num_forks = 20; + + // Both threads fork repeatedly so that a shared round-robin counter would be very likely to + // hand the same stream to both of them. + auto collect = [](std::unordered_set& out, std::latch& ready) { + ready.arrive_and_wait(); + for (auto fork = 0; fork < num_forks; fork++) { + auto const streams = fork_and_collect(num_streams); + out.insert(streams.begin(), streams.end()); + } + }; + + std::unordered_set first_streams; + std::unordered_set second_streams; + std::latch ready{2}; + + std::thread first(collect, std::ref(first_streams), std::ref(ready)); + std::thread second(collect, std::ref(second_streams), std::ref(ready)); + first.join(); + second.join(); + + auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); + EXPECT_GE(first_streams.size(), num_streams); + EXPECT_LE(first_streams.size(), pool_size); + EXPECT_TRUE(std::none_of(first_streams.begin(), first_streams.end(), [&](auto stream) { + return second_streams.contains(stream); + })); +} + +TEST_F(StreamPoolTest, NestedForkDoesNotReuseOuterStreams) +{ + auto constexpr outer_count = 4; + auto constexpr inner_count = 2; + + std::vector outer; + std::vector inner; + + std::thread worker([&]() { + auto const outer_streams = cudf::detail::fork_streams(cudf::get_default_stream(), outer_count); + // Fork again while the outer streams are still in use, as decompression does while its caller + // is working on forked streams of its own. + inner = fork_and_collect(inner_count); + std::transform( + outer_streams.begin(), outer_streams.end(), std::back_inserter(outer), [](auto stream) { + return stream.value(); + }); + cudf::detail::join_streams(outer_streams, cudf::get_default_stream()); + }); + worker.join(); + + EXPECT_TRUE(std::none_of(inner.begin(), inner.end(), [&](auto stream) { + return std::find(outer.begin(), outer.end(), stream) != outer.end(); + })); +} + +TEST_F(StreamPoolTest, PoolGrowsToBoundedHighWaterMark) +{ + auto constexpr count = 4; + auto constexpr forks = 10; + + std::unordered_set all_streams; + std::vector> per_fork; + + std::thread worker([&]() { + for (auto fork = 0; fork < forks; fork++) { + per_fork.push_back(fork_and_collect(count)); + all_streams.insert(per_fork.back().begin(), per_fork.back().end()); + } + }); + worker.join(); + + // Consecutive requests are served from different streams, ... + for (auto fork = 1; fork < forks; fork++) { + auto const& previous = per_fork[fork - 1]; + EXPECT_TRUE(std::none_of(per_fork[fork].begin(), per_fork[fork].end(), [&](auto stream) { + return std::find(previous.begin(), previous.end(), stream) != previous.end(); + })); + } + + // ... but repeated requests cycle through the streams already created rather than making more. + auto early = std::unordered_set{}; + for (auto fork = 0; fork < forks / 2; fork++) { + early.insert(per_fork[fork].begin(), per_fork[fork].end()); + } + for (auto fork = forks / 2; fork < forks; fork++) { + EXPECT_TRUE(std::all_of(per_fork[fork].begin(), per_fork[fork].end(), [&](auto stream) { + return early.contains(stream); + })); + } + EXPECT_LE(all_streams.size(), cudf::detail::global_cuda_stream_pool().get_stream_pool_size()); +} + +TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) +{ + auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); + auto const count = pool_size + 4; + + auto const streams = fork_and_collect(count); + EXPECT_EQ(streams.size(), count); + + auto const unique = std::unordered_set(streams.begin(), streams.end()); + EXPECT_EQ(unique.size(), pool_size); +} + +TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) +{ + auto constexpr num_streams = 4; + + // Each thread cycles through its whole pool, so an adopting thread observes exactly the streams + // its predecessor used. A thread that created a fresh pool would observe entirely new ones. + auto cycle_pool = [](std::unordered_set& out) { + for (auto fork = 0; fork < 4; fork++) { + auto const streams = fork_and_collect(num_streams); + out.insert(streams.begin(), streams.end()); + } + }; + + std::unordered_set first_thread_streams; + std::thread first(cycle_pool, std::ref(first_thread_streams)); + first.join(); + + std::unordered_set second_thread_streams; + std::thread second(cycle_pool, std::ref(second_thread_streams)); + second.join(); + + EXPECT_EQ(first_thread_streams, second_thread_streams); +} From 71dd8f9700d2bd7f4030a384e0b0f792fe123114 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 17 Aug 2026 20:05:04 +0000 Subject: [PATCH 02/32] Address review: correct the pool thread-safety contract and harden tests The interface documented every accessor as thread safe, which the unsynchronized per-thread implementation does not provide; state instead that a pool is owned by one thread at a time. The tests derived their expectations from an assumed pool size, so they depended on which recycled pool a thread happened to adopt and on the configured cap; they now derive sizes from the pool and skip where a property is undefined. Adds coverage for the per-device pools. --- .../cudf/detail/utilities/stream_pool.hpp | 14 +-- cpp/src/utilities/stream_pool.cpp | 11 +- .../utilities_tests/stream_pool_tests.cpp | 118 +++++++++++++----- 3 files changed, 100 insertions(+), 43 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 4dcadd05eedd..14f31615de63 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -16,6 +16,13 @@ namespace CUDF_EXPORT cudf { namespace detail { +/** + * @brief Interface for a pool of CUDA streams. + * + * Implementations are not required to be thread safe. A pool is owned by a single thread at a time, + * which is how `global_cuda_stream_pool()` hands them out, so an implementation may keep + * unsynchronized state. Sharing one pool between threads requires external synchronization. + */ class cuda_stream_pool { public: // matching type used in rmm::cuda_stream_pool::get_stream(stream_id) @@ -30,8 +37,6 @@ class cuda_stream_pool { /** * @brief Get a `cuda_stream_view` of a stream in the pool. * - * This function is thread safe with respect to other calls to the same function. - * * @return Stream view. */ virtual rmm::cuda_stream_view get_stream() = 0; @@ -40,7 +45,6 @@ class cuda_stream_pool { * @brief Get a `cuda_stream_view` of the stream associated with `stream_id`. * * Equivalent values of `stream_id` return a `cuda_stream_view` to the same underlying stream. - * This function is thread safe with respect to other calls to the same function. * * @param stream_id Unique identifier for the desired stream * @return Requested stream view. @@ -58,8 +62,6 @@ class cuda_stream_pool { * guaranteed: once the pool has reached its maximum size the assignment wraps around, so a * request for more than half the pool can overlap with the streams the caller holds. * - * This function is thread safe with respect to other calls to the same function. - * * @param count The number of stream views to return. * @return Vector containing `count` stream views. */ @@ -68,8 +70,6 @@ class cuda_stream_pool { /** * @brief Get the maximum number of unique stream objects the pool can provide. * - * This function is thread safe with respect to other calls to the same function. - * * @return the maximum number of stream objects in the pool */ [[nodiscard]] virtual std::size_t get_stream_pool_size() const = 0; diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index e9cbfc3b63d3..36bba5449405 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -19,9 +19,14 @@ namespace cudf::detail { -// Maximum number of streams a single thread's pool will create. Sized to cover the largest -// number of streams requested by a single `fork_streams` call in libcudf, which is the number -// of distinct parquet decode kernels (see `decode_kernel_mask`). +// Maximum number of streams a single thread's pool will create, for a single device. Sized to cover +// the largest number of streams requested by a single `fork_streams` call in libcudf, which is the +// number of distinct parquet decode kernels (see `decode_kernel_mask`). +// +// This is a per-thread bound, not a process-wide one, so the streams an application holds scale +// with the number of threads that call into libcudf. Pools only grow on demand and are recycled +// when a thread exits, so the steady-state total is bounded by the peak number of concurrent +// threads rather than by the number of threads created. std::size_t constexpr STREAM_POOL_SIZE = 32; // FIXME: "borrowed" from rmm...remove when this stream pool is moved there diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index f2a587c6dd3e..6b326798707c 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -7,26 +7,37 @@ #include #include +#include + +#include +#include #include +#include +#include #include #include #include #include -#include #include class StreamPoolTest : public cudf::test::BaseFixture {}; namespace { -std::vector fork_and_collect(std::size_t count) +std::vector values_of(std::vector const& streams) { - auto const streams = cudf::detail::fork_streams(cudf::get_default_stream(), count); - auto values = std::vector{}; + auto values = std::vector{}; std::transform(streams.begin(), streams.end(), std::back_inserter(values), [](auto stream) { return stream.value(); }); + return values; +} + +std::vector fork_and_collect(std::size_t count) +{ + auto const streams = cudf::detail::fork_streams(cudf::get_default_stream(), count); + auto const values = values_of(streams); cudf::detail::join_streams(streams, cudf::get_default_stream()); return values; } @@ -35,12 +46,13 @@ std::vector fork_and_collect(std::size_t count) TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) { - auto constexpr num_streams = 8; - auto constexpr num_forks = 20; + auto constexpr num_forks = 20; + auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); + auto const num_streams = std::min(8, pool_size); // Both threads fork repeatedly so that a shared round-robin counter would be very likely to // hand the same stream to both of them. - auto collect = [](std::unordered_set& out, std::latch& ready) { + auto collect = [num_streams](std::unordered_set& out, std::latch& ready) { ready.arrive_and_wait(); for (auto fork = 0; fork < num_forks; fork++) { auto const streams = fork_and_collect(num_streams); @@ -57,7 +69,6 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) first.join(); second.join(); - auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); EXPECT_GE(first_streams.size(), num_streams); EXPECT_LE(first_streams.size(), pool_size); EXPECT_TRUE(std::none_of(first_streams.begin(), first_streams.end(), [&](auto stream) { @@ -70,6 +81,14 @@ TEST_F(StreamPoolTest, NestedForkDoesNotReuseOuterStreams) auto constexpr outer_count = 4; auto constexpr inner_count = 2; + // Disjointness is only guaranteed while the two requests together fit in the pool; beyond that + // the rotation wraps, as `get_streams` documents. + auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); + if (pool_size < outer_count + inner_count) { + GTEST_SKIP() << "Pool of " << pool_size << " streams is too small to fork " << outer_count + << " and then " << inner_count; + } + std::vector outer; std::vector inner; @@ -78,10 +97,7 @@ TEST_F(StreamPoolTest, NestedForkDoesNotReuseOuterStreams) // Fork again while the outer streams are still in use, as decompression does while its caller // is working on forked streams of its own. inner = fork_and_collect(inner_count); - std::transform( - outer_streams.begin(), outer_streams.end(), std::back_inserter(outer), [](auto stream) { - return stream.value(); - }); + outer = values_of(outer_streams); cudf::detail::join_streams(outer_streams, cudf::get_default_stream()); }); worker.join(); @@ -93,14 +109,22 @@ TEST_F(StreamPoolTest, NestedForkDoesNotReuseOuterStreams) TEST_F(StreamPoolTest, PoolGrowsToBoundedHighWaterMark) { - auto constexpr count = 4; - auto constexpr forks = 10; + auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); + if (pool_size < 4) { + GTEST_SKIP() << "Pool of " << pool_size << " streams is too small to cycle"; + } + + // Request half the pool at a time so that consecutive requests can be disjoint, and fork enough + // times to cycle any pool the worker might adopt. Pools are recycled between threads, so the size + // of the one this thread is handed is not known here. + auto const count = pool_size / 2; + auto const forks = 2 * pool_size / count + 2; std::unordered_set all_streams; std::vector> per_fork; std::thread worker([&]() { - for (auto fork = 0; fork < forks; fork++) { + for (auto fork = 0u; fork < forks; fork++) { per_fork.push_back(fork_and_collect(count)); all_streams.insert(per_fork.back().begin(), per_fork.back().end()); } @@ -108,24 +132,23 @@ TEST_F(StreamPoolTest, PoolGrowsToBoundedHighWaterMark) worker.join(); // Consecutive requests are served from different streams, ... - for (auto fork = 1; fork < forks; fork++) { + for (auto fork = 1u; fork < forks; fork++) { auto const& previous = per_fork[fork - 1]; EXPECT_TRUE(std::none_of(per_fork[fork].begin(), per_fork[fork].end(), [&](auto stream) { return std::find(previous.begin(), previous.end(), stream) != previous.end(); })); } - // ... but repeated requests cycle through the streams already created rather than making more. - auto early = std::unordered_set{}; - for (auto fork = 0; fork < forks / 2; fork++) { - early.insert(per_fork[fork].begin(), per_fork[fork].end()); + // ... but the pool stops growing: once it has been cycled, later requests only return streams + // that were handed out before. + auto seen = std::unordered_set{}; + for (auto fork = 0u; fork + 1 < forks; fork++) { + seen.insert(per_fork[fork].begin(), per_fork[fork].end()); } - for (auto fork = forks / 2; fork < forks; fork++) { - EXPECT_TRUE(std::all_of(per_fork[fork].begin(), per_fork[fork].end(), [&](auto stream) { - return early.contains(stream); - })); - } - EXPECT_LE(all_streams.size(), cudf::detail::global_cuda_stream_pool().get_stream_pool_size()); + EXPECT_TRUE(std::all_of(per_fork.back().begin(), per_fork.back().end(), [&](auto stream) { + return seen.contains(stream); + })); + EXPECT_LE(all_streams.size(), pool_size); } TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) @@ -143,23 +166,52 @@ TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) { auto constexpr num_streams = 4; + auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); - // Each thread cycles through its whole pool, so an adopting thread observes exactly the streams - // its predecessor used. A thread that created a fresh pool would observe entirely new ones. - auto cycle_pool = [](std::unordered_set& out) { - for (auto fork = 0; fork < 4; fork++) { + auto collect = [](std::unordered_set& out, std::size_t forks) { + for (auto fork = 0u; fork < forks; fork++) { auto const streams = fork_and_collect(num_streams); out.insert(streams.begin(), streams.end()); } }; std::unordered_set first_thread_streams; - std::thread first(cycle_pool, std::ref(first_thread_streams)); + std::thread first(collect, std::ref(first_thread_streams), 2); first.join(); + // The second thread cycles the whole pool, so if it adopted the retired pool it observes + // everything its predecessor used. A thread that created a fresh pool would observe entirely + // different streams, since the rotation offset carries over but the streams themselves would not. std::unordered_set second_thread_streams; - std::thread second(cycle_pool, std::ref(second_thread_streams)); + std::thread second(collect, std::ref(second_thread_streams), 2 * pool_size / num_streams + 2); second.join(); - EXPECT_EQ(first_thread_streams, second_thread_streams); + EXPECT_FALSE(first_thread_streams.empty()); + EXPECT_TRUE(std::all_of(first_thread_streams.begin(), + first_thread_streams.end(), + [&](auto stream) { return second_thread_streams.contains(stream); })); +} + +TEST_F(StreamPoolTest, EachDeviceHasItsOwnPool) +{ + auto num_devices = 0; + CUDF_CUDA_TRY(cudaGetDeviceCount(&num_devices)); + if (num_devices < 2) { GTEST_SKIP() << "Requires more than one visible device"; } + + auto constexpr count = 2; + auto const first = values_of(cudf::detail::global_cuda_stream_pool().get_streams(count)); + + auto second = std::vector{}; + { + // Streams are bound to the device that was current when they were created, so the pool for + // another device must not hand out the streams of this one. + rmm::cuda_set_device_raii const device{rmm::cuda_device_id{1}}; + second = values_of(cudf::detail::global_cuda_stream_pool().get_streams(count)); + } + + EXPECT_EQ(first.size(), count); + EXPECT_EQ(second.size(), count); + EXPECT_TRUE(std::none_of(second.begin(), second.end(), [&](auto stream) { + return std::find(first.begin(), first.end(), stream) != first.end(); + })); } From 73ce4778f7e9607e46eb45f78ac018984e7b5f3b Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 05:55:35 +0000 Subject: [PATCH 03/32] Use cuda::stream_ref and cuda::stream in the stream pool The pool interface now hands out cuda::stream_ref instead of rmm::cuda_stream_view, and owns cuda::stream rather than rmm::cuda_stream. cuda::stream is always non-blocking and takes the device explicitly, which suits a per-device pool, and its destructor uses the driver API so it does not depend on the current device. Callers are unaffected where they pass streams on, since the two view types convert implicitly; the changes elsewhere are value() to get() at kernel launches and synchronize() to sync(). --- .../cudf/detail/utilities/stream_pool.hpp | 14 +++---- cpp/src/io/parquet/page_enc.cu | 24 +++++------ cpp/src/io/parquet/page_string_decode.cu | 6 +-- cpp/src/io/text/multibyte_split.cu | 12 +++--- cpp/src/io/utilities/datasource.cpp | 2 +- cpp/src/utilities/host_memory.cpp | 2 +- cpp/src/utilities/stream_pool.cpp | 41 ++++++++++--------- cpp/tests/streams/pool_test.cu | 6 +-- cpp/tests/utilities/identify_stream_usage.cpp | 8 ++-- .../utilities_tests/stream_pool_tests.cpp | 7 ++-- 10 files changed, 63 insertions(+), 59 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 14f31615de63..c9398e3392e0 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include #include @@ -39,7 +39,7 @@ class cuda_stream_pool { * * @return Stream view. */ - virtual rmm::cuda_stream_view get_stream() = 0; + virtual cuda::stream_ref get_stream() = 0; /** * @brief Get a `cuda_stream_view` of the stream associated with `stream_id`. @@ -49,7 +49,7 @@ class cuda_stream_pool { * @param stream_id Unique identifier for the desired stream * @return Requested stream view. */ - virtual rmm::cuda_stream_view get_stream(stream_id_type stream_id) = 0; + virtual cuda::stream_ref get_stream(stream_id_type stream_id) = 0; /** * @brief Get a set of `cuda_stream_view` objects from the pool. @@ -65,7 +65,7 @@ class cuda_stream_pool { * @param count The number of stream views to return. * @return Vector containing `count` stream views. */ - virtual std::vector get_streams(std::size_t count) = 0; + virtual std::vector get_streams(std::size_t count) = 0; /** * @brief Get the maximum number of unique stream objects the pool can provide. @@ -122,8 +122,8 @@ cuda_stream_pool& global_cuda_stream_pool(); * @param count The number of `cuda_stream_view` objects to return. * @return Vector containing `count` stream views. */ -[[nodiscard]] std::vector fork_streams(rmm::cuda_stream_view stream, - std::size_t count); +[[nodiscard]] std::vector fork_streams(cuda::stream_ref stream, + std::size_t count); /** * @brief Synchronize a stream to an event on a set of streams. @@ -131,7 +131,7 @@ cuda_stream_pool& global_cuda_stream_pool(); * @param streams Streams to wait on. * @param stream Joined stream that synchronizes with the waited-on streams. */ -void join_streams(host_span streams, rmm::cuda_stream_view stream); +void join_streams(host_span streams, cuda::stream_ref stream); } // namespace detail } // namespace CUDF_EXPORT cudf diff --git a/cpp/src/io/parquet/page_enc.cu b/cpp/src/io/parquet/page_enc.cu index 4c5c5523313c..87d569abd583 100644 --- a/cpp/src/io/parquet/page_enc.cu +++ b/cpp/src/io/parquet/page_enc.cu @@ -3510,55 +3510,55 @@ void EncodePages(device_span pages, int s_idx = 0; if (BitAnd(kernel_mask, encode_kernel_mask::PLAIN) != 0) { auto const strm = streams[s_idx++]; - gpuEncodePageLevels<<>>( + gpuEncodePageLevels<<>>( pages, write_v2_headers, encode_kernel_mask::PLAIN); CUDF_CUDA_TRY(cudaGetLastError()); - gpuEncodePages<<>>( + gpuEncodePages<<>>( pages, comp_in, comp_out, comp_results, write_v2_headers, false); CUDF_CUDA_TRY(cudaGetLastError()); } if (BitAnd(kernel_mask, encode_kernel_mask::BYTE_STREAM_SPLIT) != 0) { auto const strm = streams[s_idx++]; - gpuEncodePageLevels<<>>( + gpuEncodePageLevels<<>>( pages, write_v2_headers, encode_kernel_mask::BYTE_STREAM_SPLIT); CUDF_CUDA_TRY(cudaGetLastError()); - gpuEncodePages<<>>( + gpuEncodePages<<>>( pages, comp_in, comp_out, comp_results, write_v2_headers, true); CUDF_CUDA_TRY(cudaGetLastError()); } if (BitAnd(kernel_mask, encode_kernel_mask::DELTA_BINARY) != 0) { auto const strm = streams[s_idx++]; - gpuEncodePageLevels<<>>( + gpuEncodePageLevels<<>>( pages, write_v2_headers, encode_kernel_mask::DELTA_BINARY); CUDF_CUDA_TRY(cudaGetLastError()); gpuEncodeDeltaBinaryPages - <<>>(pages, comp_in, comp_out, comp_results); + <<>>(pages, comp_in, comp_out, comp_results); CUDF_CUDA_TRY(cudaGetLastError()); } if (BitAnd(kernel_mask, encode_kernel_mask::DELTA_LENGTH_BA) != 0) { auto const strm = streams[s_idx++]; - gpuEncodePageLevels<<>>( + gpuEncodePageLevels<<>>( pages, write_v2_headers, encode_kernel_mask::DELTA_LENGTH_BA); CUDF_CUDA_TRY(cudaGetLastError()); gpuEncodeDeltaLengthByteArrayPages - <<>>(pages, comp_in, comp_out, comp_results); + <<>>(pages, comp_in, comp_out, comp_results); CUDF_CUDA_TRY(cudaGetLastError()); } if (BitAnd(kernel_mask, encode_kernel_mask::DELTA_BYTE_ARRAY) != 0) { auto const strm = streams[s_idx++]; - gpuEncodePageLevels<<>>( + gpuEncodePageLevels<<>>( pages, write_v2_headers, encode_kernel_mask::DELTA_BYTE_ARRAY); CUDF_CUDA_TRY(cudaGetLastError()); gpuEncodeDeltaByteArrayPages - <<>>(pages, comp_in, comp_out, comp_results); + <<>>(pages, comp_in, comp_out, comp_results); CUDF_CUDA_TRY(cudaGetLastError()); } if (BitAnd(kernel_mask, encode_kernel_mask::DICTIONARY) != 0) { auto const strm = streams[s_idx++]; - gpuEncodePageLevels<<>>( + gpuEncodePageLevels<<>>( pages, write_v2_headers, encode_kernel_mask::DICTIONARY); CUDF_CUDA_TRY(cudaGetLastError()); - gpuEncodeDictPages<<>>( + gpuEncodeDictPages<<>>( pages, comp_in, comp_out, comp_results, write_v2_headers); CUDF_CUDA_TRY(cudaGetLastError()); } diff --git a/cpp/src/io/parquet/page_string_decode.cu b/cpp/src/io/parquet/page_string_decode.cu index fb316bb5a208..0a25f487b8b9 100644 --- a/cpp/src/io/parquet/page_string_decode.cu +++ b/cpp/src/io/parquet/page_string_decode.cu @@ -974,7 +974,7 @@ void compute_page_string_sizes_pass1(cudf::detail::hostdevice_span pag int s_idx = 0; if (BitAnd(kernel_mask, decode_kernel_mask::DELTA_BYTE_ARRAY) != 0) { dim3 dim_delta(delta_preproc_block_size, 1); - compute_delta_page_string_sizes_kernel<<>>( + compute_delta_page_string_sizes_kernel<<>>( pages.device_ptr(), chunks, page_mask, min_row, num_rows); CUDF_CUDA_TRY(cudaGetLastError()); } @@ -983,12 +983,12 @@ void compute_page_string_sizes_pass1(cudf::detail::hostdevice_span pag compute_delta_length_page_string_sizes_kernel<<>>( + streams[s_idx++].get()>>>( pages.device_ptr(), chunks, page_mask, min_row, num_rows); CUDF_CUDA_TRY(cudaGetLastError()); } if (BitAnd(kernel_mask, STRINGS_MASK_NON_DELTA) != 0) { - compute_page_string_sizes_kernel<<>>( + compute_page_string_sizes_kernel<<>>( pages.device_ptr(), chunks, page_mask, page_string_offset_indices, min_row, num_rows); CUDF_CUDA_TRY(cudaGetLastError()); } diff --git a/cpp/src/io/text/multibyte_split.cu b/cpp/src/io/text/multibyte_split.cu index 8812077a1890..dab9ae33fb53 100644 --- a/cpp/src/io/text/multibyte_split.cu +++ b/cpp/src/io/text/multibyte_split.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -405,20 +405,20 @@ std::unique_ptr multibyte_split(cudf::io::text::data_chunk_source multibyte_split_init_kernel<<>>( // + scan_stream.get()>>>( // base_tile_idx, tiles_in_launch, tile_multistates, tile_offsets); - CUDF_CUDA_TRY(cudaStreamWaitEvent(scan_stream.value(), last_launch_event)); + CUDF_CUDA_TRY(cudaStreamWaitEvent(scan_stream.get(), last_launch_event)); if (delimiter.size() == 1) { // the single-byte case allows for a much more efficient kernel, so we special-case it byte_split_kernel<<>>( // + scan_stream.get()>>>( // base_tile_idx, chunk_offset, row_offset_storage.size(), @@ -431,7 +431,7 @@ std::unique_ptr multibyte_split(cudf::io::text::data_chunk_source multibyte_split_kernel<<>>( // + scan_stream.get()>>>( // base_tile_idx, chunk_offset, row_offset_storage.size(), @@ -492,7 +492,7 @@ std::unique_ptr multibyte_split(cudf::io::text::data_chunk_source char_storage.advance_output(output_size, scan_stream); } - CUDF_CUDA_TRY(cudaEventRecord(last_launch_event, scan_stream.value())); + CUDF_CUDA_TRY(cudaEventRecord(last_launch_event, scan_stream.get())); std::swap(read_stream, scan_stream); base_tile_idx += tiles_in_launch; diff --git a/cpp/src/io/utilities/datasource.cpp b/cpp/src/io/utilities/datasource.cpp index 05908cb9ee5e..8212beaa3e90 100644 --- a/cpp/src/io/utilities/datasource.cpp +++ b/cpp/src/io/utilities/datasource.cpp @@ -217,7 +217,7 @@ class device_buffer_source final : public datasource { auto const stream = cudf::detail::global_cuda_stream_pool().get_stream(); auto h_data = cudf::detail::make_host_vector_async( cudf::device_span{_d_buffer.data() + offset, count}, stream); - stream.synchronize(); + stream.sync(); return std::make_unique>>(std::move(h_data)); } diff --git a/cpp/src/utilities/host_memory.cpp b/cpp/src/utilities/host_memory.cpp index 3d6e56471c50..ece8971b44ce 100644 --- a/cpp/src/utilities/host_memory.cpp +++ b/cpp/src/utilities/host_memory.cpp @@ -90,7 +90,7 @@ class pinned_pool_with_fallback_memory_resource { size_t max_pool_size_{0}; // Raw pointer to avoid a segfault when the pool is destroyed on exit host_pooled_mr* pool_{nullptr}; - cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream().value()}; + cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream()}; // Wrapped in shared_ptr so the outer class is copyable (required by any_resource) std::shared_ptr fallback_{std::make_shared()}; diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 36bba5449405..8bbd29e55123 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -9,7 +9,7 @@ #include #include -#include +#include #include #include @@ -123,7 +123,7 @@ cudaEvent_t event_for_thread() * `stream_pool_size()`. */ class growing_cuda_stream_pool : public cuda_stream_pool { - std::vector _streams; + std::vector _streams; std::size_t _next_stream{0}; /** @@ -136,31 +136,34 @@ class growing_cuda_stream_pool : public cuda_stream_pool { */ void grow_to(std::size_t count) { + // A pool is only ever used with the device it was created on current, so the streams it creates + // belong to that device. `cuda::stream` is always non-blocking. + auto const device = cuda::device_ref{get_current_cuda_device().value()}; auto const target = std::min(2 * count, stream_pool_size()); while (_streams.size() < target) { - _streams.emplace_back(rmm::cuda_stream::flags::non_blocking); + _streams.emplace_back(device); } } public: - rmm::cuda_stream_view get_stream() override { return get_streams(1).front(); } + cuda::stream_ref get_stream() override { return get_streams(1).front(); } - rmm::cuda_stream_view get_stream(stream_id_type stream_id) override + cuda::stream_ref get_stream(stream_id_type stream_id) override { // The id maps to the same stream on every call: growing for `stream_id` leaves the pool either // larger than `stream_id` or at exactly `stream_pool_size()`, so the modulus below is fixed. grow_to(stream_id + 1); - return _streams[stream_id % _streams.size()].view(); + return _streams[stream_id % _streams.size()]; } - std::vector get_streams(std::size_t count) override + std::vector get_streams(std::size_t count) override { grow_to(count); auto const first = std::exchange(_next_stream, _next_stream + count); - auto streams = std::vector(); + auto streams = std::vector(); streams.reserve(count); for (std::size_t i = 0; i < count; i++) { - streams.emplace_back(_streams[(first + i) % _streams.size()].view()); + streams.emplace_back(_streams[(first + i) % _streams.size()]); } return streams; } @@ -173,15 +176,15 @@ class growing_cuda_stream_pool : public cuda_stream_pool { */ class debug_cuda_stream_pool : public cuda_stream_pool { public: - rmm::cuda_stream_view get_stream() override { return cudf::get_default_stream(); } - rmm::cuda_stream_view get_stream(stream_id_type stream_id) override + cuda::stream_ref get_stream() override { return cudf::get_default_stream(); } + cuda::stream_ref get_stream(stream_id_type stream_id) override { return cudf::get_default_stream(); } - std::vector get_streams(std::size_t count) override + std::vector get_streams(std::size_t count) override { - return std::vector(count, cudf::get_default_stream()); + return std::vector(count, cudf::get_default_stream()); } [[nodiscard]] std::size_t get_stream_pool_size() const override { return 1UL; } @@ -292,23 +295,23 @@ cuda_stream_pool& global_cuda_stream_pool() return pools.pool_for(get_current_cuda_device()); } -std::vector fork_streams(rmm::cuda_stream_view stream, std::size_t count) +std::vector fork_streams(cuda::stream_ref stream, std::size_t count) { auto const streams = global_cuda_stream_pool().get_streams(count); auto const event = event_for_thread(); - CUDF_CUDA_TRY(cudaEventRecord(event, stream)); + CUDF_CUDA_TRY(cudaEventRecord(event, stream.get())); std::for_each(streams.begin(), streams.end(), [&](auto& strm) { - CUDF_CUDA_TRY(cudaStreamWaitEvent(strm, event, 0)); + CUDF_CUDA_TRY(cudaStreamWaitEvent(strm.get(), event, 0)); }); return streams; } -void join_streams(host_span streams, rmm::cuda_stream_view stream) +void join_streams(host_span streams, cuda::stream_ref stream) { auto const event = event_for_thread(); std::for_each(streams.begin(), streams.end(), [&](auto& strm) { - CUDF_CUDA_TRY(cudaEventRecord(event, strm)); - CUDF_CUDA_TRY(cudaStreamWaitEvent(stream, event, 0)); + CUDF_CUDA_TRY(cudaEventRecord(event, strm.get())); + CUDF_CUDA_TRY(cudaStreamWaitEvent(stream.get(), event, 0)); }); } diff --git a/cpp/tests/streams/pool_test.cu b/cpp/tests/streams/pool_test.cu index 172282d7c6a0..49620c3e61fe 100644 --- a/cpp/tests/streams/pool_test.cu +++ b/cpp/tests/streams/pool_test.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,7 +9,7 @@ #include -#include +#include class StreamPoolTest : public cudf::test::BaseFixture {}; @@ -19,7 +19,7 @@ TEST_F(StreamPoolTest, ForkStreams) { auto streams = cudf::detail::fork_streams(cudf::test::get_default_stream(), 2); for (auto& stream : streams) { - do_nothing_kernel<<<1, 32, 0, stream.value()>>>(); + do_nothing_kernel<<<1, 32, 0, stream.get()>>>(); } } diff --git a/cpp/tests/utilities/identify_stream_usage.cpp b/cpp/tests/utilities/identify_stream_usage.cpp index 00249b830346..c1f9f1fce8bb 100644 --- a/cpp/tests/utilities/identify_stream_usage.cpp +++ b/cpp/tests/utilities/identify_stream_usage.cpp @@ -67,15 +67,15 @@ namespace detail { */ class test_cuda_stream_pool : public cuda_stream_pool { public: - rmm::cuda_stream_view get_stream() override { return cudf::test::get_default_stream(); } - [[maybe_unused]] rmm::cuda_stream_view get_stream(stream_id_type stream_id) override + cuda::stream_ref get_stream() override { return cudf::test::get_default_stream(); } + [[maybe_unused]] cuda::stream_ref get_stream(stream_id_type stream_id) override { return cudf::test::get_default_stream(); } - std::vector get_streams(std::size_t count) override + std::vector get_streams(std::size_t count) override { - return std::vector(count, cudf::test::get_default_stream()); + return std::vector(count, cudf::test::get_default_stream()); } [[nodiscard]] std::size_t get_stream_pool_size() const override { return 1UL; } diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 6b326798707c..dffb9bf4c559 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -10,7 +10,8 @@ #include #include -#include + +#include #include #include @@ -25,11 +26,11 @@ class StreamPoolTest : public cudf::test::BaseFixture {}; namespace { -std::vector values_of(std::vector const& streams) +std::vector values_of(std::vector const& streams) { auto values = std::vector{}; std::transform(streams.begin(), streams.end(), std::back_inserter(values), [](auto stream) { - return stream.value(); + return stream.get(); }); return values; } From b72be07a70368d52a4e9f2f9164392c3140c0a40 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 06:10:29 +0000 Subject: [PATCH 04/32] Keep only the stream pool tests that catch silent regressions Drops the growth high-water-mark and nested-fork tests, which asserted a growth schedule and a rotation heuristic the interface does not guarantee, and the per-device test, which always skips because tests request a single GPU. What remains covers cross-thread disjointness, pool reuse after a thread exits, and the over-cap repeat contract; a break in any of those is otherwise silent. --- .../utilities_tests/stream_pool_tests.cpp | 105 +----------------- 1 file changed, 2 insertions(+), 103 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index dffb9bf4c559..08354675c872 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -7,9 +7,6 @@ #include #include -#include - -#include #include @@ -52,7 +49,8 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) auto const num_streams = std::min(8, pool_size); // Both threads fork repeatedly so that a shared round-robin counter would be very likely to - // hand the same stream to both of them. + // hand the same stream to both of them. The latch keeps them alive at the same time; otherwise + // the second thread could adopt the pool the first one retired and pass trivially. auto collect = [num_streams](std::unordered_set& out, std::latch& ready) { ready.arrive_and_wait(); for (auto fork = 0; fork < num_forks; fork++) { @@ -77,81 +75,6 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) })); } -TEST_F(StreamPoolTest, NestedForkDoesNotReuseOuterStreams) -{ - auto constexpr outer_count = 4; - auto constexpr inner_count = 2; - - // Disjointness is only guaranteed while the two requests together fit in the pool; beyond that - // the rotation wraps, as `get_streams` documents. - auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); - if (pool_size < outer_count + inner_count) { - GTEST_SKIP() << "Pool of " << pool_size << " streams is too small to fork " << outer_count - << " and then " << inner_count; - } - - std::vector outer; - std::vector inner; - - std::thread worker([&]() { - auto const outer_streams = cudf::detail::fork_streams(cudf::get_default_stream(), outer_count); - // Fork again while the outer streams are still in use, as decompression does while its caller - // is working on forked streams of its own. - inner = fork_and_collect(inner_count); - outer = values_of(outer_streams); - cudf::detail::join_streams(outer_streams, cudf::get_default_stream()); - }); - worker.join(); - - EXPECT_TRUE(std::none_of(inner.begin(), inner.end(), [&](auto stream) { - return std::find(outer.begin(), outer.end(), stream) != outer.end(); - })); -} - -TEST_F(StreamPoolTest, PoolGrowsToBoundedHighWaterMark) -{ - auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); - if (pool_size < 4) { - GTEST_SKIP() << "Pool of " << pool_size << " streams is too small to cycle"; - } - - // Request half the pool at a time so that consecutive requests can be disjoint, and fork enough - // times to cycle any pool the worker might adopt. Pools are recycled between threads, so the size - // of the one this thread is handed is not known here. - auto const count = pool_size / 2; - auto const forks = 2 * pool_size / count + 2; - - std::unordered_set all_streams; - std::vector> per_fork; - - std::thread worker([&]() { - for (auto fork = 0u; fork < forks; fork++) { - per_fork.push_back(fork_and_collect(count)); - all_streams.insert(per_fork.back().begin(), per_fork.back().end()); - } - }); - worker.join(); - - // Consecutive requests are served from different streams, ... - for (auto fork = 1u; fork < forks; fork++) { - auto const& previous = per_fork[fork - 1]; - EXPECT_TRUE(std::none_of(per_fork[fork].begin(), per_fork[fork].end(), [&](auto stream) { - return std::find(previous.begin(), previous.end(), stream) != previous.end(); - })); - } - - // ... but the pool stops growing: once it has been cycled, later requests only return streams - // that were handed out before. - auto seen = std::unordered_set{}; - for (auto fork = 0u; fork + 1 < forks; fork++) { - seen.insert(per_fork[fork].begin(), per_fork[fork].end()); - } - EXPECT_TRUE(std::all_of(per_fork.back().begin(), per_fork.back().end(), [&](auto stream) { - return seen.contains(stream); - })); - EXPECT_LE(all_streams.size(), pool_size); -} - TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) { auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); @@ -192,27 +115,3 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) first_thread_streams.end(), [&](auto stream) { return second_thread_streams.contains(stream); })); } - -TEST_F(StreamPoolTest, EachDeviceHasItsOwnPool) -{ - auto num_devices = 0; - CUDF_CUDA_TRY(cudaGetDeviceCount(&num_devices)); - if (num_devices < 2) { GTEST_SKIP() << "Requires more than one visible device"; } - - auto constexpr count = 2; - auto const first = values_of(cudf::detail::global_cuda_stream_pool().get_streams(count)); - - auto second = std::vector{}; - { - // Streams are bound to the device that was current when they were created, so the pool for - // another device must not hand out the streams of this one. - rmm::cuda_set_device_raii const device{rmm::cuda_device_id{1}}; - second = values_of(cudf::detail::global_cuda_stream_pool().get_streams(count)); - } - - EXPECT_EQ(first.size(), count); - EXPECT_EQ(second.size(), count); - EXPECT_TRUE(std::none_of(second.begin(), second.end(), [&](auto stream) { - return std::find(first.begin(), first.end(), stream) != first.end(); - })); -} From b3eabc84a21132d7104869d34a28e58596a16a0c Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 19:21:40 +0000 Subject: [PATCH 05/32] Rename global_cuda_stream_pool to thread_cuda_stream_pool The accessor now returns the calling thread's pool for the current device, so "global" describes the opposite of what it does. The old name stays as a deprecated inline forwarder for out-of-tree callers of this detail header. --- cpp/include/cudf/detail/utilities/stream_pool.hpp | 15 +++++++++++++-- cpp/src/io/utilities/datasource.cpp | 4 ++-- cpp/src/utilities/host_memory.cpp | 2 +- cpp/src/utilities/stream_pool.cpp | 4 ++-- cpp/tests/utilities_tests/stream_pool_tests.cpp | 6 +++--- 5 files changed, 21 insertions(+), 10 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index c9398e3392e0..73b55f8d8e36 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -20,7 +20,7 @@ namespace detail { * @brief Interface for a pool of CUDA streams. * * Implementations are not required to be thread safe. A pool is owned by a single thread at a time, - * which is how `global_cuda_stream_pool()` hands them out, so an implementation may keep + * which is how `thread_cuda_stream_pool()` hands them out, so an implementation may keep * unsynchronized state. Sharing one pool between threads requires external synchronization. */ class cuda_stream_pool { @@ -93,7 +93,18 @@ cuda_stream_pool* create_global_cuda_stream_pool(); * The returned streams may be used from any thread, but must not be used after the thread that * obtained them has exited: a pool is recycled for reuse by another thread at that point. */ -cuda_stream_pool& global_cuda_stream_pool(); +cuda_stream_pool& thread_cuda_stream_pool(); + +/** + * @brief Get the calling thread's stream pool for the current device. + * + * @deprecated Renamed to `thread_cuda_stream_pool` now that the pool is per thread and per device. + */ +[[deprecated("Use thread_cuda_stream_pool instead.")]] // +inline cuda_stream_pool& global_cuda_stream_pool() +{ + return thread_cuda_stream_pool(); +} /** * @brief Acquire a set of `cuda_stream_view` objects and synchronize them to an event on another diff --git a/cpp/src/io/utilities/datasource.cpp b/cpp/src/io/utilities/datasource.cpp index 11382b0d8998..4ea3c707f486 100644 --- a/cpp/src/io/utilities/datasource.cpp +++ b/cpp/src/io/utilities/datasource.cpp @@ -200,7 +200,7 @@ class device_buffer_source final : public datasource { size_t host_read(size_t offset, size_t size, uint8_t* dst) override { auto const count = std::min(size, this->size() - offset); - auto const stream = cudf::detail::global_cuda_stream_pool().get_stream(); + auto const stream = cudf::detail::thread_cuda_stream_pool().get_stream(); cudf::detail::cuda_memcpy(host_span{dst, count}, device_span{ reinterpret_cast(_d_buffer.data() + offset), count}, @@ -211,7 +211,7 @@ class device_buffer_source final : public datasource { std::unique_ptr host_read(size_t offset, size_t size) override { auto const count = std::min(size, this->size() - offset); - auto const stream = cudf::detail::global_cuda_stream_pool().get_stream(); + auto const stream = cudf::detail::thread_cuda_stream_pool().get_stream(); auto h_data = cudf::detail::make_host_vector_async( cudf::device_span{_d_buffer.data() + offset, count}, stream); stream.sync(); diff --git a/cpp/src/utilities/host_memory.cpp b/cpp/src/utilities/host_memory.cpp index ece8971b44ce..3b5bd7254828 100644 --- a/cpp/src/utilities/host_memory.cpp +++ b/cpp/src/utilities/host_memory.cpp @@ -90,7 +90,7 @@ class pinned_pool_with_fallback_memory_resource { size_t max_pool_size_{0}; // Raw pointer to avoid a segfault when the pool is destroyed on exit host_pooled_mr* pool_{nullptr}; - cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream()}; + cuda::stream_ref stream_{cudf::detail::thread_cuda_stream_pool().get_stream()}; // Wrapped in shared_ptr so the outer class is copyable (required by any_resource) std::shared_ptr fallback_{std::make_shared()}; diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 8bbd29e55123..0d332042f9ca 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -289,7 +289,7 @@ class thread_stream_pools { * @brief Returns a reference to the calling thread's stream pool for the current device. * @return `cuda_stream_pool` owned by the current thread and valid on the current device. */ -cuda_stream_pool& global_cuda_stream_pool() +cuda_stream_pool& thread_cuda_stream_pool() { thread_local thread_stream_pools pools; return pools.pool_for(get_current_cuda_device()); @@ -297,7 +297,7 @@ cuda_stream_pool& global_cuda_stream_pool() std::vector fork_streams(cuda::stream_ref stream, std::size_t count) { - auto const streams = global_cuda_stream_pool().get_streams(count); + auto const streams = thread_cuda_stream_pool().get_streams(count); auto const event = event_for_thread(); CUDF_CUDA_TRY(cudaEventRecord(event, stream.get())); std::for_each(streams.begin(), streams.end(), [&](auto& strm) { diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 08354675c872..9088bda5535b 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -45,7 +45,7 @@ std::vector fork_and_collect(std::size_t count) TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) { auto constexpr num_forks = 20; - auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); + auto const pool_size = cudf::detail::thread_cuda_stream_pool().get_stream_pool_size(); auto const num_streams = std::min(8, pool_size); // Both threads fork repeatedly so that a shared round-robin counter would be very likely to @@ -77,7 +77,7 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) { - auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); + auto const pool_size = cudf::detail::thread_cuda_stream_pool().get_stream_pool_size(); auto const count = pool_size + 4; auto const streams = fork_and_collect(count); @@ -90,7 +90,7 @@ TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) { auto constexpr num_streams = 4; - auto const pool_size = cudf::detail::global_cuda_stream_pool().get_stream_pool_size(); + auto const pool_size = cudf::detail::thread_cuda_stream_pool().get_stream_pool_size(); auto collect = [](std::unordered_set& out, std::size_t forks) { for (auto fork = 0u; fork < forks; fork++) { From 3b0ac5d5d15be7c9db764c8930ce8176e26d2598 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 20:59:48 +0000 Subject: [PATCH 06/32] Describe what pool streams actually guarantee The accessor documented growth as part of its contract even though it returns the abstract interface, and it claimed streams must not be used after the obtaining thread exits. They stay valid; what is lost is the isolation, because the pool can be handed to another thread. --- cpp/include/cudf/detail/utilities/stream_pool.hpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 73b55f8d8e36..52c59c89bda0 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -87,11 +87,12 @@ cuda_stream_pool* create_global_cuda_stream_pool(); * @brief Get the calling thread's stream pool for the current device. * * Each thread has its own pool for each device it uses, so concurrent threads are handed distinct - * streams. Pools are created empty and grow on demand up to a maximum that can be configured with - * the `LIBCUDF_STREAM_POOL_SIZE` environment variable. + * streams. The maximum number of streams a pool provides can be configured with the + * `LIBCUDF_STREAM_POOL_SIZE` environment variable. * - * The returned streams may be used from any thread, but must not be used after the thread that - * obtained them has exited: a pool is recycled for reuse by another thread at that point. + * The returned streams stay valid for the lifetime of the process and may be used from any thread. + * Once the thread that obtained them exits its pool is recycled, so another thread can be handed + * the same streams; holding on to them past that point gives up the isolation the pool provides. */ cuda_stream_pool& thread_cuda_stream_pool(); @@ -115,7 +116,8 @@ inline cuda_stream_pool& global_cuda_stream_pool() * version that always returns the stream returned by `cudf::get_default_stream()`. To use this * debugging version, set the environment variable `LIBCUDF_USE_DEBUG_STREAM_POOL`. * - * The returned streams must not be used after the calling thread has exited. + * The returned streams stay valid after the calling thread exits, but its pool is recycled at that + * point, so they may then be handed to another thread as well. * * Example usage: * @code{.cpp} From 7a08f9ca58f759b7e0db78093006359c1a1a498a Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 20:59:56 +0000 Subject: [PATCH 07/32] Drop the stream pool size accessor The growing pool reported its cap rather than the number of streams it holds, and no production code queried it; the tests that did now derive their expectations from observable behavior instead. --- .../cudf/detail/utilities/stream_pool.hpp | 14 ++---- cpp/src/utilities/stream_pool.cpp | 4 -- cpp/tests/utilities/identify_stream_usage.cpp | 2 - .../utilities_tests/stream_pool_tests.cpp | 43 ++++++++----------- 4 files changed, 23 insertions(+), 40 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 52c59c89bda0..825c552cc070 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -54,8 +54,8 @@ class cuda_stream_pool { /** * @brief Get a set of `cuda_stream_view` objects from the pool. * - * The returned streams are distinct unless `count` is greater than the value returned by - * `get_stream_pool_size()`, in which case streams are repeated. + * The returned streams are distinct unless `count` is greater than the maximum number of streams + * the pool provides, in which case streams are repeated. * * Consecutive calls are served from different streams where the pool is large enough, so a * nested call generally does not return streams that its caller is already using. This is not @@ -67,13 +67,6 @@ class cuda_stream_pool { */ virtual std::vector get_streams(std::size_t count) = 0; - /** - * @brief Get the maximum number of unique stream objects the pool can provide. - * - * @return the maximum number of stream objects in the pool - */ - [[nodiscard]] virtual std::size_t get_stream_pool_size() const = 0; - protected: cuda_stream_pool() = default; }; @@ -102,7 +95,8 @@ cuda_stream_pool& thread_cuda_stream_pool(); * @deprecated Renamed to `thread_cuda_stream_pool` now that the pool is per thread and per device. */ [[deprecated("Use thread_cuda_stream_pool instead.")]] // -inline cuda_stream_pool& global_cuda_stream_pool() +inline cuda_stream_pool& +global_cuda_stream_pool() { return thread_cuda_stream_pool(); } diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 0d332042f9ca..1c458044370e 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -167,8 +167,6 @@ class growing_cuda_stream_pool : public cuda_stream_pool { } return streams; } - - [[nodiscard]] std::size_t get_stream_pool_size() const override { return stream_pool_size(); } }; /** @@ -186,8 +184,6 @@ class debug_cuda_stream_pool : public cuda_stream_pool { { return std::vector(count, cudf::get_default_stream()); } - - [[nodiscard]] std::size_t get_stream_pool_size() const override { return 1UL; } }; cuda_stream_pool* create_global_cuda_stream_pool() diff --git a/cpp/tests/utilities/identify_stream_usage.cpp b/cpp/tests/utilities/identify_stream_usage.cpp index c1f9f1fce8bb..15f3cabe905d 100644 --- a/cpp/tests/utilities/identify_stream_usage.cpp +++ b/cpp/tests/utilities/identify_stream_usage.cpp @@ -77,8 +77,6 @@ class test_cuda_stream_pool : public cuda_stream_pool { { return std::vector(count, cudf::test::get_default_stream()); } - - [[nodiscard]] std::size_t get_stream_pool_size() const override { return 1UL; } }; cuda_stream_pool* create_global_cuda_stream_pool() { return new test_cuda_stream_pool(); } diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 9088bda5535b..e36e01974cba 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -44,14 +44,13 @@ std::vector fork_and_collect(std::size_t count) TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) { - auto constexpr num_forks = 20; - auto const pool_size = cudf::detail::thread_cuda_stream_pool().get_stream_pool_size(); - auto const num_streams = std::min(8, pool_size); + auto constexpr num_forks = 20; + auto constexpr num_streams = 8; // Both threads fork repeatedly so that a shared round-robin counter would be very likely to // hand the same stream to both of them. The latch keeps them alive at the same time; otherwise // the second thread could adopt the pool the first one retired and pass trivially. - auto collect = [num_streams](std::unordered_set& out, std::latch& ready) { + auto collect = [](std::unordered_set& out, std::latch& ready) { ready.arrive_and_wait(); for (auto fork = 0; fork < num_forks; fork++) { auto const streams = fork_and_collect(num_streams); @@ -68,8 +67,7 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) first.join(); second.join(); - EXPECT_GE(first_streams.size(), num_streams); - EXPECT_LE(first_streams.size(), pool_size); + EXPECT_FALSE(first_streams.empty()); EXPECT_TRUE(std::none_of(first_streams.begin(), first_streams.end(), [&](auto stream) { return second_streams.contains(stream); })); @@ -77,37 +75,34 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) { - auto const pool_size = cudf::detail::thread_cuda_stream_pool().get_stream_pool_size(); - auto const count = pool_size + 4; + // The pool is capped, so a request this large is served by repeating streams instead of creating + // one stream per request. + auto constexpr count = 256; auto const streams = fork_and_collect(count); EXPECT_EQ(streams.size(), count); auto const unique = std::unordered_set(streams.begin(), streams.end()); - EXPECT_EQ(unique.size(), pool_size); + EXPECT_LT(unique.size(), count); } TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) { - auto constexpr num_streams = 4; - auto const pool_size = cudf::detail::thread_cuda_stream_pool().get_stream_pool_size(); - - auto collect = [](std::unordered_set& out, std::size_t forks) { - for (auto fork = 0u; fork < forks; fork++) { - auto const streams = fork_and_collect(num_streams); - out.insert(streams.begin(), streams.end()); - } - }; - std::unordered_set first_thread_streams; - std::thread first(collect, std::ref(first_thread_streams), 2); + std::thread first([&] { + auto const streams = fork_and_collect(4); + first_thread_streams.insert(streams.begin(), streams.end()); + }); first.join(); - // The second thread cycles the whole pool, so if it adopted the retired pool it observes - // everything its predecessor used. A thread that created a fresh pool would observe entirely - // different streams, since the rotation offset carries over but the streams themselves would not. + // Requesting more streams than the pool can hold cycles through all of them, so a thread that + // adopted its predecessor's retired pool observes every stream that predecessor used. One that + // created a fresh pool would observe entirely different streams. std::unordered_set second_thread_streams; - std::thread second(collect, std::ref(second_thread_streams), 2 * pool_size / num_streams + 2); + std::thread second([&] { + auto const streams = fork_and_collect(256); + second_thread_streams.insert(streams.begin(), streams.end()); + }); second.join(); EXPECT_FALSE(first_thread_streams.empty()); From a3400fa62877318247f005a1ae86900959125167 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 21:12:58 +0000 Subject: [PATCH 08/32] Drop the stream pool id-based accessor Nothing has ever called it; the identity mapping it provided only made sense for the fixed-size RMM pool it was modeled on. --- cpp/include/cudf/detail/utilities/stream_pool.hpp | 13 ------------- cpp/src/utilities/stream_pool.cpp | 12 ------------ cpp/tests/utilities/identify_stream_usage.cpp | 4 ---- 3 files changed, 29 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 825c552cc070..84a5cc96d2a8 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -25,9 +25,6 @@ namespace detail { */ class cuda_stream_pool { public: - // matching type used in rmm::cuda_stream_pool::get_stream(stream_id) - using stream_id_type = std::size_t; - virtual ~cuda_stream_pool() = default; cuda_stream_pool(cuda_stream_pool const&) = delete; cuda_stream_pool(cuda_stream_pool&&) = delete; @@ -41,16 +38,6 @@ class cuda_stream_pool { */ virtual cuda::stream_ref get_stream() = 0; - /** - * @brief Get a `cuda_stream_view` of the stream associated with `stream_id`. - * - * Equivalent values of `stream_id` return a `cuda_stream_view` to the same underlying stream. - * - * @param stream_id Unique identifier for the desired stream - * @return Requested stream view. - */ - virtual cuda::stream_ref get_stream(stream_id_type stream_id) = 0; - /** * @brief Get a set of `cuda_stream_view` objects from the pool. * diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 1c458044370e..4794a9f433b1 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -148,14 +148,6 @@ class growing_cuda_stream_pool : public cuda_stream_pool { public: cuda::stream_ref get_stream() override { return get_streams(1).front(); } - cuda::stream_ref get_stream(stream_id_type stream_id) override - { - // The id maps to the same stream on every call: growing for `stream_id` leaves the pool either - // larger than `stream_id` or at exactly `stream_pool_size()`, so the modulus below is fixed. - grow_to(stream_id + 1); - return _streams[stream_id % _streams.size()]; - } - std::vector get_streams(std::size_t count) override { grow_to(count); @@ -175,10 +167,6 @@ class growing_cuda_stream_pool : public cuda_stream_pool { class debug_cuda_stream_pool : public cuda_stream_pool { public: cuda::stream_ref get_stream() override { return cudf::get_default_stream(); } - cuda::stream_ref get_stream(stream_id_type stream_id) override - { - return cudf::get_default_stream(); - } std::vector get_streams(std::size_t count) override { diff --git a/cpp/tests/utilities/identify_stream_usage.cpp b/cpp/tests/utilities/identify_stream_usage.cpp index 15f3cabe905d..c6774902e54e 100644 --- a/cpp/tests/utilities/identify_stream_usage.cpp +++ b/cpp/tests/utilities/identify_stream_usage.cpp @@ -68,10 +68,6 @@ namespace detail { class test_cuda_stream_pool : public cuda_stream_pool { public: cuda::stream_ref get_stream() override { return cudf::test::get_default_stream(); } - [[maybe_unused]] cuda::stream_ref get_stream(stream_id_type stream_id) override - { - return cudf::test::get_default_stream(); - } std::vector get_streams(std::size_t count) override { From 445f39a4c4bad6538c7cc0aeb8818139fe6ab7aa Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 21:27:50 +0000 Subject: [PATCH 09/32] Point callers that need several streams at get_streams Repeated single-stream requests rotate over a two-stream pool, which reads like a way to spread work across the pool but serializes it. --- cpp/include/cudf/detail/utilities/stream_pool.hpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 84a5cc96d2a8..d8c3530cb4c9 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -32,7 +32,11 @@ class cuda_stream_pool { cuda_stream_pool& operator=(cuda_stream_pool&&) = delete; /** - * @brief Get a `cuda_stream_view` of a stream in the pool. + * @brief Get a single stream from the pool. + * + * @note Use `get_streams` to obtain multiple streams. The pool grows to serve the largest request + * it has seen, so repeated single-stream requests rotate over a pool of two streams and the work + * enqueued on them does not run concurrently. * * @return Stream view. */ From 3726678f933767815ad2ce237a8c6d298b21d3ed Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 22:17:00 +0000 Subject: [PATCH 10/32] docs --- cpp/include/cudf/detail/utilities/stream_pool.hpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index d8c3530cb4c9..e28765e4d22d 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -34,9 +34,8 @@ class cuda_stream_pool { /** * @brief Get a single stream from the pool. * - * @note Use `get_streams` to obtain multiple streams. The pool grows to serve the largest request - * it has seen, so repeated single-stream requests rotate over a pool of two streams and the work - * enqueued on them does not run concurrently. + * @note Use `get_streams` to obtain multiple streams. Repeated single-stream requests are not + * guaranteed to return different streams. * * @return Stream view. */ From 1a2fb3ec9666c38511349c2b83240d8236d997a7 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 22:30:19 +0000 Subject: [PATCH 11/32] Rename create_global_cuda_stream_pool to create_cuda_stream_pool The factory produces one pool per thread and device, and its doc did not mention that it exists as an override point for the stream identification utilities. --- cpp/include/cudf/detail/utilities/stream_pool.hpp | 14 +++++++------- cpp/src/utilities/stream_pool.cpp | 4 ++-- cpp/tests/utilities/identify_stream_usage.cpp | 2 +- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index e28765e4d22d..632fdfc69ee7 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -47,11 +47,6 @@ class cuda_stream_pool { * The returned streams are distinct unless `count` is greater than the maximum number of streams * the pool provides, in which case streams are repeated. * - * Consecutive calls are served from different streams where the pool is large enough, so a - * nested call generally does not return streams that its caller is already using. This is not - * guaranteed: once the pool has reached its maximum size the assignment wraps around, so a - * request for more than half the pool can overlap with the streams the caller holds. - * * @param count The number of stream views to return. * @return Vector containing `count` stream views. */ @@ -62,9 +57,14 @@ class cuda_stream_pool { }; /** - * @brief Initialize global stream pool. + * @brief Create a stream pool for a thread to use with one device. + * + * Overridden by the stream identification test utilities to substitute a pool that always returns + * the default stream. + * + * @return An owning pointer to a new pool. */ -cuda_stream_pool* create_global_cuda_stream_pool(); +cuda_stream_pool* create_cuda_stream_pool(); /** * @brief Get the calling thread's stream pool for the current device. diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 4794a9f433b1..ae5846afb6d7 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -174,7 +174,7 @@ class debug_cuda_stream_pool : public cuda_stream_pool { } }; -cuda_stream_pool* create_global_cuda_stream_pool() +cuda_stream_pool* create_cuda_stream_pool() { if (getenv("LIBCUDF_USE_DEBUG_STREAM_POOL")) return new debug_cuda_stream_pool(); return new growing_cuda_stream_pool(); @@ -213,7 +213,7 @@ class stream_pool_registry { return pool; } } - return create_global_cuda_stream_pool(); + return create_cuda_stream_pool(); } /** diff --git a/cpp/tests/utilities/identify_stream_usage.cpp b/cpp/tests/utilities/identify_stream_usage.cpp index c6774902e54e..1410dc33a4ae 100644 --- a/cpp/tests/utilities/identify_stream_usage.cpp +++ b/cpp/tests/utilities/identify_stream_usage.cpp @@ -75,7 +75,7 @@ class test_cuda_stream_pool : public cuda_stream_pool { } }; -cuda_stream_pool* create_global_cuda_stream_pool() { return new test_cuda_stream_pool(); } +cuda_stream_pool* create_cuda_stream_pool() { return new test_cuda_stream_pool(); } } // namespace detail #endif From f1c16c5dbf675a0851d7906b7af0125a7f0d3c09 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 22:36:22 +0000 Subject: [PATCH 12/32] Test the stream pool directly instead of through fork_streams The contracts under test belong to the pool, and the events fork_streams records around each request only add noise. pool_test.cu already covers fork_streams routing to the thread pool. --- .../utilities_tests/stream_pool_tests.cpp | 30 +++++++------------ 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index e36e01974cba..123e973c2631 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -6,7 +6,6 @@ #include #include -#include #include @@ -23,37 +22,30 @@ class StreamPoolTest : public cudf::test::BaseFixture {}; namespace { -std::vector values_of(std::vector const& streams) +std::vector get_streams(std::size_t count) { - auto values = std::vector{}; + auto const streams = cudf::detail::thread_cuda_stream_pool().get_streams(count); + auto values = std::vector{}; std::transform(streams.begin(), streams.end(), std::back_inserter(values), [](auto stream) { return stream.get(); }); return values; } -std::vector fork_and_collect(std::size_t count) -{ - auto const streams = cudf::detail::fork_streams(cudf::get_default_stream(), count); - auto const values = values_of(streams); - cudf::detail::join_streams(streams, cudf::get_default_stream()); - return values; -} - } // namespace TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) { - auto constexpr num_forks = 20; - auto constexpr num_streams = 8; + auto constexpr num_requests = 20; + auto constexpr num_streams = 8; - // Both threads fork repeatedly so that a shared round-robin counter would be very likely to + // Both threads request repeatedly so that a shared round-robin counter would be very likely to // hand the same stream to both of them. The latch keeps them alive at the same time; otherwise // the second thread could adopt the pool the first one retired and pass trivially. auto collect = [](std::unordered_set& out, std::latch& ready) { ready.arrive_and_wait(); - for (auto fork = 0; fork < num_forks; fork++) { - auto const streams = fork_and_collect(num_streams); + for (auto request = 0; request < num_requests; request++) { + auto const streams = get_streams(num_streams); out.insert(streams.begin(), streams.end()); } }; @@ -79,7 +71,7 @@ TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) // one stream per request. auto constexpr count = 256; - auto const streams = fork_and_collect(count); + auto const streams = get_streams(count); EXPECT_EQ(streams.size(), count); auto const unique = std::unordered_set(streams.begin(), streams.end()); @@ -90,7 +82,7 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) { std::unordered_set first_thread_streams; std::thread first([&] { - auto const streams = fork_and_collect(4); + auto const streams = get_streams(4); first_thread_streams.insert(streams.begin(), streams.end()); }); first.join(); @@ -100,7 +92,7 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) // created a fresh pool would observe entirely different streams. std::unordered_set second_thread_streams; std::thread second([&] { - auto const streams = fork_and_collect(256); + auto const streams = get_streams(256); second_thread_streams.insert(streams.begin(), streams.end()); }); second.join(); From 3ba29c00ec1957cbef0a30030b127f52c023e019 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 22:38:21 +0000 Subject: [PATCH 13/32] Name the test helper after what it returns The helper unwraps stream refs into values the tests can hash and compare; naming it get_streams made it look like a layer over the pool. --- cpp/tests/utilities_tests/stream_pool_tests.cpp | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 123e973c2631..b7fcbd12c8f9 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -22,7 +22,11 @@ class StreamPoolTest : public cudf::test::BaseFixture {}; namespace { -std::vector get_streams(std::size_t count) +/** + * @brief Requests `count` streams from the calling thread's pool, as values that can be compared + * and hashed. + */ +std::vector stream_values(std::size_t count) { auto const streams = cudf::detail::thread_cuda_stream_pool().get_streams(count); auto values = std::vector{}; @@ -45,7 +49,7 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) auto collect = [](std::unordered_set& out, std::latch& ready) { ready.arrive_and_wait(); for (auto request = 0; request < num_requests; request++) { - auto const streams = get_streams(num_streams); + auto const streams = stream_values(num_streams); out.insert(streams.begin(), streams.end()); } }; @@ -71,7 +75,7 @@ TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) // one stream per request. auto constexpr count = 256; - auto const streams = get_streams(count); + auto const streams = stream_values(count); EXPECT_EQ(streams.size(), count); auto const unique = std::unordered_set(streams.begin(), streams.end()); @@ -82,7 +86,7 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) { std::unordered_set first_thread_streams; std::thread first([&] { - auto const streams = get_streams(4); + auto const streams = stream_values(4); first_thread_streams.insert(streams.begin(), streams.end()); }); first.join(); @@ -92,7 +96,7 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) // created a fresh pool would observe entirely different streams. std::unordered_set second_thread_streams; std::thread second([&] { - auto const streams = get_streams(256); + auto const streams = stream_values(256); second_thread_streams.insert(streams.begin(), streams.end()); }); second.join(); From f90ba7dc3fe0c33636ceef9cf939fd5055f2b535 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 22:43:26 +0000 Subject: [PATCH 14/32] Correct why the concurrency test needs the latch Without it the second thread would adopt the retired pool and see the same streams, failing the test rather than passing it trivially. --- .../utilities_tests/stream_pool_tests.cpp | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index b7fcbd12c8f9..1d1cf989fb44 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -22,11 +22,7 @@ class StreamPoolTest : public cudf::test::BaseFixture {}; namespace { -/** - * @brief Requests `count` streams from the calling thread's pool, as values that can be compared - * and hashed. - */ -std::vector stream_values(std::size_t count) +std::vector get_streams_from_pool(std::size_t count) { auto const streams = cudf::detail::thread_cuda_stream_pool().get_streams(count); auto values = std::vector{}; @@ -44,12 +40,13 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) auto constexpr num_streams = 8; // Both threads request repeatedly so that a shared round-robin counter would be very likely to - // hand the same stream to both of them. The latch keeps them alive at the same time; otherwise - // the second thread could adopt the pool the first one retired and pass trivially. + // hand the same stream to both of them. The latch makes them overlap, which is the case under + // test: a thread that started after the other exited would correctly adopt the retired pool and + // observe the very streams this test requires to be disjoint. auto collect = [](std::unordered_set& out, std::latch& ready) { ready.arrive_and_wait(); for (auto request = 0; request < num_requests; request++) { - auto const streams = stream_values(num_streams); + auto const streams = get_streams_from_pool(num_streams); out.insert(streams.begin(), streams.end()); } }; @@ -75,7 +72,7 @@ TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) // one stream per request. auto constexpr count = 256; - auto const streams = stream_values(count); + auto const streams = get_streams_from_pool(count); EXPECT_EQ(streams.size(), count); auto const unique = std::unordered_set(streams.begin(), streams.end()); @@ -86,7 +83,7 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) { std::unordered_set first_thread_streams; std::thread first([&] { - auto const streams = stream_values(4); + auto const streams = get_streams_from_pool(4); first_thread_streams.insert(streams.begin(), streams.end()); }); first.join(); @@ -96,7 +93,7 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) // created a fresh pool would observe entirely different streams. std::unordered_set second_thread_streams; std::thread second([&] { - auto const streams = stream_values(256); + auto const streams = get_streams_from_pool(256); second_thread_streams.insert(streams.begin(), streams.end()); }); second.join(); From 2b461f5c45c376f6cf6026d2271c164379a6771c Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 22:52:52 +0000 Subject: [PATCH 15/32] Name the accessor for the caller, not the pool's scope current_cuda_stream_pool describes the pool a caller should use, and stays accurate if the pool ever becomes shared again. --- .../cudf/detail/utilities/stream_pool.hpp | 18 +++++++++--------- cpp/src/io/utilities/datasource.cpp | 4 ++-- cpp/src/utilities/host_memory.cpp | 2 +- cpp/src/utilities/stream_pool.cpp | 4 ++-- .../utilities_tests/stream_pool_tests.cpp | 2 +- 5 files changed, 15 insertions(+), 15 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 632fdfc69ee7..07ad965f7d56 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -20,7 +20,7 @@ namespace detail { * @brief Interface for a pool of CUDA streams. * * Implementations are not required to be thread safe. A pool is owned by a single thread at a time, - * which is how `thread_cuda_stream_pool()` hands them out, so an implementation may keep + * which is how `current_cuda_stream_pool()` hands them out, so an implementation may keep * unsynchronized state. Sharing one pool between threads requires external synchronization. */ class cuda_stream_pool { @@ -67,28 +67,28 @@ class cuda_stream_pool { cuda_stream_pool* create_cuda_stream_pool(); /** - * @brief Get the calling thread's stream pool for the current device. + * @brief Get the stream pool the calling thread should use for the current device. * - * Each thread has its own pool for each device it uses, so concurrent threads are handed distinct - * streams. The maximum number of streams a pool provides can be configured with the + * Each thread currently has its own pool for each device it uses, so concurrent threads are handed + * distinct streams. The maximum number of streams a pool provides can be configured with the * `LIBCUDF_STREAM_POOL_SIZE` environment variable. * * The returned streams stay valid for the lifetime of the process and may be used from any thread. * Once the thread that obtained them exits its pool is recycled, so another thread can be handed * the same streams; holding on to them past that point gives up the isolation the pool provides. */ -cuda_stream_pool& thread_cuda_stream_pool(); +cuda_stream_pool& current_cuda_stream_pool(); /** - * @brief Get the calling thread's stream pool for the current device. + * @brief Get the stream pool the calling thread should use for the current device. * - * @deprecated Renamed to `thread_cuda_stream_pool` now that the pool is per thread and per device. + * @deprecated Renamed to `current_cuda_stream_pool`, which does not imply a process-wide pool. */ -[[deprecated("Use thread_cuda_stream_pool instead.")]] // +[[deprecated("Use current_cuda_stream_pool instead.")]] // inline cuda_stream_pool& global_cuda_stream_pool() { - return thread_cuda_stream_pool(); + return current_cuda_stream_pool(); } /** diff --git a/cpp/src/io/utilities/datasource.cpp b/cpp/src/io/utilities/datasource.cpp index 4ea3c707f486..01fb1e6eb5a1 100644 --- a/cpp/src/io/utilities/datasource.cpp +++ b/cpp/src/io/utilities/datasource.cpp @@ -200,7 +200,7 @@ class device_buffer_source final : public datasource { size_t host_read(size_t offset, size_t size, uint8_t* dst) override { auto const count = std::min(size, this->size() - offset); - auto const stream = cudf::detail::thread_cuda_stream_pool().get_stream(); + auto const stream = cudf::detail::current_cuda_stream_pool().get_stream(); cudf::detail::cuda_memcpy(host_span{dst, count}, device_span{ reinterpret_cast(_d_buffer.data() + offset), count}, @@ -211,7 +211,7 @@ class device_buffer_source final : public datasource { std::unique_ptr host_read(size_t offset, size_t size) override { auto const count = std::min(size, this->size() - offset); - auto const stream = cudf::detail::thread_cuda_stream_pool().get_stream(); + auto const stream = cudf::detail::current_cuda_stream_pool().get_stream(); auto h_data = cudf::detail::make_host_vector_async( cudf::device_span{_d_buffer.data() + offset, count}, stream); stream.sync(); diff --git a/cpp/src/utilities/host_memory.cpp b/cpp/src/utilities/host_memory.cpp index 3b5bd7254828..8b338d45eda2 100644 --- a/cpp/src/utilities/host_memory.cpp +++ b/cpp/src/utilities/host_memory.cpp @@ -90,7 +90,7 @@ class pinned_pool_with_fallback_memory_resource { size_t max_pool_size_{0}; // Raw pointer to avoid a segfault when the pool is destroyed on exit host_pooled_mr* pool_{nullptr}; - cuda::stream_ref stream_{cudf::detail::thread_cuda_stream_pool().get_stream()}; + cuda::stream_ref stream_{cudf::detail::current_cuda_stream_pool().get_stream()}; // Wrapped in shared_ptr so the outer class is copyable (required by any_resource) std::shared_ptr fallback_{std::make_shared()}; diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index ae5846afb6d7..16af14ed1bba 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -273,7 +273,7 @@ class thread_stream_pools { * @brief Returns a reference to the calling thread's stream pool for the current device. * @return `cuda_stream_pool` owned by the current thread and valid on the current device. */ -cuda_stream_pool& thread_cuda_stream_pool() +cuda_stream_pool& current_cuda_stream_pool() { thread_local thread_stream_pools pools; return pools.pool_for(get_current_cuda_device()); @@ -281,7 +281,7 @@ cuda_stream_pool& thread_cuda_stream_pool() std::vector fork_streams(cuda::stream_ref stream, std::size_t count) { - auto const streams = thread_cuda_stream_pool().get_streams(count); + auto const streams = current_cuda_stream_pool().get_streams(count); auto const event = event_for_thread(); CUDF_CUDA_TRY(cudaEventRecord(event, stream.get())); std::for_each(streams.begin(), streams.end(), [&](auto& strm) { diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 1d1cf989fb44..f5fe0ae0c3dc 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -24,7 +24,7 @@ namespace { std::vector get_streams_from_pool(std::size_t count) { - auto const streams = cudf::detail::thread_cuda_stream_pool().get_streams(count); + auto const streams = cudf::detail::current_cuda_stream_pool().get_streams(count); auto values = std::vector{}; std::transform(streams.begin(), streams.end(), std::back_inserter(values), [](auto stream) { return stream.get(); From f28f9c2d88a2757fc66b67c205147044a4ce0065 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 23:15:02 +0000 Subject: [PATCH 16/32] Store the pool cap in the pool Reading the environment through a function-local static put a lazy initialization guard on the growth path and left the value observable from nowhere. A pool's cap is fixed once it is created. --- cpp/src/utilities/stream_pool.cpp | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 16af14ed1bba..92ed48f252f9 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -86,13 +86,11 @@ rmm::cuda_device_id get_current_cuda_device() } /** - * @brief Returns the maximum number of streams a single thread's pool will hold. + * @brief Returns the configured maximum number of streams a single pool will hold. */ -std::size_t stream_pool_size() +std::size_t configured_max_pool_size() { - static std::size_t const size = - std::max(1, getenv_or("LIBCUDF_STREAM_POOL_SIZE", STREAM_POOL_SIZE)); - return size; + return std::max(1, getenv_or("LIBCUDF_STREAM_POOL_SIZE", STREAM_POOL_SIZE)); } /** @@ -119,12 +117,12 @@ cudaEvent_t event_for_thread() * @brief Implementation of `cuda_stream_pool` that creates streams on demand. * * Instances are owned by a single thread at a time, so no synchronization is needed. The pool - * never shrinks; it grows to the largest number of streams requested so far, up to - * `stream_pool_size()`. + * never shrinks; it grows to the largest number of streams requested so far, up to `_max_size`. */ class growing_cuda_stream_pool : public cuda_stream_pool { std::vector _streams; std::size_t _next_stream{0}; + std::size_t const _max_size{configured_max_pool_size()}; /** * @brief Creates streams until the pool can serve `count` streams with room to spare. @@ -132,14 +130,14 @@ class growing_cuda_stream_pool : public cuda_stream_pool { * Twice the requested count is created so that consecutive requests are served from different * streams. Nested requests, such as decompression forking streams while its caller is using * forked streams of its own, then avoid colliding with the streams they are nested inside, - * except where the rotation wraps around a pool that has reached `stream_pool_size()`. + * except where the rotation wraps around a pool that has reached `_max_size`. */ void grow_to(std::size_t count) { // A pool is only ever used with the device it was created on current, so the streams it creates // belong to that device. `cuda::stream` is always non-blocking. auto const device = cuda::device_ref{get_current_cuda_device().value()}; - auto const target = std::min(2 * count, stream_pool_size()); + auto const target = std::min(2 * count, _max_size); while (_streams.size() < target) { _streams.emplace_back(device); } From de0f6e6407c8653c816438ec53e62f4d320b8750 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 23:21:56 +0000 Subject: [PATCH 17/32] Say that the large request grows the pool The comment read as though the pool were static at that point. --- cpp/tests/utilities_tests/stream_pool_tests.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index f5fe0ae0c3dc..9706e6f9845d 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -88,9 +88,9 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) }); first.join(); - // Requesting more streams than the pool can hold cycles through all of them, so a thread that - // adopted its predecessor's retired pool observes every stream that predecessor used. One that - // created a fresh pool would observe entirely different streams. + // A request this large grows the pool to its maximum and then cycles through every stream in it. + // Growth only appends, so a thread that adopted its predecessor's retired pool observes every + // stream that predecessor used, whereas one that created a fresh pool observes none of them. std::unordered_set second_thread_streams; std::thread second([&] { auto const streams = get_streams_from_pool(256); From c3fcc3e21b1673384b6045556c5dfa50df975fb7 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 23:26:56 +0000 Subject: [PATCH 18/32] Compare against the whole pool the first thread leaves behind Having the second thread make the oversized request meant comparing its 32 streams against the predecessor's 4. Reversing it compares a normal request against the pool it must have come from. --- .../utilities_tests/stream_pool_tests.cpp | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 9706e6f9845d..5e37c7c28d31 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -81,25 +81,26 @@ TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) { + // A request larger than any pool grows this thread's pool to its maximum and cycles through every + // stream in it, so the set holds the whole pool this thread leaves behind. std::unordered_set first_thread_streams; std::thread first([&] { - auto const streams = get_streams_from_pool(4); + auto const streams = get_streams_from_pool(1024); first_thread_streams.insert(streams.begin(), streams.end()); }); first.join(); - // A request this large grows the pool to its maximum and then cycles through every stream in it. - // Growth only appends, so a thread that adopted its predecessor's retired pool observes every - // stream that predecessor used, whereas one that created a fresh pool observes none of them. + // A thread that adopted the retired pool can only be handed streams from it. One that created a + // fresh pool would be handed newly created streams instead. std::unordered_set second_thread_streams; std::thread second([&] { - auto const streams = get_streams_from_pool(256); + auto const streams = get_streams_from_pool(4); second_thread_streams.insert(streams.begin(), streams.end()); }); second.join(); - EXPECT_FALSE(first_thread_streams.empty()); - EXPECT_TRUE(std::all_of(first_thread_streams.begin(), - first_thread_streams.end(), - [&](auto stream) { return second_thread_streams.contains(stream); })); + EXPECT_FALSE(second_thread_streams.empty()); + EXPECT_TRUE(std::all_of(second_thread_streams.begin(), + second_thread_streams.end(), + [&](auto stream) { return first_thread_streams.contains(stream); })); } From 226f3eb141a9be088a907e03bf610ff75bb1dcad Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 23:40:37 +0000 Subject: [PATCH 19/32] Hold the second thread's streams in a vector Only the set being searched needs to be hashed. --- cpp/tests/utilities_tests/stream_pool_tests.cpp | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 5e37c7c28d31..5cd8b5b60f07 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -85,21 +85,17 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) // stream in it, so the set holds the whole pool this thread leaves behind. std::unordered_set first_thread_streams; std::thread first([&] { - auto const streams = get_streams_from_pool(1024); + auto const streams = get_streams_from_pool(64); first_thread_streams.insert(streams.begin(), streams.end()); }); first.join(); - // A thread that adopted the retired pool can only be handed streams from it. One that created a - // fresh pool would be handed newly created streams instead. - std::unordered_set second_thread_streams; - std::thread second([&] { - auto const streams = get_streams_from_pool(4); - second_thread_streams.insert(streams.begin(), streams.end()); - }); + std::vector second_thread_streams; + std::thread second([&] { second_thread_streams = get_streams_from_pool(4); }); second.join(); EXPECT_FALSE(second_thread_streams.empty()); + // A thread that adopted the retired pool can only be handed streams from it EXPECT_TRUE(std::all_of(second_thread_streams.begin(), second_thread_streams.end(), [&](auto stream) { return first_thread_streams.contains(stream); })); From 8f7933d06523b53284a83c4093c250bb7d14c6bf Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 23:49:13 +0000 Subject: [PATCH 20/32] Share one oversized request size between the pool tests Both tests need a request that exceeds the pool's maximum: one to show repetition, the other to enumerate the pool. 256 stopped doing that once the maximum could be configured higher. --- .../utilities_tests/stream_pool_tests.cpp | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 5cd8b5b60f07..f255b145b88e 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -22,6 +22,10 @@ class StreamPoolTest : public cudf::test::BaseFixture {}; namespace { +// A thread has no use for more than 64 streams, so a request of this size exceeds the maximum of +// any pool worth configuring and cycles through every stream in it. +auto constexpr more_streams_than_any_pool = 128; + std::vector get_streams_from_pool(std::size_t count) { auto const streams = cudf::detail::current_cuda_stream_pool().get_streams(count); @@ -70,29 +74,24 @@ TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) { // The pool is capped, so a request this large is served by repeating streams instead of creating // one stream per request. - auto constexpr count = 256; - - auto const streams = get_streams_from_pool(count); - EXPECT_EQ(streams.size(), count); + auto const streams = get_streams_from_pool(more_streams_than_any_pool); + EXPECT_EQ(streams.size(), more_streams_than_any_pool); auto const unique = std::unordered_set(streams.begin(), streams.end()); - EXPECT_LT(unique.size(), count); + EXPECT_LT(unique.size(), more_streams_than_any_pool); } TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) { - // A request larger than any pool grows this thread's pool to its maximum and cycles through every - // stream in it, so the set holds the whole pool this thread leaves behind. + // A request larger than max pool size to get all streams in the pool std::unordered_set first_thread_streams; - std::thread first([&] { - auto const streams = get_streams_from_pool(64); + std::thread([&] { + auto const streams = get_streams_from_pool(more_streams_than_any_pool); first_thread_streams.insert(streams.begin(), streams.end()); - }); - first.join(); + }).join(); std::vector second_thread_streams; - std::thread second([&] { second_thread_streams = get_streams_from_pool(4); }); - second.join(); + std::thread([&] { second_thread_streams = get_streams_from_pool(4); }).join(); EXPECT_FALSE(second_thread_streams.empty()); // A thread that adopted the retired pool can only be handed streams from it From 35f71b8b82a665b03c8345cbed27fa265e8b46c8 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 18 Aug 2026 23:52:57 +0000 Subject: [PATCH 21/32] Rename the test helper and request past a 64-stream pool A request of 64 does not exceed a pool configured to hold 64, so the repetition it is meant to observe does not happen. --- .../utilities_tests/stream_pool_tests.cpp | 24 +++++++++---------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index f255b145b88e..92692b520262 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -22,11 +22,7 @@ class StreamPoolTest : public cudf::test::BaseFixture {}; namespace { -// A thread has no use for more than 64 streams, so a request of this size exceeds the maximum of -// any pool worth configuring and cycles through every stream in it. -auto constexpr more_streams_than_any_pool = 128; - -std::vector get_streams_from_pool(std::size_t count) +std::vector get_hashable_streams(std::size_t count) { auto const streams = cudf::detail::current_cuda_stream_pool().get_streams(count); auto values = std::vector{}; @@ -50,7 +46,7 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) auto collect = [](std::unordered_set& out, std::latch& ready) { ready.arrive_and_wait(); for (auto request = 0; request < num_requests; request++) { - auto const streams = get_streams_from_pool(num_streams); + auto const streams = get_hashable_streams(num_streams); out.insert(streams.begin(), streams.end()); } }; @@ -72,13 +68,15 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) { - // The pool is capped, so a request this large is served by repeating streams instead of creating - // one stream per request. - auto const streams = get_streams_from_pool(more_streams_than_any_pool); - EXPECT_EQ(streams.size(), more_streams_than_any_pool); + // The pool is capped below this, so a request this large is served by repeating streams instead + // of creating one stream per request. + auto constexpr count = 128; + + auto const streams = get_hashable_streams(count); + EXPECT_EQ(streams.size(), count); auto const unique = std::unordered_set(streams.begin(), streams.end()); - EXPECT_LT(unique.size(), more_streams_than_any_pool); + EXPECT_LT(unique.size(), count); } TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) @@ -86,12 +84,12 @@ TEST_F(StreamPoolTest, PoolIsReusedAfterThreadExits) // A request larger than max pool size to get all streams in the pool std::unordered_set first_thread_streams; std::thread([&] { - auto const streams = get_streams_from_pool(more_streams_than_any_pool); + auto const streams = get_hashable_streams(128); first_thread_streams.insert(streams.begin(), streams.end()); }).join(); std::vector second_thread_streams; - std::thread([&] { second_thread_streams = get_streams_from_pool(4); }).join(); + std::thread([&] { second_thread_streams = get_hashable_streams(4); }).join(); EXPECT_FALSE(second_thread_streams.empty()); // A thread that adopted the retired pool can only be handed streams from it From 9d8d0d2edc8b598e4c10a83e0fb9de0ecea3c8d2 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 19 Aug 2026 00:08:39 +0000 Subject: [PATCH 22/32] last of test clean up --- cpp/tests/utilities_tests/stream_pool_tests.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 92692b520262..12bd00bb7bd5 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -68,13 +68,12 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) { - // The pool is capped below this, so a request this large is served by repeating streams instead - // of creating one stream per request. auto constexpr count = 128; auto const streams = get_hashable_streams(count); EXPECT_EQ(streams.size(), count); + // Request this large is served by repeating streams instead of growing the pool auto const unique = std::unordered_set(streams.begin(), streams.end()); EXPECT_LT(unique.size(), count); } From cb741df85e85f04e31522bcfaad61b45220401c0 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 19 Aug 2026 00:42:29 +0000 Subject: [PATCH 23/32] shorten comment --- cpp/tests/utilities_tests/stream_pool_tests.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/cpp/tests/utilities_tests/stream_pool_tests.cpp b/cpp/tests/utilities_tests/stream_pool_tests.cpp index 12bd00bb7bd5..571d23d0a280 100644 --- a/cpp/tests/utilities_tests/stream_pool_tests.cpp +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -39,10 +39,9 @@ TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) auto constexpr num_requests = 20; auto constexpr num_streams = 8; - // Both threads request repeatedly so that a shared round-robin counter would be very likely to - // hand the same stream to both of them. The latch makes them overlap, which is the case under - // test: a thread that started after the other exited would correctly adopt the retired pool and - // observe the very streams this test requires to be disjoint. + // Repeated requests make a shared round-robin counter very likely to hand the same stream to both + // threads. The latch makes them overlap; a thread starting after the other exited would adopt its + // retired pool and see the same streams. auto collect = [](std::unordered_set& out, std::latch& ready) { ready.arrive_and_wait(); for (auto request = 0; request < num_requests; request++) { From 89634693b5b31f69ac895d0327d566fd444b8fae Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 19 Aug 2026 00:50:50 +0000 Subject: [PATCH 24/32] impl clean up --- cpp/src/utilities/stream_pool.cpp | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 92ed48f252f9..0c59890cacf6 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -125,20 +125,14 @@ class growing_cuda_stream_pool : public cuda_stream_pool { std::size_t const _max_size{configured_max_pool_size()}; /** - * @brief Creates streams until the pool can serve `count` streams with room to spare. - * - * Twice the requested count is created so that consecutive requests are served from different - * streams. Nested requests, such as decompression forking streams while its caller is using - * forked streams of its own, then avoid colliding with the streams they are nested inside, - * except where the rotation wraps around a pool that has reached `_max_size`. + * @brief Creates streams until the pool holds `size` of them, or has reached `_max_size`. */ - void grow_to(std::size_t count) + void grow_to(std::size_t size) { - // A pool is only ever used with the device it was created on current, so the streams it creates - // belong to that device. `cuda::stream` is always non-blocking. auto const device = cuda::device_ref{get_current_cuda_device().value()}; - auto const target = std::min(2 * count, _max_size); + auto const target = std::min(size, _max_size); while (_streams.size() < target) { + // `cuda::stream` creates non-blocking streams. _streams.emplace_back(device); } } @@ -148,7 +142,9 @@ class growing_cuda_stream_pool : public cuda_stream_pool { std::vector get_streams(std::size_t count) override { - grow_to(count); + // Growing to twice the requested count leaves room for consecutive requests to return + // different streams. + grow_to(2 * count); auto const first = std::exchange(_next_stream, _next_stream + count); auto streams = std::vector(); streams.reserve(count); From b5d2c4abb026ee8cfe2eebae215b42e7fa0b5593 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 19 Aug 2026 02:49:42 +0000 Subject: [PATCH 25/32] Build the join_streams wrapper from cuda::stream_ref The pylibcudf wrapper built an rmm::cuda_stream_view vector, which no longer converts to the span join_streams takes. --- .../pylibcudf/libcudf/detail/utilities/stream_pool.pxd | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd b/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd index 399a868db717..dde126e295d6 100644 --- a/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd @@ -12,7 +12,9 @@ cdef extern from * nogil: """ #include #include - #include + + #include + #include namespace { @@ -20,8 +22,8 @@ cdef extern from * nogil: cudf::host_span streams, cudaStream_t stream ) { - std::vector stream_views(streams.begin(), streams.end()); - cudf::detail::join_streams(stream_views, stream); + std::vector stream_refs(streams.begin(), streams.end()); + cudf::detail::join_streams(stream_refs, stream); } } """ From 5c2eca58fb56c9a0550aa3388151b844a3b6985a Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 19 Aug 2026 02:50:10 +0000 Subject: [PATCH 26/32] Update copyright header on the stream pool bindings --- .../pylibcudf/libcudf/detail/utilities/stream_pool.pxd | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd b/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd index dde126e295d6..b05edceb3abc 100644 --- a/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from cuda.bindings.cyruntime cimport cudaStream_t From 1f55ffa96fccd6e89188082c98ccd6c8081b8f8f Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 19 Aug 2026 16:35:11 +0000 Subject: [PATCH 27/32] Refresh stream pool docs for cuda::stream_ref and include The docs still described the pool in terms of cuda_stream_view after the migration to cuda::stream_ref, and the implementation relied on a transitive include for cuda::device_ref. --- .../cudf/detail/utilities/stream_pool.hpp | 18 +++++++++++------- cpp/src/utilities/stream_pool.cpp | 1 + 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 07ad965f7d56..b2a58fbc0b71 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -37,18 +37,18 @@ class cuda_stream_pool { * @note Use `get_streams` to obtain multiple streams. Repeated single-stream requests are not * guaranteed to return different streams. * - * @return Stream view. + * @return Stream reference. */ virtual cuda::stream_ref get_stream() = 0; /** - * @brief Get a set of `cuda_stream_view` objects from the pool. + * @brief Get a set of `cuda::stream_ref` objects from the pool. * * The returned streams are distinct unless `count` is greater than the maximum number of streams * the pool provides, in which case streams are repeated. * - * @param count The number of stream views to return. - * @return Vector containing `count` stream views. + * @param count The number of stream references to return. + * @return Vector containing `count` stream references. */ virtual std::vector get_streams(std::size_t count) = 0; @@ -76,6 +76,8 @@ cuda_stream_pool* create_cuda_stream_pool(); * The returned streams stay valid for the lifetime of the process and may be used from any thread. * Once the thread that obtained them exits its pool is recycled, so another thread can be handed * the same streams; holding on to them past that point gives up the isolation the pool provides. + * + * @return Reference to the calling thread's stream pool for the current device. */ cuda_stream_pool& current_cuda_stream_pool(); @@ -83,6 +85,8 @@ cuda_stream_pool& current_cuda_stream_pool(); * @brief Get the stream pool the calling thread should use for the current device. * * @deprecated Renamed to `current_cuda_stream_pool`, which does not imply a process-wide pool. + * + * @return Reference to the calling thread's stream pool for the current device. */ [[deprecated("Use current_cuda_stream_pool instead.")]] // inline cuda_stream_pool& @@ -92,7 +96,7 @@ global_cuda_stream_pool() } /** - * @brief Acquire a set of `cuda_stream_view` objects and synchronize them to an event on another + * @brief Acquire a set of `cuda::stream_ref` objects and synchronize them to an event on another * stream. * * By default the calling thread's stream pool is used to obtain the streams, so streams are not @@ -116,8 +120,8 @@ global_cuda_stream_pool() * @endcode * * @param stream Stream that the returned streams will wait on. - * @param count The number of `cuda_stream_view` objects to return. - * @return Vector containing `count` stream views. + * @param count The number of `cuda::stream_ref` objects to return. + * @return Vector containing `count` stream references. */ [[nodiscard]] std::vector fork_streams(cuda::stream_ref stream, std::size_t count); diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 0c59890cacf6..6fe71ae5c472 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include From c7565b039578d1d16eb222c9bfaaf4b859e7bfcf Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 20 Aug 2026 03:10:53 +0000 Subject: [PATCH 28/32] Drop the unused stream from the pinned memory resource Allocation and deallocation use the caller's stream; the member was read only in operator==, where it never changed the result, because two resources share a pool pointer only when one is a copy of the other. Holding it also took a stream from the pool of whichever thread first touched the process-wide resource, which that thread's successor could then be handed. --- cpp/src/utilities/host_memory.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/src/utilities/host_memory.cpp b/cpp/src/utilities/host_memory.cpp index 8b338d45eda2..026a77360500 100644 --- a/cpp/src/utilities/host_memory.cpp +++ b/cpp/src/utilities/host_memory.cpp @@ -4,7 +4,6 @@ */ #include -#include #include #include #include @@ -16,6 +15,8 @@ #include #include +#include + #include #include #include @@ -90,7 +91,6 @@ class pinned_pool_with_fallback_memory_resource { size_t max_pool_size_{0}; // Raw pointer to avoid a segfault when the pool is destroyed on exit host_pooled_mr* pool_{nullptr}; - cuda::stream_ref stream_{cudf::detail::current_cuda_stream_pool().get_stream()}; // Wrapped in shared_ptr so the outer class is copyable (required by any_resource) std::shared_ptr fallback_{std::make_shared()}; @@ -187,7 +187,7 @@ class pinned_pool_with_fallback_memory_resource { bool operator==(pinned_pool_with_fallback_memory_resource const& other) const noexcept { - return pool_ == other.pool_ && stream_ == other.stream_; + return pool_ == other.pool_; } bool operator!=(pinned_pool_with_fallback_memory_resource const& other) const noexcept From c93d3d1505b3b3713515a1d56a72b1b94e9039ea Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 20 Aug 2026 03:29:13 +0000 Subject: [PATCH 29/32] Check the batched memcpy stream argument on CUDA 12.8 and newer cudaMemcpyBatchAsync arrived in 12.8, where the callback parameters are named after that version, so guarding on 13.0 skipped the check on every supported CUDA 12 toolkit. --- cpp/tests/utilities/identify_stream_usage.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cpp/tests/utilities/identify_stream_usage.cpp b/cpp/tests/utilities/identify_stream_usage.cpp index 3ed3b410067f..6fcf82a68e9b 100644 --- a/cpp/tests/utilities/identify_stream_usage.cpp +++ b/cpp/tests/utilities/identify_stream_usage.cpp @@ -215,6 +215,9 @@ void sanitizer_subscriber::callback(Sanitizer_CallbackDomain domain, #if CUDART_VERSION >= 13000 CHECK_STREAM_ARG(cudaMemcpyBatchAsync, 13000, stream); CHECK_STREAM_ARG(cudaMemcpyBatchAsync_ptsz, 13000, stream); +#elif CUDART_VERSION >= 12080 + CHECK_STREAM_ARG(cudaMemcpyBatchAsync, 12080, stream); + CHECK_STREAM_ARG(cudaMemcpyBatchAsync_ptsz, 12080, stream); #endif CHECK_STREAM_ARG(cudaMemcpyFromSymbolAsync, 3020, stream); CHECK_STREAM_ARG(cudaMemcpyFromSymbolAsync_ptsz, 7000, stream); From 9cb1d52a1e06d9af44590217e542ec0626679f7d Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 20 Aug 2026 03:29:13 +0000 Subject: [PATCH 30/32] Take the streams to join as a std::span The developer guide reserves host_span for the cases that need its libcudf extensions, and this one only reads a contiguous range of stream references. Also says vector rather than set where the docs describe what get_streams and fork_streams return, and corrects the claim that the parquet decode kernels need the most streams of any request in libcudf: host (de)compression asks for one per dispatched chunk, which can exceed the pool maximum. --- cpp/include/cudf/detail/utilities/stream_pool.hpp | 12 ++++++------ cpp/src/utilities/stream_pool.cpp | 10 +++++++--- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index b2a58fbc0b71..090d786b8b85 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -6,11 +6,11 @@ #pragma once #include -#include #include #include +#include #include namespace CUDF_EXPORT cudf { @@ -42,7 +42,7 @@ class cuda_stream_pool { virtual cuda::stream_ref get_stream() = 0; /** - * @brief Get a set of `cuda::stream_ref` objects from the pool. + * @brief Get a vector of `cuda::stream_ref` objects from the pool. * * The returned streams are distinct unless `count` is greater than the maximum number of streams * the pool provides, in which case streams are repeated. @@ -96,8 +96,8 @@ global_cuda_stream_pool() } /** - * @brief Acquire a set of `cuda::stream_ref` objects and synchronize them to an event on another - * stream. + * @brief Acquire a vector of `cuda::stream_ref` objects and synchronize them to an event on + * another stream. * * By default the calling thread's stream pool is used to obtain the streams, so streams are not * shared with concurrently forking threads. The only other implementation at present is a debugging @@ -127,12 +127,12 @@ global_cuda_stream_pool() std::size_t count); /** - * @brief Synchronize a stream to an event on a set of streams. + * @brief Synchronize a stream to an event on each of a group of streams. * * @param streams Streams to wait on. * @param stream Joined stream that synchronizes with the waited-on streams. */ -void join_streams(host_span streams, cuda::stream_ref stream); +void join_streams(std::span streams, cuda::stream_ref stream); } // namespace detail } // namespace CUDF_EXPORT cudf diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 6fe71ae5c472..00c7c1ca1456 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -9,20 +9,24 @@ #include #include +#include + #include #include #include #include #include +#include #include #include namespace cudf::detail { // Maximum number of streams a single thread's pool will create, for a single device. Sized to cover -// the largest number of streams requested by a single `fork_streams` call in libcudf, which is the -// number of distinct parquet decode kernels (see `decode_kernel_mask`). +// the number of distinct parquet decode kernels (see `decode_kernel_mask`), the largest fixed +// number of streams a single `fork_streams` call in libcudf asks for. Host (de)compression asks for +// one stream per chunk it dispatches, which can exceed this; those requests get repeated streams. // // This is a per-thread bound, not a process-wide one, so the streams an application holds scale // with the number of threads that call into libcudf. Pools only grow on demand and are recycled @@ -285,7 +289,7 @@ std::vector fork_streams(cuda::stream_ref stream, std::size_t return streams; } -void join_streams(host_span streams, cuda::stream_ref stream) +void join_streams(std::span streams, cuda::stream_ref stream) { auto const event = event_for_thread(); std::for_each(streams.begin(), streams.end(), [&](auto& strm) { From da76fd4b0a7bb0cc1c680ad2c358faf6b9673f53 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 20 Aug 2026 03:31:57 +0000 Subject: [PATCH 31/32] Say which requests the pool maximum is sized for Parquet decode is the largest request with a fixed bound, not the largest request outright; the requests that scale with the host worker count can go past the maximum and share streams. --- cpp/src/utilities/stream_pool.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 00c7c1ca1456..5668fdc43382 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -24,9 +24,10 @@ namespace cudf::detail { // Maximum number of streams a single thread's pool will create, for a single device. Sized to cover -// the number of distinct parquet decode kernels (see `decode_kernel_mask`), the largest fixed -// number of streams a single `fork_streams` call in libcudf asks for. Host (de)compression asks for -// one stream per chunk it dispatches, which can exceed this; those requests get repeated streams. +// the largest request libcudf makes with a fixed bound, which is one stream per distinct parquet +// decode kernel (see `decode_kernel_mask`). Host (de)compression, the JSON reader and others +// instead scale their request with the host worker count, so raising `LIBCUDF_NUM_HOST_WORKERS` +// past this leaves those requests sharing streams. // // This is a per-thread bound, not a process-wide one, so the streams an application holds scale // with the number of threads that call into libcudf. Pools only grow on demand and are recycled From eb7b11c0d66aedf4ea4081bd0df6fd49219e11bd Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 20 Aug 2026 05:26:32 +0000 Subject: [PATCH 32/32] recycle events as well --- cpp/src/utilities/stream_pool.cpp | 129 ++++++++++++++++++++---------- 1 file changed, 88 insertions(+), 41 deletions(-) diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 5668fdc43382..483b18194fcc 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -99,24 +99,6 @@ std::size_t configured_max_pool_size() return std::max(1, getenv_or("LIBCUDF_STREAM_POOL_SIZE", STREAM_POOL_SIZE)); } -/** - * @brief Returns a cudaEvent_t for the current thread. - * - * The returned event is valid for the current device. - * - * @return A cudaEvent_t unique to the current thread and valid on the current device. - */ -cudaEvent_t event_for_thread() -{ - // The program may crash if this function is called from the main thread and user application - // subsequently calls cudaDeviceReset(). - // As a workaround, here we intentionally disable RAII and leak cudaEvent_t. - thread_local static std::vector thread_events(get_num_cuda_devices()); - auto const device_id = get_current_cuda_device(); - if (not thread_events[device_id.value()]) { thread_events[device_id.value()] = new cuda_event(); } - return *thread_events[device_id.value()]; -} - } // namespace /** @@ -183,26 +165,31 @@ cuda_stream_pool* create_cuda_stream_pool() namespace { /** - * @brief Free lists of pools that are not currently owned by any thread, one list per device. + * @brief Free lists of pools and events that are not currently owned by any thread, one list of + * each per device. * - * Pools are recycled instead of destroyed so that applications which create and destroy many - * threads do not accumulate streams. The registry is intentionally leaked so that its lifetime - * covers the `thread_local` destructors that push pools back into it. + * They are recycled instead of destroyed so that applications which create and destroy many + * threads do not accumulate streams and events. The registry is intentionally leaked so that its + * lifetime covers the `thread_local` destructors that push into it. * - * Pools cannot move between devices; a stream is bound to the device that was current when the - * stream was created, so each device has its own list. + * Neither pools nor events can move between devices; both are bound to the device that was current + * when they were created, so each device has its own lists. */ class stream_pool_registry { std::mutex _mutex; std::vector> _free_pools; + std::vector> _free_events; public: - stream_pool_registry() : _free_pools(get_num_cuda_devices()) {} + stream_pool_registry() + : _free_pools(get_num_cuda_devices()), _free_events(get_num_cuda_devices()) + { + } /** * @brief Takes a pool for `device_id`, reusing a retired one if there is one available. */ - cuda_stream_pool* acquire(rmm::cuda_device_id device_id) + cuda_stream_pool* acquire_pool(rmm::cuda_device_id device_id) { { std::lock_guard const lock(_mutex); @@ -216,17 +203,49 @@ class stream_pool_registry { return create_cuda_stream_pool(); } + /** + * @brief Takes an event for `device_id`, reusing a retired one if there is one available. + * + * Events are never destroyed: the program may crash if one is destroyed after the application + * calls `cudaDeviceReset()`. + */ + cuda_event* acquire_event(rmm::cuda_device_id device_id) + { + { + std::lock_guard const lock(_mutex); + auto& free_events = _free_events[device_id.value()]; + if (not free_events.empty()) { + auto* event = free_events.back(); + free_events.pop_back(); + return event; + } + } + return new cuda_event(); + } + /** * @brief Returns a pool so that another thread can reuse it. * * Called from a `thread_local` destructor, so it must not call into CUDA; destroying streams * here would race with driver teardown when the main thread exits. */ - void release(rmm::cuda_device_id device_id, cuda_stream_pool* pool) noexcept + void release_pool(rmm::cuda_device_id device_id, cuda_stream_pool* pool) noexcept { std::lock_guard const lock(_mutex); _free_pools[device_id.value()].push_back(pool); } + + /** + * @brief Returns an event so that another thread can reuse it. + * + * Called from a `thread_local` destructor, so it must not call into CUDA, for the same reason as + * `release_pool`. + */ + void release_event(rmm::cuda_device_id device_id, cuda_event* event) noexcept + { + std::lock_guard const lock(_mutex); + _free_events[device_id.value()].push_back(event); + } }; stream_pool_registry& pool_registry() @@ -236,37 +255,66 @@ stream_pool_registry& pool_registry() } /** - * @brief Owns the calling thread's pool for each device, and retires them when the thread exits. + * @brief Owns the calling thread's pool and event for each device, and retires them when the + * thread exits. */ -class thread_stream_pools { +class thread_stream_resources { std::vector _pools; + std::vector _events; public: - thread_stream_pools() : _pools(get_num_cuda_devices(), nullptr) {} + thread_stream_resources() + : _pools(get_num_cuda_devices(), nullptr), _events(get_num_cuda_devices(), nullptr) + { + } - ~thread_stream_pools() + ~thread_stream_resources() { for (rmm::cuda_device_id::value_type device = 0; std::cmp_less(device, _pools.size()); device++) { - if (_pools[device] != nullptr) { - pool_registry().release(rmm::cuda_device_id{device}, _pools[device]); - } + auto const device_id = rmm::cuda_device_id{device}; + if (_pools[device] != nullptr) { pool_registry().release_pool(device_id, _pools[device]); } + if (_events[device] != nullptr) { pool_registry().release_event(device_id, _events[device]); } } } - thread_stream_pools(thread_stream_pools const&) = delete; - thread_stream_pools& operator=(thread_stream_pools const&) = delete; - thread_stream_pools(thread_stream_pools&&) = delete; - thread_stream_pools& operator=(thread_stream_pools&&) = delete; + thread_stream_resources(thread_stream_resources const&) = delete; + thread_stream_resources& operator=(thread_stream_resources const&) = delete; + thread_stream_resources(thread_stream_resources&&) = delete; + thread_stream_resources& operator=(thread_stream_resources&&) = delete; cuda_stream_pool& pool_for(rmm::cuda_device_id device_id) { auto*& pool = _pools[device_id.value()]; - if (pool == nullptr) { pool = pool_registry().acquire(device_id); } + if (pool == nullptr) { pool = pool_registry().acquire_pool(device_id); } return *pool; } + + cuda_event& event_for(rmm::cuda_device_id device_id) + { + auto*& event = _events[device_id.value()]; + if (event == nullptr) { event = pool_registry().acquire_event(device_id); } + return *event; + } }; +thread_stream_resources& current_thread_resources() +{ + thread_local thread_stream_resources resources; + return resources; +} + +/** + * @brief Returns a cudaEvent_t the calling thread can use on the current device. + * + * The event is reused by every fork and join the thread performs on that device, and is recycled + * for another thread when this one exits. + */ +cudaEvent_t event_for_thread() +{ + return current_thread_resources().event_for(get_current_cuda_device()); +} + } // namespace /** @@ -275,8 +323,7 @@ class thread_stream_pools { */ cuda_stream_pool& current_cuda_stream_pool() { - thread_local thread_stream_pools pools; - return pools.pool_for(get_current_cuda_device()); + return current_thread_resources().pool_for(get_current_cuda_device()); } std::vector fork_streams(cuda::stream_ref stream, std::size_t count)