tracking_resource_adaptor.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/detail/error.hpp>
19 #include <rmm/detail/stack_trace.hpp>
20 #include <rmm/mr/device/device_memory_resource.hpp>
21 
22 #include <cstddef>
23 #include <map>
24 #include <mutex>
25 #include <shared_mutex>
26 #include <sstream>
27 
28 namespace rmm::mr {
44 template <typename Upstream>
46  public:
47  // can be a std::shared_mutex once C++17 is adopted
48  using read_lock_t = std::shared_lock<std::shared_timed_mutex>;
49  using write_lock_t = std::unique_lock<std::shared_timed_mutex>;
50 
57  struct allocation_info {
58  std::unique_ptr<rmm::detail::stack_trace> strace;
59  std::size_t allocation_size;
60 
61  allocation_info() = delete;
62  allocation_info(std::size_t size, bool capture_stack)
63  : strace{[&]() {
64  return capture_stack ? std::make_unique<rmm::detail::stack_trace>() : nullptr;
65  }()},
66  allocation_size{size} {};
67  };
68 
78  tracking_resource_adaptor(Upstream* upstream, bool capture_stacks = false)
79  : capture_stacks_{capture_stacks}, allocated_bytes_{0}, upstream_{upstream}
80  {
81  RMM_EXPECTS(nullptr != upstream, "Unexpected null upstream resource pointer.");
82  }
83 
84  tracking_resource_adaptor() = delete;
85  ~tracking_resource_adaptor() override = default;
86  tracking_resource_adaptor(tracking_resource_adaptor const&) = delete;
87  tracking_resource_adaptor& operator=(tracking_resource_adaptor const&) = delete;
88  tracking_resource_adaptor(tracking_resource_adaptor&&) noexcept = default;
89  tracking_resource_adaptor& operator=(tracking_resource_adaptor&&) noexcept = default;
90 
96  Upstream* get_upstream() const noexcept { return upstream_; }
97 
104  bool supports_streams() const noexcept override { return upstream_->supports_streams(); }
105 
111  bool supports_get_mem_info() const noexcept override
112  {
113  return upstream_->supports_get_mem_info();
114  }
115 
123  std::map<void*, allocation_info> const& get_outstanding_allocations() const noexcept
124  {
125  return allocations_;
126  }
127 
137  std::size_t get_allocated_bytes() const noexcept { return allocated_bytes_; }
138 
150  {
151  read_lock_t lock(mtx_);
152 
153  std::ostringstream oss;
154 
155  if (!allocations_.empty()) {
156  for (auto const& alloc : allocations_) {
157  oss << alloc.first << ": " << alloc.second.allocation_size << " B";
158  if (alloc.second.strace != nullptr) {
159  oss << " : callstack:" << std::endl << *alloc.second.strace;
160  }
161  oss << std::endl;
162  }
163  }
164 
165  return oss.str();
166  }
167 
173  {
174 #if SPDLOG_ACTIVE_LEVEL <= SPDLOG_LEVEL_DEBUG
175  RMM_LOG_DEBUG("Outstanding Allocations: {}", get_outstanding_allocations_str());
176 #endif // SPDLOG_ACTIVE_LEVEL <= SPDLOG_LEVEL_DEBUG
177  }
178 
179  private:
193  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
194  {
195  void* ptr = upstream_->allocate(bytes, stream);
196 
197  // track it.
198  {
199  write_lock_t lock(mtx_);
200  allocations_.emplace(ptr, allocation_info{bytes, capture_stacks_});
201  }
202  allocated_bytes_ += bytes;
203 
204  return ptr;
205  }
206 
216  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
217  {
218  upstream_->deallocate(ptr, bytes, stream);
219  {
220  write_lock_t lock(mtx_);
221 
222  const auto found = allocations_.find(ptr);
223 
224  // Ensure the allocation is found and the number of bytes match
225  if (found == allocations_.end()) {
226  // Don't throw but log an error. Throwing in a descructor (or any noexcept) will call
227  // std::terminate
228  RMM_LOG_ERROR(
229  "Deallocating a pointer that was not tracked. Ptr: {:p} [{}B], Current Num. Allocations: "
230  "{}",
231  fmt::ptr(ptr),
232  bytes,
233  this->allocations_.size());
234  } else {
235  allocations_.erase(found);
236 
237  auto allocated_bytes = found->second.allocation_size;
238 
239  if (allocated_bytes != bytes) {
240  // Don't throw but log an error. Throwing in a descructor (or any noexcept) will call
241  // std::terminate
242  RMM_LOG_ERROR(
243  "Alloc bytes ({}) and Dealloc bytes ({}) do not match", allocated_bytes, bytes);
244 
245  bytes = allocated_bytes;
246  }
247  }
248  }
249  allocated_bytes_ -= bytes;
250  }
251 
261  bool do_is_equal(device_memory_resource const& other) const noexcept override
262  {
263  if (this == &other) { return true; }
264  auto cast = dynamic_cast<tracking_resource_adaptor<Upstream> const*>(&other);
265  return cast != nullptr ? upstream_->is_equal(*cast->get_upstream())
266  : upstream_->is_equal(other);
267  }
268 
277  std::pair<std::size_t, std::size_t> do_get_mem_info(cuda_stream_view stream) const override
278  {
279  return upstream_->get_mem_info(stream);
280  }
281 
282  bool capture_stacks_; // whether or not to capture call stacks
283  std::map<void*, allocation_info> allocations_; // map of active allocations
284  std::atomic<std::size_t> allocated_bytes_; // number of bytes currently allocated
285  std::shared_timed_mutex mutable mtx_; // mutex for thread safe access to allocations_
286  Upstream* upstream_; // the upstream resource used for satisfying allocation requests
287 };
288 
296 template <typename Upstream>
297 tracking_resource_adaptor<Upstream> make_tracking_adaptor(Upstream* upstream)
298 {
299  return tracking_resource_adaptor<Upstream>{upstream};
300 }
301 
302 } // namespace rmm::mr
rmm::mr::tracking_resource_adaptor::log_outstanding_allocations
void log_outstanding_allocations() const
Log any outstanding allocations via RMM_LOG_DEBUG.
Definition: tracking_resource_adaptor.hpp:172
rmm::mr::tracking_resource_adaptor::allocation_info
Information stored about an allocation. Includes the size and a stack trace if the tracking_resource_...
Definition: tracking_resource_adaptor.hpp:57
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::tracking_resource_adaptor::get_allocated_bytes
std::size_t get_allocated_bytes() const noexcept
Query the number of bytes that have been allocated. Note that this can not be used to know how large ...
Definition: tracking_resource_adaptor.hpp:137
rmm::mr::tracking_resource_adaptor::tracking_resource_adaptor
tracking_resource_adaptor(Upstream *upstream, bool capture_stacks=false)
Construct a new tracking resource adaptor using upstream to satisfy allocation requests.
Definition: tracking_resource_adaptor.hpp:78
rmm::mr::tracking_resource_adaptor::get_outstanding_allocations
std::map< void *, allocation_info > const & get_outstanding_allocations() const noexcept
Get the outstanding allocations map.
Definition: tracking_resource_adaptor.hpp:123
rmm::mr::tracking_resource_adaptor::supports_streams
bool supports_streams() const noexcept override
Checks whether the upstream resource supports streams.
Definition: tracking_resource_adaptor.hpp:104
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82
rmm::mr::tracking_resource_adaptor::get_upstream
Upstream * get_upstream() const noexcept
Return pointer to the upstream resource.
Definition: tracking_resource_adaptor.hpp:96
rmm::mr::tracking_resource_adaptor
Resource that uses Upstream to allocate memory and tracks allocations.
Definition: tracking_resource_adaptor.hpp:45
rmm::mr::tracking_resource_adaptor::get_outstanding_allocations_str
std::string get_outstanding_allocations_str() const
Gets a string containing the outstanding allocation pointers, their size, and optionally the stack tr...
Definition: tracking_resource_adaptor.hpp:149
rmm::mr::tracking_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: tracking_resource_adaptor.hpp:111