pool_memory_resource.hpp

前往此文件的文档。

/*
 * Copyright (c) 2020-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <rmm/aligned.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/format.hpp>
#include <rmm/detail/logging_assert.hpp>
#include <rmm/detail/thrust_namespace.h>
#include <rmm/logger.hpp>
#include <rmm/mr/device/detail/coalescing_free_list.hpp>
#include <rmm/mr/device/detail/stream_ordered_memory_resource.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda/std/type_traits>
#include <cuda_runtime_api.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>

#include <algorithm>
#include <cstddef>
#include <mutex>
#include <optional>
#include <set>

namespace RMM_NAMESPACE {
namespace mr {
namespace detail {
template <class PoolResource, class Upstream, class Property, class = void>
struct maybe_remove_property {};

template <class PoolResource, class Upstream, class Property>
struct maybe_remove_property<PoolResource,
 Upstream,
 Property,
 cuda::std::enable_if_t<!cuda::has_property<Upstream, Property>>> {
#if defined(__GNUC__) && !defined(__clang__) // GCC warns about compatibility
 // issues with pre ISO C++ code
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wnon-template-friend"
#endif // __GNUC__ and not __clang__
 friend void get_property(const PoolResource&, Property) = delete;
#if defined(__GNUC__) && !defined(__clang__)
#pragma GCC diagnostic pop
#endif // __GNUC__ and not __clang__
};
} // namespace detail

template <typename Upstream>
class pool_memory_resource final
 : public detail:
 maybe_remove_property<pool_memory_resource<Upstream>, Upstream, cuda::mr::device_accessible>,
 public detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
 detail::coalescing_free_list>,
 public cuda::forward_property<pool_memory_resource<Upstream>, Upstream> {
 public
 friend class detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
 detail::coalescing_free_list>;

 explicit pool_memory_resource(device_async_resource_ref upstream_mr,
 std::size_t initial_pool_size,
 std::optional<std::size_t> maximum_pool_size = std::nullopt)
 : upstream_mr_{upstream_mr}
 {
 RMM_EXPECTS(rmm::is_aligned(initial_pool_size, rmm::CUDA_ALLOCATION_ALIGNMENT),
 "Error, Initial pool size required to be a multiple of 256 bytes");
 RMM_EXPECTS(rmm::is_aligned(maximum_pool_size.value_or(0), rmm::CUDA_ALLOCATION_ALIGNMENT),
 "Error, Maximum pool size required to be a multiple of 256 bytes");

 initialize_pool(initial_pool_size, maximum_pool_size);
 }

 explicit pool_memory_resource(Upstream* upstream_mr,
 std::size_t initial_pool_size,
 std::optional<std::size_t> maximum_pool_size = std::nullopt)
 : upstream_mr_{to_device_async_resource_ref_checked(upstream_mr)}
 {
 RMM_EXPECTS(rmm::is_aligned(initial_pool_size, rmm::CUDA_ALLOCATION_ALIGNMENT),
 "Error, Initial pool size required to be a multiple of 256 bytes");
 RMM_EXPECTS(rmm::is_aligned(maximum_pool_size.value_or(0), rmm::CUDA_ALLOCATION_ALIGNMENT),
 "Error, Maximum pool size required to be a multiple of 256 bytes");

 initialize_pool(initial_pool_size, maximum_pool_size);
 }

 template <typename Upstream2 = Upstream,
 cuda::std::enable_if_t<cuda::mr::async_resource<Upstream2>, int> = 0>
 explicit pool_memory_resource(Upstream2& upstream_mr,
 std::size_t initial_pool_size,
 std::optional<std::size_t> maximum_pool_size = std::nullopt)
 : pool_memory_resource(cuda::std::addressof(upstream_mr), initial_pool_size, maximum_pool_size)
 {
 }

 ~pool_memory_resource() override { release(); }

 pool_memory_resource() = delete;
 pool_memory_resource(pool_memory_resource const&) = delete;
 pool_memory_resource(pool_memory_resource&&) = delete;
 pool_memory_resource& operator=(pool_memory_resource const&) = delete;
 pool_memory_resource& operator=(pool_memory_resource&&) = delete;

 [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
 {
 return upstream_mr_;
 }

 [[nodiscard]] std::size_t pool_size() const noexcept { return current_pool_size_; }

 protected
 using free_list = detail::coalescing_free_list;
 using block_type = free_list::block_type;
 using typename detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
 detail::coalescing_free_list>::split_block;
 using lock_guard = std::lock_guard<std::mutex>

 [[nodiscard]] std::size_t get_maximum_allocation_size() const
 {
 return std::numeric_limits<std::size_t>::max();
 }

 block_type try_to_expand(std::size_t try_size, std::size_t min_size, cuda_stream_view stream)
 {
 auto report_error = [&](const char* reason) {
 RMM_LOG_ERROR("[A][Stream %s][Upstream %zuB][FAILURE maximum pool size exceeded: %s]",
 rmm::detail::format_stream(stream),
 min_size,
 reason);
 auto const msg = std::string("Maximum pool size exceeded (failed to allocate ") +
 rmm::detail::format_bytes(min_size) + std::string("): ") + reason;
 RMM_FAIL(msg.c_str(), rmm::out_of_memory);
 };

 while (try_size >= min_size) {
 try {
 auto block = block_from_upstream(try_size, stream);
 current_pool_size_ += block.size();
 return block;
 } catch (std::exception const& e) {
 if (try_size == min_size) { report_error(e.what()); }
 }
 try_size = std::max(min_size, try_size / 2);
 }

 auto const max_size = maximum_pool_size_.value_or(std::numeric_limits<std::size_t>::max());
 auto const msg = std::string("Not enough room to grow, current/max/try size = ") +
 rmm::detail::format_bytes(pool_size()) + ", " +
 rmm::detail::format_bytes(max_size) + ", " +
 rmm::detail::format_bytes(min_size);
 report_error(msg.c_str());
 return {};
 }

 void initialize_pool(std::size_t initial_size, std::optional<std::size_t> maximum_size)
 {
 current_pool_size_ = 0; // try_to_expand will set this if it succeeds
 maximum_pool_size_ = maximum_size;

 RMM_EXPECTS(
 initial_size <= maximum_pool_size_.value_or(std::numeric_limits<std::size_t>::max()),
 "Initial pool size exceeds the maximum pool size!");

 if (initial_size > 0) {
 auto const block = try_to_expand(initial_size, initial_size, cuda_stream_legacy);
 this->insert_block(block, cuda_stream_legacy);
 }
 }

 block_type expand_pool(std::size_t size, free_list& blocks, cuda_stream_view stream)
 {
 // Strategy: If maximum_pool_size_ is set, then grow geometrically, e.g. by halfway to the
 // limit each time. If it is not set, grow exponentially, e.g. by doubling the pool size each
 // time. Upon failure, attempt to back off exponentially, e.g. by half the attempted size,
 // until either success or the attempt is less than the requested size.

 return try_to_expand(size_to_grow(size), size, stream);
 }

 [[nodiscard]] std::size_t size_to_grow(std::size_t size) const
 {
 if (maximum_pool_size_.has_value()) {
 auto const unaligned_remaining = maximum_pool_size_.value() - pool_size();
 using rmm::align_up;
 auto const remaining = align_up(unaligned_remaining, rmm::CUDA_ALLOCATION_ALIGNMENT);
 auto const aligned_size = align_up(size, rmm::CUDA_ALLOCATION_ALIGNMENT);
 return (aligned_size <= remaining) ? std::max(aligned_size, remaining / 2) : 0;
 }
 return std::max(size, pool_size());
 };

 block_type block_from_upstream(std::size_t size, cuda_stream_view stream)
 {
 RMM_LOG_DEBUG("[A][Stream %s][Upstream %zuB]", rmm::detail::format_stream(stream), size);

 if (size == 0) { return {}; }

 void* ptr = get_upstream_resource().allocate_async(size, stream);
 return *upstream_blocks_.emplace(static_cast<char*>(ptr), size, true).first;
 }

 split_block allocate_from_block(block_type const& block, std::size_t size)
 {
 block_type const alloc{block.pointer(), size, block.is_head()};
#ifdef RMM_POOL_TRACK_ALLOCATIONS
 allocated_blocks_.insert(alloc);
#endif

 auto rest = (block.size() > size)
 // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
 ? block_type{block.pointer() + size, block.size() - size, false}
 : block_type{};
 return {alloc, rest};
 }

 block_type free_block(void* ptr, std::size_t size) noexcept
 {
#ifdef RMM_POOL_TRACK_ALLOCATIONS
 if (ptr == nullptr) return block_type{};
 auto const iter = allocated_blocks_.find(static_cast<char*>(ptr));
 RMM_LOGGING_ASSERT(iter != allocated_blocks_.end());

 auto block = *iter;
 RMM_LOGGING_ASSERT(block.size() == rmm::align_up(size, allocation_alignment));
 allocated_blocks_.erase(iter);

 return block;
#else
 auto const iter = upstream_blocks_.find(static_cast<char*>(ptr));
 return block_type{static_cast<char*>(ptr), size, (iter != upstream_blocks_.end())};
#endif
 }

 void release()
 {
 lock_guard lock(this->get_mutex());

 for (auto block : upstream_blocks_) {
 get_upstream_resource().deallocate(block.pointer(), block.size());
 }
 upstream_blocks_.clear();
#ifdef RMM_POOL_TRACK_ALLOCATIONS
 allocated_blocks_.clear();
#endif

 current_pool_size_ = 0;
 }

#ifdef RMM_DEBUG_PRINT
 void print()
 {
 lock_guard lock(this->get_mutex());

 auto const [free, total] = rmm::available_device_memory();
 std::cout << "GPU free memory: " << free << " total: " << total << "\n";

 std::cout << "upstream_blocks: " << upstream_blocks_.size() << "\n";
 std::size_t upstream_total{0};

 for (auto blocks : upstream_blocks_) {
 blocks.print();
 upstream_total += blocks.size();
 }
 std::cout << "total upstream: " << upstream_total << " B\n";

#ifdef RMM_POOL_TRACK_ALLOCATIONS
 std::cout << "allocated_blocks: " << allocated_blocks_.size() << "\n";
 for (auto block : allocated_blocks_)
 block.print();
#endif

 this->print_free_blocks();
 }
#endif

 std::pair<std::size_t, std::size_t> free_list_summary(free_list const& blocks)
 {
 std::size_t largest{};
 std::size_t total{};
 std::for_each(blocks.cbegin(), blocks.cend(), [&largest, &total](auto const& block) {
 total += block.size();
 largest = std::max(largest, block.size());
 });
 return {largest, total};
 }

 private
 // The "heap" to allocate the pool from
 device_async_resource_ref upstream_mr_;
 std::size_t current_pool_size_{};
 std::optional<std::size_t> maximum_pool_size_{};

#ifdef RMM_POOL_TRACK_ALLOCATIONS
 std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> allocated_blocks_;
#endif

 // blocks allocated from upstream
 std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> upstream_blocks_;
}; // namespace mr
 // end of group
} // namespace mr
} // namespace RMM_NAMESPACE