Add CUDA stream pool. (#11458)

Author: trivialfis

src/common/cuda_stream_pool.cuh

@@ -0,0 +1,29 @@
+/**
+ * Copyright 2025, XGBoost contributors
+ */
+#pragma once
+#include <atomic>   // for atomic
+#include <cstddef>  // for size_t
+#include <vector>   // for vector
+
+#include "device_helpers.cuh"  // for CUDAStreamView, CUDAStream
+
+namespace xgboost::curt {
+// rmm cuda_stream_pool
+class StreamPool {
+  mutable std::atomic<std::size_t> next_{0};
+  std::vector<dh::CUDAStream> stream_;
+
+ public:
+  explicit StreamPool(std::size_t n) : stream_(n) {}
+  ~StreamPool() = default;
+  StreamPool(StreamPool const& that) = delete;
+  StreamPool& operator=(StreamPool const& that) = delete;
+
+  [[nodiscard]] dh::CUDAStreamView operator[](std::size_t i) const { return stream_[i].View(); }
+  [[nodiscard]] dh::CUDAStreamView Next() const {
+    return stream_[(next_++) % stream_.size()].View();
+  }
+  [[nodiscard]] std::size_t Size() const { return stream_.size(); }
+};
+}  // namespace xgboost::curt

src/data/ellpack_page_source.cu

@@ -10,6 +10,7 @@
 
 #include "../common/common.h"               // for HumanMemUnit, safe_cuda
 #include "../common/cuda_rt_utils.h"        // for SetDevice
+#include "../common/cuda_stream_pool.cuh"   // for StreamPool
 #include "../common/device_helpers.cuh"     // for CUDAStreamView, DefaultStream
 #include "../common/ref_resource_view.cuh"  // for MakeFixedVecWithCudaMalloc
 #include "../common/resource.cuh"           // for PrivateCudaMmapConstStream
@@ -25,11 +26,12 @@ namespace xgboost::data {
 /**
  * Cache
  */
-EllpackMemCache::EllpackMemCache(EllpackCacheInfo cinfo)
+EllpackMemCache::EllpackMemCache(EllpackCacheInfo cinfo, std::int32_t n_workers)
     : cache_mapping{std::move(cinfo.cache_mapping)},
       buffer_bytes{std::move(cinfo.buffer_bytes)},
       buffer_rows{std::move(cinfo.buffer_rows)},
-      cache_host_ratio{cinfo.cache_host_ratio} {
+      cache_host_ratio{cinfo.cache_host_ratio},
+      streams{std::make_unique<curt::StreamPool>(n_workers)} {
   CHECK_EQ(buffer_bytes.size(), buffer_rows.size());
   CHECK(!detail::HostRatioIsAuto(this->cache_host_ratio));
   CHECK_GE(this->cache_host_ratio, 0) << error::CacheHostRatioInvalid();
@@ -289,7 +291,8 @@ EllpackCacheStreamPolicy<S, F>::CreateWriter(StringView, std::uint32_t iter) {
     CHECK(!detail::HostRatioIsAuto(this->CacheInfo().cache_host_ratio));
     CHECK_GE(this->CacheInfo().cache_host_ratio, 0.0);
     CHECK_LE(this->CacheInfo().cache_host_ratio, 1.0);
-    this->p_cache_ = std::make_unique<EllpackMemCache>(this->CacheInfo());
+    constexpr std::int32_t kMaxGpuExtMemWorkers = 4;
+    this->p_cache_ = std::make_unique<EllpackMemCache>(this->CacheInfo(), kMaxGpuExtMemWorkers);
   }
   auto fo = std::make_unique<EllpackHostCacheStream>(this->p_cache_);
   if (iter == 0) {

src/data/ellpack_page_source.h

@@ -23,6 +23,10 @@
 #include "xgboost/data.h"                   // for BatchParam
 #include "xgboost/span.h"                   // for Span
 
+namespace xgboost::curt {
+class StreamPool;
+}
+
 namespace xgboost::data {
 struct EllpackCacheInfo {
   BatchParam param;
@@ -66,7 +70,9 @@ struct EllpackMemCache {
   std::vector<bst_idx_t> const buffer_rows;
   float const cache_host_ratio;
 
-  explicit EllpackMemCache(EllpackCacheInfo cinfo);
+  std::unique_ptr<curt::StreamPool> streams;
+
+  explicit EllpackMemCache(EllpackCacheInfo cinfo, std::int32_t n_workers);
   ~EllpackMemCache();
 
   // The number of bytes of the entire cache.

tests/cpp/common/test_cuda_rt_utils.cu

@@ -0,0 +1,31 @@
+/**
+ * Copyright 2025, XGBoost contributors
+ */
+
+#include <gtest/gtest.h>
+
+#include <cstdint>  // for int32_t
+#include <set>      // for set
+
+#include "../../../src/common/cuda_stream_pool.cuh"
+
+namespace xgboost::curt {
+TEST(RtUtils, StreamPool) {
+  auto n_streams = 16;
+  auto pool = std::make_unique<StreamPool>(n_streams);
+  std::set<cudaStream_t> hdls;
+
+  for (std::int32_t i = 0; i < n_streams; ++i) {
+    hdls.insert(cudaStream_t{pool->Next()});
+  }
+
+  ASSERT_EQ(hdls.size(), n_streams);
+  ASSERT_EQ(hdls.size(), pool->Size());
+
+  for (std::int32_t i = 0; i < n_streams; ++i) {
+    hdls.insert(cudaStream_t{pool->Next()});
+  }
+  ASSERT_EQ(hdls.size(), n_streams);
+  ASSERT_EQ(hdls.size(), pool->Size());
+}
+}  // namespace xgboost::curt