diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index d68527d4c6ef..090d786b8b85 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -1,26 +1,30 @@ /* - * 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 */ #pragma once #include -#include -#include +#include #include +#include #include 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 `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 { 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; @@ -28,70 +32,80 @@ class cuda_stream_pool { cuda_stream_pool& operator=(cuda_stream_pool&&) = delete; /** - * @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; - - /** - * @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. - */ - virtual rmm::cuda_stream_view get_stream(stream_id_type stream_id) = 0; - - /** - * @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()`. + * @brief Get a single stream from the pool. * - * This function is thread safe with respect to other calls to the same function. + * @note Use `get_streams` to obtain multiple streams. Repeated single-stream requests are not + * guaranteed to return different streams. * - * @param count The number of stream views to return. - * @return Vector containing `count` stream views. + * @return Stream reference. */ - virtual std::vector get_streams(std::size_t count) = 0; + virtual cuda::stream_ref get_stream() = 0; /** - * @brief Get the number of unique stream objects in the pool. + * @brief Get a vector of `cuda::stream_ref` objects from the pool. * - * This function is thread safe with respect to other calls to the same function. + * The returned streams are distinct unless `count` is greater than the maximum number of streams + * the pool provides, in which case streams are repeated. * - * @return the number of stream objects in the pool + * @param count The number of stream references to return. + * @return Vector containing `count` stream references. */ - [[nodiscard]] virtual std::size_t get_stream_pool_size() const = 0; + virtual std::vector get_streams(std::size_t count) = 0; protected: cuda_stream_pool() = default; }; /** - * @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 global stream pool. + * @brief Get the stream pool the calling thread should use for the current device. + * + * 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. + * + * @return Reference to the calling thread's stream pool for the current device. */ -cuda_stream_pool& global_cuda_stream_pool(); +cuda_stream_pool& current_cuda_stream_pool(); /** - * @brief Acquire a set of `cuda_stream_view` objects and synchronize them to an event on another - * stream. + * @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& +global_cuda_stream_pool() +{ + return current_cuda_stream_pool(); +} + +/** + * @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 + * 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 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`. + * 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} @@ -106,19 +120,19 @@ cuda_stream_pool& 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(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. + * @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, rmm::cuda_stream_view stream); +void join_streams(std::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 7a5de2721392..7972b8067741 100644 --- a/cpp/src/io/parquet/page_enc.cu +++ b/cpp/src/io/parquet/page_enc.cu @@ -3488,55 +3488,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 ae615db090b9..4dad0efa6692 100644 --- a/cpp/src/io/parquet/page_string_decode.cu +++ b/cpp/src/io/parquet/page_string_decode.cu @@ -976,7 +976,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()); } @@ -985,12 +985,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 a9076e415544..f2d9de7224ff 100644 --- a/cpp/src/io/text/multibyte_split.cu +++ b/cpp/src/io/text/multibyte_split.cu @@ -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 7cb277f6797a..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::global_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,10 +211,10 @@ 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::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.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..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::global_cuda_stream_pool().get_stream().value()}; // 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 diff --git a/cpp/src/utilities/stream_pool.cpp b/cpp/src/utilities/stream_pool.cpp index 6c65a4067347..5668fdc43382 100644 --- a/cpp/src/utilities/stream_pool.cpp +++ b/cpp/src/utilities/stream_pool.cpp @@ -3,26 +3,36 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include #include #include #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, for a single device. Sized to cover +// 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 +// 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 @@ -81,6 +91,14 @@ rmm::cuda_device_id get_current_cuda_device() return rmm::cuda_device_id{device_id}; } +/** + * @brief Returns the configured maximum number of streams a single pool will hold. + */ +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. * @@ -102,29 +120,45 @@ 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 `_max_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}; + std::size_t const _max_size{configured_max_pool_size()}; - 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(stream_id_type stream_id) override + /** + * @brief Creates streams until the pool holds `size` of them, or has reached `_max_size`. + */ + void grow_to(std::size_t size) { - return _pool.get_stream(stream_id); + auto const device = cuda::device_ref{get_current_cuda_device().value()}; + auto const target = std::min(size, _max_size); + while (_streams.size() < target) { + // `cuda::stream` creates non-blocking streams. + _streams.emplace_back(device); + } } - std::vector get_streams(std::size_t count) override + public: + cuda::stream_ref get_stream() override { return get_streams(1).front(); } + + 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()); + // 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); + for (std::size_t i = 0; i < count; i++) { + streams.emplace_back(_streams[(first + i) % _streams.size()]); } return streams; } - - [[nodiscard]] std::size_t get_stream_pool_size() const override { return STREAM_POOL_SIZE; } }; /** @@ -132,62 +166,136 @@ class rmm_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(); } + + std::vector get_streams(std::size_t count) override { - return cudf::get_default_stream(); + return std::vector(count, cudf::get_default_stream()); } +}; + +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(); +} + +namespace { + +/** + * @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. + */ +class stream_pool_registry { + std::mutex _mutex; + std::vector> _free_pools; + + public: + stream_pool_registry() : _free_pools(get_num_cuda_devices()) {} - std::vector get_streams(std::size_t count) override + /** + * @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) { - return std::vector(count, cudf::get_default_stream()); + { + 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_cuda_stream_pool(); } - [[nodiscard]] std::size_t get_stream_pool_size() const override { return 1UL; } + /** + * @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); + } }; -cuda_stream_pool* create_global_cuda_stream_pool() +stream_pool_registry& pool_registry() { - if (getenv("LIBCUDF_USE_DEBUG_STREAM_POOL")) return new debug_cuda_stream_pool(); - return new rmm_cuda_stream_pool(); + static auto* registry = new stream_pool_registry(); + return *registry; } /** - * @brief Returns a reference to the global stream pool for the current device. - * @return `cuda_stream_pool` valid on the current device. + * @brief Owns the calling thread's pool for each device, and retires them when the thread exits. */ -cuda_stream_pool& global_cuda_stream_pool() -{ - // 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(); +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]); + } + } + } - std::lock_guard const lock(mutex); - if (pools[device_id.value()] == nullptr) { - pools[device_id.value()] = create_global_cuda_stream_pool(); + 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; + + 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& current_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) +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 = current_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(std::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/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/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 c80650cdf22f..6fcf82a68e9b 100644 --- a/cpp/tests/utilities/identify_stream_usage.cpp +++ b/cpp/tests/utilities/identify_stream_usage.cpp @@ -67,21 +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 - { - return cudf::test::get_default_stream(); - } + cuda::stream_ref get_stream() 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; } }; -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 @@ -221,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); 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..571d23d0a280 --- /dev/null +++ b/cpp/tests/utilities_tests/stream_pool_tests.cpp @@ -0,0 +1,97 @@ +/* + * 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 +#include + +class StreamPoolTest : public cudf::test::BaseFixture {}; + +namespace { + +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{}; + std::transform(streams.begin(), streams.end(), std::back_inserter(values), [](auto stream) { + return stream.get(); + }); + return values; +} + +} // namespace + +TEST_F(StreamPoolTest, ConcurrentThreadsGetDistinctStreams) +{ + auto constexpr num_requests = 20; + auto constexpr num_streams = 8; + + // 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++) { + auto const streams = get_hashable_streams(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(); + + EXPECT_FALSE(first_streams.empty()); + EXPECT_TRUE(std::none_of(first_streams.begin(), first_streams.end(), [&](auto stream) { + return second_streams.contains(stream); + })); +} + +TEST_F(StreamPoolTest, RequestLargerThanPoolRepeatsStreams) +{ + 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); +} + +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_hashable_streams(128); + first_thread_streams.insert(streams.begin(), streams.end()); + }).join(); + + std::vector second_thread_streams; + 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 + EXPECT_TRUE(std::all_of(second_thread_streams.begin(), + second_thread_streams.end(), + [&](auto stream) { return first_thread_streams.contains(stream); })); +} diff --git a/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd b/python/pylibcudf/pylibcudf/libcudf/detail/utilities/stream_pool.pxd index 399a868db717..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 @@ -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); } } """