arena_memory_resource.hpp

前往此文件的文档。

/*
 * 版权所有 (c) 2020-2025, NVIDIA CORPORATION。
 *
 * 根据 Apache 许可证 2.0 版（“许可证”）获得许可；
 * 除非符合许可证的规定，否则您不得使用此文件。
 * 您可以获取许可证的副本，网址为
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，根据许可证分发的软件按“原样”分发，
 * 不附带任何明示或暗示的保证或条件。
 * 请参阅许可证以了解管理权限和限制的特定语言。
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <rmm/aligned.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/format.hpp>
#include <rmm/detail/logging_assert.hpp>
#include <rmm/logger.hpp>
#include <rmm/mr/device/detail/arena.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda_runtime_api.h>

#include <cstddef>
#include <map>
#include <shared_mutex>
#include <thread>

namespace RMM_NAMESPACE {
namespace mr {
template <typename Upstream>
class arena_memory_resource final : public device_memory_resource {
 public
 /// 构造一个 arena_memory_resource。
 explicit arena_memory_resource(device_async_resource_ref upstream_mr,
 std::optional<std::size_t> arena_size = std::nullopt,
 bool dump_log_on_failure = false)
 : global_arena_{upstream_mr, arena_size}, dump_log_on_failure_{dump_log_on_failure}
 {
 if (dump_log_on_failure_) {
 logger_ =
 std::make_shared<rapids_logger::logger>("arena_memory_dump", "rmm_arena_memory_dump.log");
 // Set the level to `debug` for more detailed output.
 logger_->set_level(rapids_logger::level_enum::info);
 }
 }
 /// 构造一个 arena_memory_resource。
 explicit arena_memory_resource(Upstream* upstream_mr,
 std::optional<std::size_t> arena_size = std::nullopt,
 bool dump_log_on_failure = false)
 : arena_memory_resource{
 to_device_async_resource_ref_checked(upstream_mr), arena_size, dump_log_on_failure}
 {
 }

 ~arena_memory_resource() override = default;

 // 禁用复制（和移动）语义。
 arena_memory_resource(arena_memory_resource const&) = delete;
 arena_memory_resource& operator=(arena_memory_resource const&) = delete;
 arena_memory_resource(arena_memory_resource&&) noexcept = delete;
 arena_memory_resource& operator=(arena_memory_resource&&) noexcept = delete;

 private
 using global_arena = rmm::mr::detail::arena::global_arena;
 using arena = rmm::mr::detail::arena::arena;
 /// 从资源分配大小为 `bytes` 的内存。
 void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
 {
 if (bytes <= 0) { return nullptr; }
#ifdef RMM_ARENA_USE_SIZE_CLASSES
 bytes = rmm::mr::detail::arena::align_to_size_class(bytes);
#else
 bytes = rmm::align_up(bytes, rmm::CUDA_ALLOCATION_ALIGNMENT);
#endif
 auto& arena = get_arena(stream);

 {
 std::shared_lock lock(mtx_);
 void* pointer = arena.allocate(bytes);
 if (pointer != nullptr) { return pointer; }
 }

 {
 std::unique_lock lock(mtx_);
 defragment();
 void* pointer = arena.allocate(bytes);
 if (pointer == nullptr) {
 if (dump_log_on_failure_) { dump_memory_log(bytes); }
 auto const msg = std::string("超出最大池大小（分配失败") +
 rmm::detail::format_bytes(bytes) + "): 竞技场中没有空间。";
 RMM_FAIL(msg.c_str(), rmm::out_of_memory);
 }
 return pointer;
 }
 }
 /// 碎片整理所有竞技场。
 void defragment()
 {
 RMM_CUDA_TRY(cudaDeviceSynchronize());
 for (auto& thread_arena : thread_arenas_) {
 thread_arena.second->clean();
 }
 for (auto& stream_arena : stream_arenas_) {
 stream_arena.second.clean();
 }
 }
 /// 释放指向 `ptr` 的大小为 `bytes` 的内存。
 void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
 {
 if (ptr == nullptr || bytes <= 0) { return; }
#ifdef RMM_ARENA_USE_SIZE_CLASSES
 bytes = rmm::mr::detail::arena::align_to_size_class(bytes);
#else
 bytes = rmm::align_up(bytes, rmm::CUDA_ALLOCATION_ALIGNMENT);
#endif
 auto& arena = get_arena(stream);

 {
 std::shared_lock lock(mtx_);
 // 如果要释放的内存不属于当前竞技场，则以下操作将返回 false。
 if (arena.deallocate(ptr, bytes, stream)) { return; }
 }

 {
 // 由于我们将此内存返回到另一个流，因此需要确保当前流
 // 已完成。
 stream.synchronize_no_throw();

 std::unique_lock lock(mtx_);
 deallocate_from_other_arena(ptr, bytes, stream);
 }
 }
 /// 从不属于调用线程/流的竞技场中释放内存。
 void deallocate_from_other_arena(void* ptr, std::size_t bytes, cuda_stream_view stream)
 {
 if (use_per_thread_arena(stream)) {
 for (auto const& thread_arena : thread_arenas_) {
 if (thread_arena.second->deallocate(ptr, bytes)) { return; }
 }
 } else {
 for (auto& stream_arena : stream_arenas_) {
 if (stream_arena.second.deallocate(ptr, bytes)) { return; }
 }
 }

 if (!global_arena_.deallocate(ptr, bytes)) {
 // 可以同时使用每个线程的默认流和另一个流池。
 // 这意味着分配有可能从线程或流竞技场移动
 // 在碎片整理期间回到全局竞技场，然后再移至另一个竞技场
 // 类型。例如，线程竞技场 -> 全局竞技场 -> 流竞技场。如果发生这种情况并且
 // 在它还是线程竞技场时曾有分配来自它，我们现在需要检查
 // 分配是否属于流竞技场，反之亦然。
 // 仅在特殊情况下执行此操作，以免影响性能并必须检查所有
 // 竞技场。
 if (use_per_thread_arena(stream)) {
 for (auto& stream_arena : stream_arenas_) {
 if (stream_arena.second.deallocate(ptr, bytes)) { return; }
 }
 } else {
 for (auto const& thread_arena : thread_arenas_) {
 if (thread_arena.second->deallocate(ptr, bytes)) { return; }
 }
 }
 RMM_FAIL("未找到分配");
 }
 }
 /// 返回指定流的竞技场。
 arena& get_arena(cuda_stream_view stream)
 {
 if (use_per_thread_arena(stream)) { return get_thread_arena(); }
 return get_stream_arena(stream);
 }
 /// 返回当前线程的竞技场。
 arena& get_thread_arena()
 {
 auto const thread_id = std::this_thread::get_id();
 {
 std::shared_lock lock(map_mtx_);
 auto const iter = thread_arenas_.find(thread_id);
 if (iter != thread_arenas_.end()) { return *iter->second; }
 }
 {
 std::unique_lock lock(map_mtx_);
 auto thread_arena = std::make_shared<arena>(global_arena_);
 thread_arenas_.emplace(thread_id, thread_arena);
 thread_local detail::arena::arena_cleaner cleaner{thread_arena};
 return *thread_arena;
 }
 }
 /// 返回指定流的竞技场。
 arena& get_stream_arena(cuda_stream_view stream)
 {
 RMM_LOGGING_ASSERT(!use_per_thread_arena(stream));
 {
 std::shared_lock lock(map_mtx_);
 auto const iter = stream_arenas_.find(stream.value());
 if (iter != stream_arenas_.end()) { return iter->second; }
 }
 {
 std::unique_lock lock(map_mtx_);
 stream_arenas_.emplace(stream.value(), global_arena_);
 return stream_arenas_.at(stream.value());
 }
 }
 /// 转储全局竞技场的内存日志。
 void dump_memory_log(size_t bytes)
 {
 logger_->info("**************************************************");
 logger_->info("尝试分配 %s 时内存不足。", rmm::detail::format_bytes(bytes));
 logger_->info("**************************************************");
 logger_->info("全局竞技场：");
 global_arena_.dump_memory_log(logger_);
 logger_->flush();
 }
 /// 如果流是每线程默认流，则返回 true。
 static bool use_per_thread_arena(cuda_stream_view stream)
 {
 return stream.is_per_thread_default();
 }
 /// 全局竞技场。
 /// 从线程 ID 到线程竞技场的映射。
 /// 从流到流竞技场的映射。
 /// 失败时转储日志。
 /// 用于内存转储的记录器。
 /// 用于映射访问的互斥锁。
 /// 用于竞技场访问的互斥锁。
};
 // 组结束
} // namespace mr
} // namespace RMM_NAMESPACE