cuda_stream_pool.hpp

转到此文件的文档。

/*
 * Copyright (c) 2020-2024, 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
 *
 * https://apache.ac.cn/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 <rmm/detail/export.hpp>

#include <atomic>
#include <cstddef>
#include <vector>

namespace RMM_NAMESPACE {
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, "流池大小必须大于零");
 }
 ~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{};
};
 // 组结束
} // 命名空间 RMM_NAMESPACE