Adds a linear accessor to RMM cuda stream pool

include/rmm/cuda_stream_pool.hpp

@@ -18,6 +18,7 @@
 
 #include <rmm/cuda_stream.hpp>
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/detail/error.hpp>
 
 #include <atomic>
 #include <vector>
@@ -61,6 +62,33 @@ class cuda_stream_pool {
     return streams_[(next_stream++) % streams_.size()].view();
   }
 
+  /**
+   * @brief Get a `cuda_stream_view` of the stream at stream_id index the pool.
+   *
+   * This function is thread safe with respect to other calls to the same function.
+   *
+   * @throws rmm::out_of_range exception if `stream_index >= size()`
+   *
+   * @param stream_id The index of the stream in the pool
+   *
+   * @return rmm::cuda_stream_view
+   */
+  rmm::cuda_stream_view get_stream(std::size_t stream_index) const
+  {
+    RMM_EXPECTS(
+      stream_index < streams_.size(), rmm::out_of_range, "Attempt to access out of bounds stream.");
+    return streams_[stream_index].view();
+  }
+
+  /**
+   * @brief Get the number of streams in the pool.
+   *
+   * This function is thread safe with respect to other calls to the same function.
+   *
+   * @return the number of streams in the pool
+   */
+  size_t get_pool_size() const noexcept { return streams_.size(); }
+
  private:
   std::vector<rmm::cuda_stream> streams_;
   mutable std::atomic_size_t next_stream{};

tests/cuda_stream_pool_tests.cpp

@@ -54,7 +54,26 @@ TEST_F(CudaStreamPoolTest, ValidStreams)
   RMM_CUDA_TRY(cudaMemsetAsync(v.data(), 0xcc, 100, stream_a.value()));
   stream_a.synchronize();
 
-  auto v2 = rmm::device_uvector<uint8_t>{v, stream_b};
+  auto v2 = rmm::device_uvector<std::uint8_t>{v, stream_b};
   auto x  = v2.front_element(stream_b);
   EXPECT_EQ(x, 0xcc);
 }
+
+TEST_F(CudaStreamPoolTest, PoolSize) { EXPECT_GE(this->pool.get_pool_size(), 1); }
+
+TEST_F(CudaStreamPoolTest, OutOfBoundLinearAccess)
+{
+  EXPECT_NO_THROW(this->pool.get_stream(this->pool.get_pool_size() - 1));
+  EXPECT_THROW(this->pool.get_stream(this->pool.get_pool_size()), rmm::out_of_range);
+}
+
+TEST_F(CudaStreamPoolTest, ValidLinearAccess)
+{
+  auto const stream_a = this->pool.get_stream(0);
+  auto const stream_b = this->pool.get_stream(1);
+  EXPECT_NE(stream_a, stream_b);
+  EXPECT_FALSE(stream_a.is_default());
+  EXPECT_FALSE(stream_a.is_per_thread_default());
+  EXPECT_FALSE(stream_b.is_default());
+  EXPECT_FALSE(stream_b.is_per_thread_default());
+}