owning_wrapper.hpp

前往此文件的文档。

/*
 * 版权所有 (c) 2020-2025, NVIDIA CORPORATION.
 *
 * 根据 Apache 许可证 2.0 版（“许可证”）获得许可；
 * 除非遵守许可证的规定，否则您不得使用此文件。
 * 您可以获取许可证的副本，网址为
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则，软件
 * 根据许可证分发，是按“现状”基础分发，
 * 不带任何明示或暗示的担保或条件。
 * 有关管理权限和
 * 限制的特定语言，请参见许可证。
 */
#pragma once

#include <rmm/detail/export.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>

#include <memory>
#include <utility>

namespace RMM_NAMESPACE {
namespace mr {
namespace detail {
template <typename Resource, typename UpstreamTuple, std::size_t... Indices, typename... Args>
auto make_resource_impl(UpstreamTuple const& upstreams,
 std::index_sequence<Indices...>,
 Args&&... args)
{
 return std::make_unique<Resource>(std::get<Indices>(upstreams).get()...,
 std::forward<Args>(args)...);
}

template <typename Resource, typename... Upstreams, typename... Args>
auto make_resource(std::tuple<std::shared_ptr<Upstreams>...> const& upstreams, Args&&... args)
{
 return make_resource_impl<Resource>(
 upstreams, std::index_sequence_for<Upstreams...>{}, std::forward<Args>(args)...);
}
} // namespace detail

template <typename Resource, typename... Upstreams>
class owning_wrapper : public device_memory_resource {
 public
 using upstream_tuple =
 std::tuple<std::shared_ptr<Upstreams>...>;

 template <typename... Args>
 owning_wrapper(upstream_tuple upstreams, Args&&... args)
 : upstreams_{std::move(upstreams)},
 wrapped_{detail::make_resource<Resource>(upstreams_, std::forward<Args>(args)...)}
 {
 }

 [[nodiscard]] Resource const& wrapped() const noexcept { return *wrapped_; }

 [[nodiscard]] Resource& wrapped() noexcept { return *wrapped_; }

 private
 void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
 {
 return wrapped().allocate(bytes, stream);
 }

 void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
 {
 wrapped().deallocate(ptr, bytes, stream);
 }

 [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
 {
 if (this == &other) { return true; }
 auto casted = dynamic_cast<owning_wrapper<Resource, Upstreams...> const*>(&other);
 if (nullptr != casted) { return wrapped().is_equal(casted->wrapped()); }
 return wrapped().is_equal(other);
 }

 upstream_tuple upstreams_;
 std::unique_ptr<Resource> wrapped_;
};

template <template <typename...> class Resource, typename... Upstreams, typename... Args>
auto make_owning_wrapper(std::tuple<std::shared_ptr<Upstreams>...> upstreams, Args&&... args)
{
 return std::make_shared<owning_wrapper<Resource<Upstreams...>, Upstreams...>>(
 std::move(upstreams), std::forward<Args>(args)...);
}

template <template <typename> class Resource, typename Upstream, typename... Args>
auto make_owning_wrapper(std::shared_ptr<Upstream> upstream, Args&&... args)
{
 return make_owning_wrapper<Resource>(std::make_tuple(std::move(upstream)),
 std::forward<Args>(args)...);
}
 // end of group
} // namespace mr
} // namespace RMM_NAMESPACE