All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
pool_memory_resource.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/aligned.hpp>
19 #include <rmm/cuda_stream_view.hpp>
20 #include <rmm/detail/error.hpp>
21 #include <rmm/detail/export.hpp>
22 #include <rmm/detail/format.hpp>
23 #include <rmm/detail/logging_assert.hpp>
24 #include <rmm/detail/thrust_namespace.h>
25 #include <rmm/logger.hpp>
26 #include <rmm/mr/device/detail/coalescing_free_list.hpp>
27 #include <rmm/mr/device/detail/stream_ordered_memory_resource.hpp>
30 #include <rmm/resource_ref.hpp>
31 
32 #include <cuda/std/type_traits>
33 #include <cuda_runtime_api.h>
34 #include <thrust/iterator/counting_iterator.h>
35 #include <thrust/iterator/transform_iterator.h>
36 
37 #include <algorithm>
38 #include <cstddef>
39 #include <mutex>
40 #include <optional>
41 #include <set>
42 
43 namespace RMM_NAMESPACE {
44 namespace mr {
50 namespace detail {
62 template <class PoolResource, class Upstream, class Property, class = void>
64 
68 template <class PoolResource, class Upstream, class Property>
69 struct maybe_remove_property<PoolResource,
70  Upstream,
71  Property,
72  cuda::std::enable_if_t<!cuda::has_property<Upstream, Property>>> {
73 #if defined(__GNUC__) && !defined(__clang__) // GCC warns about compatibility
74  // issues with pre ISO C++ code
75 #pragma GCC diagnostic push
76 #pragma GCC diagnostic ignored "-Wnon-template-friend"
77 #endif // __GNUC__ and not __clang__
82  friend void get_property(const PoolResource&, Property) = delete;
83 #if defined(__GNUC__) && !defined(__clang__)
84 #pragma GCC diagnostic pop
85 #endif // __GNUC__ and not __clang__
86 };
87 } // namespace detail
88 
99 template <typename Upstream>
101  : public detail::
102  maybe_remove_property<pool_memory_resource<Upstream>, Upstream, cuda::mr::device_accessible>,
103  public detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
104  detail::coalescing_free_list>,
105  public cuda::forward_property<pool_memory_resource<Upstream>, Upstream> {
106  public:
107  friend class detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
108  detail::coalescing_free_list>;
109 
124  explicit pool_memory_resource(device_async_resource_ref upstream_mr,
125  std::size_t initial_pool_size,
126  std::optional<std::size_t> maximum_pool_size = std::nullopt)
127  : upstream_mr_{upstream_mr}
128  {
129  RMM_EXPECTS(rmm::is_aligned(initial_pool_size, rmm::CUDA_ALLOCATION_ALIGNMENT),
130  "Error, Initial pool size required to be a multiple of 256 bytes");
131  RMM_EXPECTS(rmm::is_aligned(maximum_pool_size.value_or(0), rmm::CUDA_ALLOCATION_ALIGNMENT),
132  "Error, Maximum pool size required to be a multiple of 256 bytes");
133 
134  initialize_pool(initial_pool_size, maximum_pool_size);
135  }
136 
152  explicit pool_memory_resource(Upstream* upstream_mr,
153  std::size_t initial_pool_size,
154  std::optional<std::size_t> maximum_pool_size = std::nullopt)
155  : upstream_mr_{to_device_async_resource_ref_checked(upstream_mr)}
156  {
157  RMM_EXPECTS(rmm::is_aligned(initial_pool_size, rmm::CUDA_ALLOCATION_ALIGNMENT),
158  "Error, Initial pool size required to be a multiple of 256 bytes");
159  RMM_EXPECTS(rmm::is_aligned(maximum_pool_size.value_or(0), rmm::CUDA_ALLOCATION_ALIGNMENT),
160  "Error, Maximum pool size required to be a multiple of 256 bytes");
161 
162  initialize_pool(initial_pool_size, maximum_pool_size);
163  }
164 
180  template <typename Upstream2 = Upstream,
181  cuda::std::enable_if_t<cuda::mr::async_resource<Upstream2>, int> = 0>
182  explicit pool_memory_resource(Upstream2& upstream_mr,
183  std::size_t initial_pool_size,
184  std::optional<std::size_t> maximum_pool_size = std::nullopt)
185  : pool_memory_resource(cuda::std::addressof(upstream_mr), initial_pool_size, maximum_pool_size)
186  {
187  }
188 
193  ~pool_memory_resource() override { release(); }
194 
195  pool_memory_resource() = delete;
198  pool_memory_resource& operator=(pool_memory_resource const&) = delete;
199  pool_memory_resource& operator=(pool_memory_resource&&) = delete;
200 
204  [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
205  {
206  return upstream_mr_;
207  }
208 
216  [[nodiscard]] std::size_t pool_size() const noexcept { return current_pool_size_; }
217 
218  protected:
219  using free_list = detail::coalescing_free_list;
220  using block_type = free_list::block_type;
221  using typename detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
222  detail::coalescing_free_list>::split_block;
223  using lock_guard = std::lock_guard<std::mutex>;
224 
233  [[nodiscard]] std::size_t get_maximum_allocation_size() const
234  {
235  return std::numeric_limits<std::size_t>::max();
236  }
237 
253  block_type try_to_expand(std::size_t try_size, std::size_t min_size, cuda_stream_view stream)
254  {
255  auto report_error = [&](const char* reason) {
256  RMM_LOG_ERROR("[A][Stream %s][Upstream %zuB][FAILURE maximum pool size exceeded: %s]",
257  rmm::detail::format_stream(stream),
258  min_size,
259  reason);
260  auto const msg = std::string("Maximum pool size exceeded (failed to allocate ") +
261  rmm::detail::format_bytes(min_size) + std::string("): ") + reason;
262  RMM_FAIL(msg.c_str(), rmm::out_of_memory);
263  };
264 
265  while (try_size >= min_size) {
266  try {
267  auto block = block_from_upstream(try_size, stream);
268  current_pool_size_ += block.size();
269  return block;
270  } catch (std::exception const& e) {
271  if (try_size == min_size) { report_error(e.what()); }
272  }
273  try_size = std::max(min_size, try_size / 2);
274  }
275 
276  auto const max_size = maximum_pool_size_.value_or(std::numeric_limits<std::size_t>::max());
277  auto const msg = std::string("Not enough room to grow, current/max/try size = ") +
278  rmm::detail::format_bytes(pool_size()) + ", " +
279  rmm::detail::format_bytes(max_size) + ", " +
280  rmm::detail::format_bytes(min_size);
281  report_error(msg.c_str());
282  return {};
283  }
284 
293  void initialize_pool(std::size_t initial_size, std::optional<std::size_t> maximum_size)
294  {
295  current_pool_size_ = 0; // try_to_expand will set this if it succeeds
296  maximum_pool_size_ = maximum_size;
297 
298  RMM_EXPECTS(
299  initial_size <= maximum_pool_size_.value_or(std::numeric_limits<std::size_t>::max()),
300  "Initial pool size exceeds the maximum pool size!");
301 
302  if (initial_size > 0) {
303  auto const block = try_to_expand(initial_size, initial_size, cuda_stream_legacy);
304  this->insert_block(block, cuda_stream_legacy);
305  }
306  }
307 
317  block_type expand_pool(std::size_t size, free_list& blocks, cuda_stream_view stream)
318  {
319  // Strategy: If maximum_pool_size_ is set, then grow geometrically, e.g. by halfway to the
320  // limit each time. If it is not set, grow exponentially, e.g. by doubling the pool size each
321  // time. Upon failure, attempt to back off exponentially, e.g. by half the attempted size,
322  // until either success or the attempt is less than the requested size.
323 
324  return try_to_expand(size_to_grow(size), size, stream);
325  }
326 
339  [[nodiscard]] std::size_t size_to_grow(std::size_t size) const
340  {
341  if (maximum_pool_size_.has_value()) {
342  auto const unaligned_remaining = maximum_pool_size_.value() - pool_size();
343  using rmm::align_up;
344  auto const remaining = align_up(unaligned_remaining, rmm::CUDA_ALLOCATION_ALIGNMENT);
345  auto const aligned_size = align_up(size, rmm::CUDA_ALLOCATION_ALIGNMENT);
346  return (aligned_size <= remaining) ? std::max(aligned_size, remaining / 2) : 0;
347  }
348  return std::max(size, pool_size());
349  };
350 
360  {
361  RMM_LOG_DEBUG("[A][Stream %s][Upstream %zuB]", rmm::detail::format_stream(stream), size);
362 
363  if (size == 0) { return {}; }
364 
365  void* ptr = get_upstream_resource().allocate_async(size, stream);
366  return *upstream_blocks_.emplace(static_cast<char*>(ptr), size, true).first;
367  }
368 
379  split_block allocate_from_block(block_type const& block, std::size_t size)
380  {
381  block_type const alloc{block.pointer(), size, block.is_head()};
382 #ifdef RMM_POOL_TRACK_ALLOCATIONS
383  allocated_blocks_.insert(alloc);
384 #endif
385 
386  auto rest = (block.size() > size)
387  // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
388  ? block_type{block.pointer() + size, block.size() - size, false}
389  : block_type{};
390  return {alloc, rest};
391  }
392 
401  block_type free_block(void* ptr, std::size_t size) noexcept
402  {
403 #ifdef RMM_POOL_TRACK_ALLOCATIONS
404  if (ptr == nullptr) return block_type{};
405  auto const iter = allocated_blocks_.find(static_cast<char*>(ptr));
406  RMM_LOGGING_ASSERT(iter != allocated_blocks_.end());
407 
408  auto block = *iter;
409  RMM_LOGGING_ASSERT(block.size() == rmm::align_up(size, allocation_alignment));
410  allocated_blocks_.erase(iter);
411 
412  return block;
413 #else
414  auto const iter = upstream_blocks_.find(static_cast<char*>(ptr));
415  return block_type{static_cast<char*>(ptr), size, (iter != upstream_blocks_.end())};
416 #endif
417  }
418 
423  void release()
424  {
425  lock_guard lock(this->get_mutex());
426 
427  for (auto block : upstream_blocks_) {
428  get_upstream_resource().deallocate(block.pointer(), block.size());
429  }
430  upstream_blocks_.clear();
431 #ifdef RMM_POOL_TRACK_ALLOCATIONS
432  allocated_blocks_.clear();
433 #endif
434 
435  current_pool_size_ = 0;
436  }
437 
438 #ifdef RMM_DEBUG_PRINT
445  void print()
446  {
447  lock_guard lock(this->get_mutex());
448 
449  auto const [free, total] = rmm::available_device_memory();
450  std::cout << "GPU free memory: " << free << " total: " << total << "\n";
451 
452  std::cout << "upstream_blocks: " << upstream_blocks_.size() << "\n";
453  std::size_t upstream_total{0};
454 
455  for (auto blocks : upstream_blocks_) {
456  blocks.print();
457  upstream_total += blocks.size();
458  }
459  std::cout << "total upstream: " << upstream_total << " B\n";
460 
461 #ifdef RMM_POOL_TRACK_ALLOCATIONS
462  std::cout << "allocated_blocks: " << allocated_blocks_.size() << "\n";
463  for (auto block : allocated_blocks_)
464  block.print();
465 #endif
466 
467  this->print_free_blocks();
468  }
469 #endif
470 
479  std::pair<std::size_t, std::size_t> free_list_summary(free_list const& blocks)
480  {
481  std::size_t largest{};
482  std::size_t total{};
483  std::for_each(blocks.cbegin(), blocks.cend(), [&largest, &total](auto const& block) {
484  total += block.size();
485  largest = std::max(largest, block.size());
486  });
487  return {largest, total};
488  }
489 
490  private:
491  // The "heap" to allocate the pool from
492  device_async_resource_ref upstream_mr_;
493  std::size_t current_pool_size_{};
494  std::optional<std::size_t> maximum_pool_size_{};
495 
496 #ifdef RMM_POOL_TRACK_ALLOCATIONS
497  std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> allocated_blocks_;
498 #endif
499 
500  // blocks allocated from upstream
501  std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> upstream_blocks_;
502 }; // namespace mr
503  // end of group
505 } // namespace mr
506 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
A coalescing best-fit suballocator which uses a pool of memory allocated from an upstream memory_reso...
Definition: pool_memory_resource.hpp:105
block_type free_block(void *ptr, std::size_t size) noexcept
Finds, frees and returns the block associated with pointer ptr.
Definition: pool_memory_resource.hpp:401
void initialize_pool(std::size_t initial_size, std::optional< std::size_t > maximum_size)
Allocate initial memory for the pool.
Definition: pool_memory_resource.hpp:293
split_block allocate_from_block(block_type const &block, std::size_t size)
Splits block if necessary to return a pointer to memory of size bytes.
Definition: pool_memory_resource.hpp:379
device_async_resource_ref get_upstream_resource() const noexcept
rmm::device_async_resource_ref to the upstream resource
Definition: pool_memory_resource.hpp:204
std::size_t size_to_grow(std::size_t size) const
Given a minimum size, computes an appropriate size to grow the pool.
Definition: pool_memory_resource.hpp:339
free_list::block_type block_type
The type of block returned by the free list.
Definition: pool_memory_resource.hpp:220
std::pair< std::size_t, std::size_t > free_list_summary(free_list const &blocks)
Get the largest available block size and total free size in the specified free list.
Definition: pool_memory_resource.hpp:479
std::size_t get_maximum_allocation_size() const
Get the maximum size of allocations supported by this memory resource.
Definition: pool_memory_resource.hpp:233
block_type expand_pool(std::size_t size, free_list &blocks, cuda_stream_view stream)
Allocate space from upstream to supply the suballocation pool and return a sufficiently sized block.
Definition: pool_memory_resource.hpp:317
pool_memory_resource(Upstream2 &upstream_mr, std::size_t initial_pool_size, std::optional< std::size_t > maximum_pool_size=std::nullopt)
Construct a pool_memory_resource and allocate the initial device memory pool using upstream_mr.
Definition: pool_memory_resource.hpp:182
void release()
Free all memory allocated from the upstream memory_resource.
Definition: pool_memory_resource.hpp:423
block_type try_to_expand(std::size_t try_size, std::size_t min_size, cuda_stream_view stream)
Try to expand the pool by allocating a block of at least min_size bytes from upstream.
Definition: pool_memory_resource.hpp:253
std::lock_guard< std::mutex > lock_guard
Type of lock used to synchronize access.
Definition: pool_memory_resource.hpp:223
std::size_t pool_size() const noexcept
Computes the size of the current pool.
Definition: pool_memory_resource.hpp:216
~pool_memory_resource() override
Destroy the pool_memory_resource and deallocate all memory it allocated using the upstream resource.
Definition: pool_memory_resource.hpp:193
detail::coalescing_free_list free_list
The free list implementation.
Definition: pool_memory_resource.hpp:219
pool_memory_resource(Upstream *upstream_mr, std::size_t initial_pool_size, std::optional< std::size_t > maximum_pool_size=std::nullopt)
Construct a pool_memory_resource and allocate the initial device memory pool using upstream_mr.
Definition: pool_memory_resource.hpp:152
block_type block_from_upstream(std::size_t size, cuda_stream_view stream)
Allocate a block from upstream to expand the suballocation pool.
Definition: pool_memory_resource.hpp:359
Exception thrown when RMM runs out of memory.
Definition: error.hpp:87
std::pair< std::size_t, std::size_t > available_device_memory()
Returns the available and total device memory in bytes for the current device.
Definition: cuda_device.hpp:123
static const cuda_stream_view cuda_stream_legacy
Static cuda_stream_view of cudaStreamLegacy, for convenience.
Definition: cuda_stream_view.hpp:131
cuda::mr::async_resource_ref< cuda::mr::device_accessible > device_async_resource_ref
Alias for a cuda::mr::async_resource_ref with the property cuda::mr::device_accessible.
Definition: resource_ref.hpp:41
device_async_resource_ref to_device_async_resource_ref_checked(Resource *res)
Convert pointer to memory resource into device_async_resource_ref, checking for nullptr
Definition: resource_ref.hpp:79
static constexpr std::size_t CUDA_ALLOCATION_ALIGNMENT
Default alignment used for CUDA memory allocation.
Definition: aligned.hpp:43
constexpr bool is_aligned(std::size_t value, std::size_t alignment) noexcept
Checks whether a value is aligned to a multiple of a specified power of 2.
Definition: aligned.hpp:105
constexpr std::size_t align_up(std::size_t value, std::size_t alignment) noexcept
Align up to nearest multiple of specified power of 2.
Definition: aligned.hpp:77
Management of per-device device_memory_resources.
A helper class to remove the device_accessible property.
Definition: pool_memory_resource.hpp:63