arena_memory_resource.hpp

/*
 * Copyright (c) 2020-2022, 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
 *
 *     http://www.apache.org/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/detail/error.hpp>
#include <rmm/logger.hpp>
#include <rmm/mr/device/detail/arena.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
 
#include <cuda_runtime_api.h>
 
#include <spdlog/common.h>
#include <spdlog/fmt/ostr.h>
 
#include <cstddef>
#include <map>
#include <shared_mutex>
#include <thread>
 
namespace rmm::mr {
 
template <typename Upstream>
class arena_memory_resource final : public device_memory_resource {
 public:
  explicit arena_memory_resource(Upstream* 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_ = spdlog::basic_logger_mt("arena_memory_dump", "rmm_arena_memory_dump.log");
      // Set the level to `debug` for more detailed output.
      logger_->set_level(spdlog::level::info);
    }
  }
 
  ~arena_memory_resource() override = default;
 
  // Disable copy (and move) semantics.
  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;
 
  bool supports_streams() const noexcept override { return true; }
 
  bool supports_get_mem_info() const noexcept override { return false; }
 
 private:
  using global_arena = rmm::mr::detail::arena::global_arena<Upstream>;
  using arena        = rmm::mr::detail::arena::arena<Upstream>;
 
  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::detail::align_up(bytes, rmm::detail::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); }
        RMM_FAIL("Maximum pool size exceeded", 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();
    }
  }
 
  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::detail::align_up(bytes, rmm::detail::CUDA_ALLOCATION_ALIGNMENT);
#endif
    auto& arena = get_arena(stream);
 
    {
      std::shared_lock lock(mtx_);
      // If the memory being freed does not belong to the arena, the following will return false.
      if (arena.deallocate(ptr, bytes, stream)) { return; }
    }
 
    {
      // Since we are returning this memory to another stream, we need to make sure the current
      // stream is caught up.
      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)) { RMM_FAIL("allocation not found"); }
  }
 
  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<Upstream> 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());
    }
  }
 
  std::pair<std::size_t, std::size_t> do_get_mem_info(cuda_stream_view stream) const override
  {
    return std::make_pair(0, 0);
  }
 
  void dump_memory_log(size_t bytes)
  {
    logger_->info("**************************************************");
    logger_->info("Ran out of memory trying to allocate {}.", rmm::detail::bytes{bytes});
    logger_->info("**************************************************");
    logger_->info("Global arena:");
    global_arena_.dump_memory_log(logger_);
    logger_->flush();
  }
 
  static bool use_per_thread_arena(cuda_stream_view stream)
  {
    return stream.is_per_thread_default();
  }
 
  global_arena global_arena_;
  std::map<std::thread::id, std::shared_ptr<arena>> thread_arenas_;
  std::map<cudaStream_t, arena> stream_arenas_;
  bool dump_log_on_failure_{};
  std::shared_ptr<spdlog::logger> logger_{};
  mutable std::shared_mutex map_mtx_;
  mutable std::shared_mutex mtx_;
};
 
}  // namespace rmm::mr