owning_wrapper.hpp
1 /*
2  * Copyright (c) 2020-2021, 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/mr/device/device_memory_resource.hpp>
19 
20 #include <functional>
21 #include <iostream>
22 #include <memory>
23 #include <utility>
24 
25 namespace rmm::mr {
26 namespace detail {
28 template <typename Resource, typename UpstreamTuple, std::size_t... Indices, typename... Args>
29 auto make_resource_impl(UpstreamTuple const& upstreams,
30  std::index_sequence<Indices...>,
31  Args&&... args)
32 {
33  return std::make_unique<Resource>(std::get<Indices>(upstreams).get()...,
34  std::forward<Args>(args)...);
35 }
36 
37 template <typename Resource, typename... Upstreams, typename... Args>
38 auto make_resource(std::tuple<std::shared_ptr<Upstreams>...> const& upstreams, Args&&... args)
39 {
40  return make_resource_impl<Resource>(
41  upstreams, std::index_sequence_for<Upstreams...>{}, std::forward<Args>(args)...);
42 }
43 } // namespace detail
44 
73 template <typename Resource, typename... Upstreams>
75  public:
76  using upstream_tuple = std::tuple<std::shared_ptr<Upstreams>...>;
77 
111  template <typename... Args>
112  owning_wrapper(upstream_tuple upstreams, Args&&... args)
113  : upstreams_{std::move(upstreams)},
114  wrapped_{detail::make_resource<Resource>(upstreams_, std::forward<Args>(args)...)}
115  {
116  }
117 
122  [[nodiscard]] Resource const& wrapped() const noexcept { return *wrapped_; }
123 
128  [[nodiscard]] Resource& wrapped() noexcept { return *wrapped_; }
129 
133  [[nodiscard]] bool supports_streams() const noexcept override
134  {
135  return wrapped().supports_streams();
136  }
137 
143  [[nodiscard]] bool supports_get_mem_info() const noexcept override
144  {
145  return wrapped().supports_get_mem_info();
146  }
147 
148  private:
159  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
160  {
161  return wrapped().allocate(bytes, stream);
162  }
163 
175  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
176  {
177  wrapped().deallocate(ptr, bytes, stream);
178  }
179 
191  [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
192  {
193  if (this == &other) { return true; }
194  auto casted = dynamic_cast<owning_wrapper<Resource, Upstreams...> const*>(&other);
195  if (nullptr != casted) { return wrapped().is_equal(casted->wrapped()); }
196  return wrapped().is_equal(other);
197  }
198 
207  [[nodiscard]] std::pair<std::size_t, std::size_t> do_get_mem_info(
208  cuda_stream_view stream) const override
209  {
210  return wrapped().get_mem_info(stream);
211  }
212 
213  upstream_tuple upstreams_;
214  std::unique_ptr<Resource> wrapped_;
215 };
216 
249 template <template <typename...> class Resource, typename... Upstreams, typename... Args>
250 auto make_owning_wrapper(std::tuple<std::shared_ptr<Upstreams>...> upstreams, Args&&... args)
251 {
252  return std::make_shared<owning_wrapper<Resource<Upstreams...>, Upstreams...>>(
253  std::move(upstreams), std::forward<Args>(args)...);
254 }
255 
271 template <template <typename> class Resource, typename Upstream, typename... Args>
272 auto make_owning_wrapper(std::shared_ptr<Upstream> upstream, Args&&... args)
273 {
274  return make_owning_wrapper<Resource>(std::make_tuple(std::move(upstream)),
275  std::forward<Args>(args)...);
276 }
277 
278 } // namespace rmm::mr
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::owning_wrapper::wrapped
Resource & wrapped() noexcept
Returns reference to the wrapped resource.
Definition: owning_wrapper.hpp:128
rmm::mr::owning_wrapper::owning_wrapper
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:112
rmm::mr::owning_wrapper
Resource adaptor that maintains the lifetime of upstream resources.
Definition: owning_wrapper.hpp:74
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82
rmm::mr::owning_wrapper::wrapped
Resource const & wrapped() const noexcept
Returns a constant reference to the wrapped resource.
Definition: owning_wrapper.hpp:122
rmm::mr::owning_wrapper::supports_get_mem_info
bool supports_get_mem_info() const noexcept override
Query whether the resource supports the get_mem_info API.
Definition: owning_wrapper.hpp:143
rmm::mr::owning_wrapper::supports_streams
bool supports_streams() const noexcept override
Query whether the resource supports use of non-null CUDA streams for allocation/deallocation.
Definition: owning_wrapper.hpp:133