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>
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>
43 namespace RMM_NAMESPACE {
62 template <
class PoolResource,
class Upstream,
class Property,
class =
void>
68 template <
class PoolResource,
class Upstream,
class Property>
72 cuda::std::enable_if_t<!cuda::has_property<Upstream, Property>>> {
73 #if defined(__GNUC__) && !defined(__clang__)
75 #pragma GCC diagnostic push
76 #pragma GCC diagnostic ignored "-Wnon-template-friend"
82 friend void get_property(const PoolResource&, Property) = delete;
83 #if defined(__GNUC__) && !defined(__clang__)
84 #pragma GCC diagnostic pop
99 template <
typename Upstream>
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> {
108 detail::coalescing_free_list>;
125 std::size_t initial_pool_size,
126 std::optional<std::size_t> maximum_pool_size = std::nullopt)
127 : upstream_mr_{upstream_mr}
130 "Error, Initial pool size required to be a multiple of 256 bytes");
132 "Error, Maximum pool size required to be a multiple of 256 bytes");
134 initialize_pool(initial_pool_size, maximum_pool_size);
153 std::size_t initial_pool_size,
154 std::optional<std::size_t> maximum_pool_size = std::nullopt)
158 "Error, Initial pool size required to be a multiple of 256 bytes");
160 "Error, Maximum pool size required to be a multiple of 256 bytes");
162 initialize_pool(initial_pool_size, maximum_pool_size);
180 template <
typename Upstream2 = Upstream,
181 cuda::std::enable_if_t<cuda::mr::async_resource<Upstream2>,
int> = 0>
183 std::size_t initial_pool_size,
184 std::optional<std::size_t> maximum_pool_size = std::nullopt)
216 [[nodiscard]] std::size_t
pool_size() const noexcept {
return current_pool_size_; }
221 using typename detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
222 detail::coalescing_free_list>::split_block;
235 return std::numeric_limits<std::size_t>::max();
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),
260 auto const msg = std::string(
"Maximum pool size exceeded (failed to allocate ") +
261 rmm::detail::format_bytes(min_size) + std::string(
"): ") + reason;
265 while (try_size >= min_size) {
267 auto block = block_from_upstream(try_size, stream);
268 current_pool_size_ += block.size();
270 }
catch (std::exception
const& e) {
271 if (try_size == min_size) { report_error(e.what()); }
273 try_size = std::max(min_size, try_size / 2);
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());
293 void initialize_pool(std::size_t initial_size, std::optional<std::size_t> maximum_size)
295 current_pool_size_ = 0;
296 maximum_pool_size_ = maximum_size;
299 initial_size <= maximum_pool_size_.value_or(std::numeric_limits<std::size_t>::max()),
300 "Initial pool size exceeds the maximum pool size!");
302 if (initial_size > 0) {
324 return try_to_expand(size_to_grow(size), size, stream);
341 if (maximum_pool_size_.has_value()) {
342 auto const unaligned_remaining = maximum_pool_size_.value() - pool_size();
346 return (aligned_size <= remaining) ? std::max(aligned_size, remaining / 2) : 0;
348 return std::max(size, pool_size());
361 RMM_LOG_DEBUG(
"[A][Stream %s][Upstream %zuB]", rmm::detail::format_stream(stream), size);
363 if (size == 0) {
return {}; }
365 void* ptr = get_upstream_resource().allocate_async(size, stream);
366 return *upstream_blocks_.emplace(
static_cast<char*
>(ptr), size,
true).first;
381 block_type const alloc{block.pointer(), size, block.is_head()};
382 #ifdef RMM_POOL_TRACK_ALLOCATIONS
383 allocated_blocks_.insert(alloc);
386 auto rest = (block.size() > size)
388 ?
block_type{block.pointer() + size, block.size() - size,
false}
390 return {alloc, rest};
403 #ifdef RMM_POOL_TRACK_ALLOCATIONS
405 auto const iter = allocated_blocks_.find(
static_cast<char*
>(ptr));
406 RMM_LOGGING_ASSERT(iter != allocated_blocks_.end());
409 RMM_LOGGING_ASSERT(block.size() ==
rmm::align_up(size, allocation_alignment));
410 allocated_blocks_.erase(iter);
414 auto const iter = upstream_blocks_.find(
static_cast<char*
>(ptr));
415 return block_type{
static_cast<char*
>(ptr), size, (iter != upstream_blocks_.end())};
427 for (
auto block : upstream_blocks_) {
428 get_upstream_resource().deallocate(block.pointer(), block.size());
430 upstream_blocks_.clear();
431 #ifdef RMM_POOL_TRACK_ALLOCATIONS
432 allocated_blocks_.clear();
435 current_pool_size_ = 0;
438 #ifdef RMM_DEBUG_PRINT
447 lock_guard lock(this->get_mutex());
450 std::cout <<
"GPU free memory: " << free <<
" total: " << total <<
"\n";
452 std::cout <<
"upstream_blocks: " << upstream_blocks_.size() <<
"\n";
453 std::size_t upstream_total{0};
455 for (
auto blocks : upstream_blocks_) {
457 upstream_total += blocks.size();
459 std::cout <<
"total upstream: " << upstream_total <<
" B\n";
461 #ifdef RMM_POOL_TRACK_ALLOCATIONS
462 std::cout <<
"allocated_blocks: " << allocated_blocks_.size() <<
"\n";
463 for (
auto block : allocated_blocks_)
467 this->print_free_blocks();
481 std::size_t largest{};
483 std::for_each(blocks.cbegin(), blocks.cend(), [&largest, &total](
auto const& block) {
484 total += block.size();
485 largest = std::max(largest, block.size());
487 return {largest, total};
493 std::size_t current_pool_size_{};
494 std::optional<std::size_t> maximum_pool_size_{};
496 #ifdef RMM_POOL_TRACK_ALLOCATIONS
497 std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> allocated_blocks_;
501 std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> upstream_blocks_;
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