cuda_stream_pool.hpp

/*
 * Copyright (c) 2020-2021, 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 <rmm/detail/error.hpp>
 
#include <atomic>
#include <cstddef>
#include <vector>
 
namespace rmm {
 
class cuda_stream_pool {
 public:
  static constexpr std::size_t default_size{16};  
 
  explicit cuda_stream_pool(std::size_t pool_size = default_size) : streams_(pool_size)
  {
    RMM_EXPECTS(pool_size > 0, "Stream pool size must be greater than zero");
  }
  ~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 const&) = delete;
 
  rmm::cuda_stream_view get_stream() const noexcept
  {
    return streams_[(next_stream++) % streams_.size()].view();
  }
 
  rmm::cuda_stream_view get_stream(std::size_t stream_id) const
  {
    return streams_[stream_id % streams_.size()].view();
  }
 
  std::size_t get_pool_size() const noexcept { return streams_.size(); }
 
 private:
  std::vector<rmm::cuda_stream> streams_;
  mutable std::atomic_size_t next_stream{};
};
 
}  // namespace rmm