cuda_stream_pool.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/cuda_stream.hpp>
#include <rmm/detail/export.hpp>
 
#include <cuda/stream>
 
#include <atomic>
#include <cstddef>
#include <vector>
 
RMM_NAMESPACE_BEGIN

class cuda_stream_pool {
 public:
  static constexpr std::size_t default_size{16};  

  explicit cuda_stream_pool(std::size_t pool_size    = default_size,
                            cuda_stream::flags flags = cuda_stream::flags::sync_default);
  ~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;

  cuda::stream_ref get_stream() const noexcept;

  cuda::stream_ref get_stream(std::size_t stream_id) const;

  std::size_t get_pool_size() const noexcept;
 
 private:
  std::vector<rmm::cuda_stream> streams_;
  mutable std::atomic_size_t next_stream{};
};
  // end of group
RMM_NAMESPACE_END