limiting_resource_adaptor.hpp
1 /*
2  * Copyright (c) 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/detail/aligned.hpp>
19 #include <rmm/detail/error.hpp>
20 #include <rmm/mr/device/device_memory_resource.hpp>
21 
22 #include <cstddef>
23 
24 namespace rmm::mr {
37 template <typename Upstream>
39  public:
49  limiting_resource_adaptor(Upstream* upstream,
50  std::size_t allocation_limit,
51  std::size_t alignment = rmm::detail::CUDA_ALLOCATION_ALIGNMENT)
52  : allocation_limit_{allocation_limit},
53  allocated_bytes_(0),
54  alignment_(alignment),
55  upstream_{upstream}
56  {
57  RMM_EXPECTS(nullptr != upstream, "Unexpected null upstream resource pointer.");
58  }
59 
60  limiting_resource_adaptor() = delete;
61  ~limiting_resource_adaptor() override = default;
62  limiting_resource_adaptor(limiting_resource_adaptor const&) = delete;
63  limiting_resource_adaptor(limiting_resource_adaptor&&) noexcept = default;
64  limiting_resource_adaptor& operator=(limiting_resource_adaptor const&) = delete;
65  limiting_resource_adaptor& operator=(limiting_resource_adaptor&&) noexcept = default;
66 
72  [[nodiscard]] Upstream* get_upstream() const noexcept { return upstream_; }
73 
80  [[nodiscard]] bool supports_streams() const noexcept override
81  {
82  return upstream_->supports_streams();
83  }
84 
90  [[nodiscard]] bool supports_get_mem_info() const noexcept override
91  {
92  return upstream_->supports_get_mem_info();
93  }
94 
104  [[nodiscard]] std::size_t get_allocated_bytes() const { return allocated_bytes_; }
105 
113  [[nodiscard]] std::size_t get_allocation_limit() const { return allocation_limit_; }
114 
115  private:
129  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
130  {
131  auto const proposed_size = rmm::detail::align_up(bytes, alignment_);
132  auto const old = allocated_bytes_.fetch_add(proposed_size);
133  if (old + proposed_size <= allocation_limit_) {
134  try {
135  return upstream_->allocate(bytes, stream);
136  } catch (...) {
137  allocated_bytes_ -= proposed_size;
138  throw;
139  }
140  }
141 
142  allocated_bytes_ -= proposed_size;
143  RMM_FAIL("Exceeded memory limit", rmm::out_of_memory);
144  }
145 
155  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
156  {
157  std::size_t allocated_size = rmm::detail::align_up(bytes, alignment_);
158  upstream_->deallocate(ptr, bytes, stream);
159  allocated_bytes_ -= allocated_size;
160  }
161 
171  [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
172  {
173  if (this == &other) { return true; }
174  auto const* cast = dynamic_cast<limiting_resource_adaptor<Upstream> const*>(&other);
175  if (cast != nullptr) { return upstream_->is_equal(*cast->get_upstream()); }
176  return upstream_->is_equal(other);
177  }
178 
187  [[nodiscard]] std::pair<std::size_t, std::size_t> do_get_mem_info(
188  cuda_stream_view stream) const override
189  {
190  return {allocation_limit_ - allocated_bytes_, allocation_limit_};
191  }
192 
193  // maximum bytes this allocator is allowed to allocate.
194  std::size_t allocation_limit_;
195 
196  // number of currently-allocated bytes
197  std::atomic<std::size_t> allocated_bytes_;
198 
199  // todo: should be some way to ask the upstream...
200  std::size_t alignment_;
201 
202  Upstream* upstream_;
203 };
205 
214 template <typename Upstream>
215 limiting_resource_adaptor<Upstream> make_limiting_adaptor(Upstream* upstream,
216  std::size_t allocation_limit)
217 {
218  return limiting_resource_adaptor<Upstream>{upstream, allocation_limit};
219 }
220 
221 } // namespace rmm::mr
rmm::mr::limiting_resource_adaptor
Resource that uses Upstream to allocate memory and limits the total allocations possible.
Definition: limiting_resource_adaptor.hpp:38
rmm::mr::limiting_resource_adaptor::get_allocated_bytes
std::size_t get_allocated_bytes() const
Query the number of bytes that have been allocated. Note that this can not be used to know how large ...
Definition: limiting_resource_adaptor.hpp:104
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::limiting_resource_adaptor::get_upstream
Upstream * get_upstream() const noexcept
Return pointer to the upstream resource.
Definition: limiting_resource_adaptor.hpp:72
rmm::out_of_memory
Exception thrown when RMM runs out of memory.
Definition: error.hpp:68
rmm::mr::limiting_resource_adaptor::supports_get_mem_info
bool supports_get_mem_info() const noexcept override
Query whether the resource supports the get_mem_info API.
Definition: limiting_resource_adaptor.hpp:90
rmm::mr::limiting_resource_adaptor::supports_streams
bool supports_streams() const noexcept override
Checks whether the upstream resource supports streams.
Definition: limiting_resource_adaptor.hpp:80
rmm::mr::limiting_resource_adaptor::limiting_resource_adaptor
limiting_resource_adaptor(Upstream *upstream, std::size_t allocation_limit, std::size_t alignment=rmm::detail::CUDA_ALLOCATION_ALIGNMENT)
Construct a new limiting resource adaptor using upstream to satisfy allocation requests and limiting ...
Definition: limiting_resource_adaptor.hpp:49
rmm::mr::limiting_resource_adaptor::get_allocation_limit
std::size_t get_allocation_limit() const
Query the maximum number of bytes that this allocator is allowed to allocate. This is the limit on th...
Definition: limiting_resource_adaptor.hpp:113
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82