thread_safe_resource_adaptor.hpp

转到此文件的文档。

/*
 * Copyright (c) 2020-2024, NVIDIA CORPORATION.
 *
 * 根据 Apache License, Version 2.0 (“许可证”) 获得许可；
 * 除非符合许可证，否则您不得使用此文件。
 * 您可以获取许可证的副本
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则根据许可证分发的软件
 * 按“原样”分发，不附带任何明示或暗示的担保或条件。
 * 有关管理权限和
 * 许可证下的限制，请参阅许可证。
 */

#pragma once

#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/device_memory_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cstddef>
#include <mutex>

namespace RMM_NAMESPACE {
namespace mr {
/// @addtogroup memory_resources
/// @{template <typename Upstream>
class thread_safe_resource_adaptor final : public device_memory_resource {
 public
 /// Type of lock used to synchronize access.
 using lock_t = std::lock_guard<std::mutex>;

 /// @brief Constructs a new thread safe resource adaptor using `upstream` to satisfy
 /// allocation requests.
 thread_safe_resource_adaptor(device_async_resource_ref upstream) : upstream_{upstream} {}

 /// @brief Constructs a new thread safe resource adaptor using `upstream` to satisfy
 /// allocation requests.
 thread_safe_resource_adaptor(Upstream* upstream)
 : upstream_{to_device_async_resource_ref_checked(upstream)}
 {
 }

 thread_safe_resource_adaptor() = delete;
 ~thread_safe_resource_adaptor() override = default;
 thread_safe_resource_adaptor(thread_safe_resource_adaptor const&) = delete;
 thread_safe_resource_adaptor(thread_safe_resource_adaptor&&) = delete;
 thread_safe_resource_adaptor& operator=(thread_safe_resource_adaptor const&) = delete;
 thread_safe_resource_adaptor& operator=(thread_safe_resource_adaptor&&) = delete;

 /// @brief Returns a rmm::device_async_resource_ref to the upstream resource
 [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
 {
 return upstream_;
 }

 /// @brief Allocates memory from the upstream resource. Synchronizes access to the
 /// upstream resource.
 /// @copydoc device_memory_resource::do_allocate(std::size_t bytes, cuda_stream_view stream)
 private
 void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
 {
 lock_t lock(mtx);
 return get_upstream_resource().allocate_async(bytes, stream);
 }

 /// @brief Deallocates memory from the upstream resource. Synchronizes access to the
 /// upstream resource.
 /// @copydoc device_memory_resource::do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream)
 void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
 {
 lock_t lock(mtx);
 get_upstream_resource().deallocate_async(ptr, bytes, stream);
 }

 /// @brief Compares for equality.
 /// @copydoc device_memory_resource::do_is_equal(device_memory_resource const& other) const noexcept