Make the CUDA stream pool per-thread

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
  */
 
@@ -8,90 +8,104 @@
 #include <cudf/utilities/export.hpp>
 #include <cudf/utilities/span.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream_ref>
 
 #include <cstddef>
 #include <vector>
 
 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;
   cuda_stream_pool& operator=(cuda_stream_pool const&) = delete;
   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<rmm::cuda_stream_view> 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 set 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<cuda::stream_ref> 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
+ * @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 set of `cuda::stream_ref` 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 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<rmm::cuda_stream_view> fork_streams(rmm::cuda_stream_view stream,
-                                                              std::size_t count);
+[[nodiscard]] std::vector<cuda::stream_ref> fork_streams(cuda::stream_ref stream,
+                                                         std::size_t count);
 
 /**
  * @brief Synchronize a stream to an event on a set of streams.
  *
  * @param streams Streams to wait on.
  * @param stream Joined stream that synchronizes with the waited-on streams.
  */
-void join_streams(host_span<rmm::cuda_stream_view const> streams, rmm::cuda_stream_view stream);
+void join_streams(host_span<cuda::stream_ref const> streams, cuda::stream_ref stream);
 
 }  // namespace detail
 }  // namespace CUDF_EXPORT cudf

cpp/src/io/parquet/page_enc.cu

@@ -3488,55 +3488,55 @@ void EncodePages(device_span<EncPage> pages,
   int s_idx = 0;
   if (BitAnd(kernel_mask, encode_kernel_mask::PLAIN) != 0) {
     auto const strm = streams[s_idx++];
-    gpuEncodePageLevels<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodePageLevels<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       pages, write_v2_headers, encode_kernel_mask::PLAIN);
     CUDF_CUDA_TRY(cudaGetLastError());
-    gpuEncodePages<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodePages<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       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<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodePageLevels<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       pages, write_v2_headers, encode_kernel_mask::BYTE_STREAM_SPLIT);
     CUDF_CUDA_TRY(cudaGetLastError());
-    gpuEncodePages<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodePages<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       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<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodePageLevels<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       pages, write_v2_headers, encode_kernel_mask::DELTA_BINARY);
     CUDF_CUDA_TRY(cudaGetLastError());
     gpuEncodeDeltaBinaryPages<encode_block_size>
-      <<<num_pages, encode_block_size, 0, strm.value()>>>(pages, comp_in, comp_out, comp_results);
+      <<<num_pages, encode_block_size, 0, strm.get()>>>(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<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodePageLevels<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       pages, write_v2_headers, encode_kernel_mask::DELTA_LENGTH_BA);
     CUDF_CUDA_TRY(cudaGetLastError());
     gpuEncodeDeltaLengthByteArrayPages<encode_block_size>
-      <<<num_pages, encode_block_size, 0, strm.value()>>>(pages, comp_in, comp_out, comp_results);
+      <<<num_pages, encode_block_size, 0, strm.get()>>>(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<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodePageLevels<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       pages, write_v2_headers, encode_kernel_mask::DELTA_BYTE_ARRAY);
     CUDF_CUDA_TRY(cudaGetLastError());
     gpuEncodeDeltaByteArrayPages<encode_block_size>
-      <<<num_pages, encode_block_size, 0, strm.value()>>>(pages, comp_in, comp_out, comp_results);
+      <<<num_pages, encode_block_size, 0, strm.get()>>>(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<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodePageLevels<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       pages, write_v2_headers, encode_kernel_mask::DICTIONARY);
     CUDF_CUDA_TRY(cudaGetLastError());
-    gpuEncodeDictPages<encode_block_size><<<num_pages, encode_block_size, 0, strm.value()>>>(
+    gpuEncodeDictPages<encode_block_size><<<num_pages, encode_block_size, 0, strm.get()>>>(
       pages, comp_in, comp_out, comp_results, write_v2_headers);
     CUDF_CUDA_TRY(cudaGetLastError());
   }

cpp/src/io/parquet/page_string_decode.cu

@@ -976,7 +976,7 @@ void compute_page_string_sizes_pass1(cudf::detail::hostdevice_span<PageInfo> 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<<<dim_grid, dim_delta, 0, streams[s_idx++].value()>>>(
+    compute_delta_page_string_sizes_kernel<<<dim_grid, dim_delta, 0, streams[s_idx++].get()>>>(
       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<PageInfo> pag
     compute_delta_length_page_string_sizes_kernel<<<dim_grid,
                                                     dim_delta,
                                                     0,
-                                                    streams[s_idx++].value()>>>(
+                                                    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<<<dim_grid, dim_block, 0, streams[s_idx++].value()>>>(
+    compute_page_string_sizes_kernel<<<dim_grid, dim_block, 0, streams[s_idx++].get()>>>(
       pages.device_ptr(), chunks, page_mask, page_string_offset_indices, min_row, num_rows);
     CUDF_CUDA_TRY(cudaGetLastError());
   }

cpp/src/io/text/multibyte_split.cu

@@ -405,20 +405,20 @@ std::unique_ptr<cudf::column> multibyte_split(cudf::io::text::data_chunk_source
       multibyte_split_init_kernel<<<tiles_in_launch,
                                     THREADS_PER_TILE,
                                     0,
-                                    scan_stream.value()>>>(  //
+                                    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<<<tiles_in_launch,
                             THREADS_PER_TILE,
                             0,
-                            scan_stream.value()>>>(  //
+                            scan_stream.get()>>>(  //
           base_tile_idx,
           chunk_offset,
           row_offset_storage.size(),
@@ -431,7 +431,7 @@ std::unique_ptr<cudf::column> multibyte_split(cudf::io::text::data_chunk_source
         multibyte_split_kernel<<<tiles_in_launch,
                                  THREADS_PER_TILE,
                                  0,
-                                 scan_stream.value()>>>(  //
+                                 scan_stream.get()>>>(  //
           base_tile_idx,
           chunk_offset,
           row_offset_storage.size(),
@@ -492,7 +492,7 @@ std::unique_ptr<cudf::column> 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;

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<uint8_t>{dst, count},
                               device_span<uint8_t const>{
                                 reinterpret_cast<uint8_t const*>(_d_buffer.data() + offset), count},
@@ -211,10 +211,10 @@ class device_buffer_source final : public datasource {
   std::unique_ptr<buffer> 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<std::byte const>{_d_buffer.data() + offset, count}, stream);
-    stream.synchronize();
+    stream.sync();
     return std::make_unique<owning_buffer<cudf::detail::host_vector<std::byte>>>(std::move(h_data));
   }
 

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::current_cuda_stream_pool().get_stream()};
 
   // Wrapped in shared_ptr so the outer class is copyable (required by any_resource)
   std::shared_ptr<fallback_state> fallback_{std::make_shared<fallback_state>()};