cuda_stream_pool.hpp
1 /*
2  * Copyright (c) 2020-2021, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <rmm/cuda_stream.hpp>
20 #include <rmm/cuda_stream_view.hpp>
21 #include <rmm/detail/error.hpp>
22 
23 #include <atomic>
24 #include <cstddef>
25 #include <vector>
26 
27 namespace rmm {
28 
38  public:
39  static constexpr std::size_t default_size{16};
40 
47  explicit cuda_stream_pool(std::size_t pool_size = default_size) : streams_(pool_size)
48  {
49  RMM_EXPECTS(pool_size > 0, "Stream pool size must be greater than zero");
50  }
51  ~cuda_stream_pool() = default;
52 
54  cuda_stream_pool(cuda_stream_pool const&) = delete;
55  cuda_stream_pool& operator=(cuda_stream_pool&&) = delete;
56  cuda_stream_pool& operator=(cuda_stream_pool const&) = delete;
57 
66  {
67  return streams_[(next_stream++) % streams_.size()].view();
68  }
69 
80  rmm::cuda_stream_view get_stream(std::size_t stream_id) const
81  {
82  return streams_[stream_id % streams_.size()].view();
83  }
84 
92  std::size_t get_pool_size() const noexcept { return streams_.size(); }
93 
94  private:
95  std::vector<rmm::cuda_stream> streams_;
96  mutable std::atomic_size_t next_stream{};
97 };
98 
99 } // namespace rmm
rmm::cuda_stream_pool
A pool of CUDA streams.
Definition: cuda_stream_pool.hpp:37
rmm::cuda_stream_pool::cuda_stream_pool
cuda_stream_pool(std::size_t pool_size=default_size)
Construct a new cuda stream pool object of the given non-zero size.
Definition: cuda_stream_pool.hpp:47
rmm::cuda_stream_pool::get_stream
rmm::cuda_stream_view get_stream() const noexcept
Get a cuda_stream_view of a stream in the pool.
Definition: cuda_stream_pool.hpp:65
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::cuda_stream_pool::get_stream
rmm::cuda_stream_view get_stream(std::size_t stream_id) const
Get a cuda_stream_view of the stream associated with stream_id. Equivalent values of stream_id return...
Definition: cuda_stream_pool.hpp:80
rmm::cuda_stream_pool::default_size
static constexpr std::size_t default_size
Default stream pool size.
Definition: cuda_stream_pool.hpp:39
rmm::cuda_stream_pool::get_pool_size
std::size_t get_pool_size() const noexcept
Get the number of streams in the pool.
Definition: cuda_stream_pool.hpp:92