owning_wrapper.hpp

Go to the documentation of this file.

/*
 * 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
 *
 *     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/detail/export.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
 
#include <functional>
#include <iostream>
#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