Add a CUDA stream pool

conda/recipes/librmm/meta.yaml

@@ -36,6 +36,7 @@ test:
     - test -f $PREFIX/include/rmm/logger.hpp
     - test -f $PREFIX/include/rmm/cuda_stream.hpp
     - test -f $PREFIX/include/rmm/cuda_stream_view.hpp
+    - test -f $PREFIX/include/rmm/cuda_stream_pool.hpp
     - test -f $PREFIX/include/rmm/device_uvector.hpp
     - test -f $PREFIX/include/rmm/device_scalar.hpp
     - test -f $PREFIX/include/rmm/device_buffer.hpp

include/rmm/cuda_stream_pool.hpp

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <rmm/cuda_stream.hpp>
+#include <rmm/cuda_stream_view.hpp>
+
+#include <atomic>
+#include <vector>
+
+namespace rmm {
+
+/**
+ * @brief A pool of CUDA streams.
+ *
+ * Provides fast access to a circular buffer of `cuda_stream` objects.
+ *
+ */
+class cuda_stream_pool {
+ public:
+  static constexpr std::size_t default_size{16};  ///< Default stream pool size
+
+  /**
+   * @brief Construct a new cuda stream pool object of the given size
+   *
+   * @param pool_size The number of streams in the pool
+   */
+  explicit cuda_stream_pool(std::size_t pool_size = default_size) : streams_{pool_size} {}
+  ~cuda_stream_pool() = default;
+
+  cuda_stream_pool(cuda_stream_pool&&)      = delete;
+  cuda_stream_pool(cuda_stream_pool const&) = delete;
+  cuda_stream_pool& operator=(cuda_stream_pool&&) = 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 rmm::cuda_stream_view
+   */
+  rmm::cuda_stream_view get_stream() { return streams_[(next_stream++) % streams_.size()].view(); }
+
+ private:
+  std::vector<rmm::cuda_stream> streams_;
+  std::atomic_size_t next_stream{};
+};
+
+}  // namespace rmm

tests/CMakeLists.txt

@@ -131,8 +131,8 @@ ConfigureTest(HOST_MR_TEST "${HOST_MR_TEST_SRC}")
 
 # cuda stream tests
 
-set(CUDA_STREAM_TEST_SRC "${CMAKE_CURRENT_SOURCE_DIR}/cuda_stream_tests.cpp")
-
+set(CUDA_STREAM_TEST_SRC "${CMAKE_CURRENT_SOURCE_DIR}/cuda_stream_tests.cpp"
+                         "${CMAKE_CURRENT_SOURCE_DIR}/cuda_stream_pool_tests.cpp")
 ConfigureTest(CUDA_STREAM_TEST "${CUDA_STREAM_TEST_SRC}")
 
 # device buffer tests

tests/cuda_stream_pool_tests.cpp

@@ -0,0 +1,60 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <rmm/cuda_stream_pool.hpp>
+#include <rmm/detail/error.hpp>
+#include <rmm/device_uvector.hpp>
+
+#include <gtest/gtest.h>
+
+#include <cuda_runtime_api.h>
+
+struct CudaStreamPoolTest : public ::testing::Test {
+  rmm::cuda_stream_pool pool{};
+};
+
+TEST_F(CudaStreamPoolTest, Unequal)
+{
+  auto const stream_a = this->pool.get_stream();
+  auto const stream_b = this->pool.get_stream();
+
+  EXPECT_NE(stream_a, stream_b);
+}
+
+TEST_F(CudaStreamPoolTest, Nondefault)
+{
+  auto const stream_a = this->pool.get_stream();
+  auto const stream_b = this->pool.get_stream();
+
+  // pool streams are explicit, non-default streams
+  EXPECT_FALSE(stream_a.is_default());
+  EXPECT_FALSE(stream_a.is_per_thread_default());
+}
+
+TEST_F(CudaStreamPoolTest, ValidStreams)
+{
+  auto const stream_a = this->pool.get_stream();
+  auto const stream_b = this->pool.get_stream();
+
+  // Operations on the streams should work correctly and without throwing exceptions
+  auto v = rmm::device_uvector<std::uint8_t>{100, stream_a};
+  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 x  = v2.front_element(stream_b);
+  EXPECT_EQ(x, 0xcc);
+}

tests/cuda_stream_tests.cpp

@@ -14,14 +14,14 @@
  * limitations under the License.
  */
 
-#include "gtest/gtest.h"
-
 #include <rmm/cuda_stream.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 
 #include <cuda_runtime_api.h>
 
+#include <gtest/gtest.h>
+
 struct CudaStreamTest : public ::testing::Test {
 };