All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
owning_wrapper.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2020-2025, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #pragma once
17 
18 #include <rmm/detail/export.hpp>
20 
21 #include <memory>
22 #include <utility>
23 
24 namespace RMM_NAMESPACE {
25 namespace mr {
26 namespace detail {
42 template <typename Resource, typename UpstreamTuple, std::size_t... Indices, typename... Args>
43 auto make_resource_impl(UpstreamTuple const& upstreams,
44  std::index_sequence<Indices...>,
45  Args&&... args)
46 {
47  return std::make_unique<Resource>(std::get<Indices>(upstreams).get()...,
48  std::forward<Args>(args)...);
49 }
50 
63 template <typename Resource, typename... Upstreams, typename... Args>
64 auto make_resource(std::tuple<std::shared_ptr<Upstreams>...> const& upstreams, Args&&... args)
65 {
66  return make_resource_impl<Resource>(
67  upstreams, std::index_sequence_for<Upstreams...>{}, std::forward<Args>(args)...);
68 }
69 } // namespace detail
70 
104 template <typename Resource, typename... Upstreams>
106  public:
108  std::tuple<std::shared_ptr<Upstreams>...>;
109 
143  template <typename... Args>
144  owning_wrapper(upstream_tuple upstreams, Args&&... args)
145  : upstreams_{std::move(upstreams)},
146  wrapped_{detail::make_resource<Resource>(upstreams_, std::forward<Args>(args)...)}
147  {
148  }
149 
153  [[nodiscard]] Resource const& wrapped() const noexcept { return *wrapped_; }
154 
158  [[nodiscard]] Resource& wrapped() noexcept { return *wrapped_; }
159 
160  private:
171  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
172  {
173  return wrapped().allocate(bytes, stream);
174  }
175 
185  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
186  {
187  wrapped().deallocate(ptr, bytes, stream);
188  }
189 
199  [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
200  {
201  if (this == &other) { return true; }
202  auto casted = dynamic_cast<owning_wrapper<Resource, Upstreams...> const*>(&other);
203  if (nullptr != casted) { return wrapped().is_equal(casted->wrapped()); }
204  return wrapped().is_equal(other);
205  }
206 
207  upstream_tuple upstreams_;
208  std::unique_ptr<Resource> wrapped_;
209 };
210 
243 template <template <typename...> class Resource, typename... Upstreams, typename... Args>
244 auto make_owning_wrapper(std::tuple<std::shared_ptr<Upstreams>...> upstreams, Args&&... args)
245 {
246  return std::make_shared<owning_wrapper<Resource<Upstreams...>, Upstreams...>>(
247  std::move(upstreams), std::forward<Args>(args)...);
248 }
249 
265 template <template <typename> class Resource, typename Upstream, typename... Args>
266 auto make_owning_wrapper(std::shared_ptr<Upstream> upstream, Args&&... args)
267 {
268  return make_owning_wrapper<Resource>(std::make_tuple(std::move(upstream)),
269  std::forward<Args>(args)...);
270 }
271  // end of group
273 } // namespace mr
274 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
Base class for all librmm device memory allocation.
Definition: device_memory_resource.hpp:93
Resource adaptor that maintains the lifetime of upstream resources.
Definition: owning_wrapper.hpp:105
Resource const & wrapped() const noexcept
A constant reference to the wrapped resource.
Definition: owning_wrapper.hpp:153
std::tuple< std::shared_ptr< Upstreams >... > upstream_tuple
Tuple of upstream memory resources.
Definition: owning_wrapper.hpp:108
Resource & wrapped() noexcept
A reference to the wrapped resource.
Definition: owning_wrapper.hpp:158
owning_wrapper(upstream_tuple upstreams, Args &&... args)
Constructs the wrapped resource using the provided upstreams and any additional arguments forwarded t...
Definition: owning_wrapper.hpp:144
auto make_owning_wrapper(std::shared_ptr< Upstream > upstream, Args &&... args)
Additional convenience factory for owning_wrapper when Resource has only a single upstream resource.
Definition: owning_wrapper.hpp:266
auto make_resource(std::tuple< std::shared_ptr< Upstreams >... > const &upstreams, Args &&... args)
Create a std::unique_ptr to a Resource with the given upstreams and arguments.
Definition: owning_wrapper.hpp:64
auto make_resource_impl(UpstreamTuple const &upstreams, std::index_sequence< Indices... >, Args &&... args)
Converts a tuple into a parameter pack.
Definition: owning_wrapper.hpp:43