memory.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <core/core.hpp>
 
#include <chrono>
#include <utility>
 
namespace iganet {
 
template <std::size_t id = 0> class MemoryDebugger {
private:
  struct MemoryObject {
    std::string name_;
    int64_t bytes_;
 
    MemoryObject(std::string name, int64_t bytes)
        : name_(std::move(name)), bytes_(bytes) {}
  };
 
  std::map<std::chrono::high_resolution_clock::time_point, MemoryObject>
      objects_;
 
  int64_t counter_;
 
  int64_t bytes_;
 
  std::chrono::high_resolution_clock::time_point init_;
 
  [[nodiscard]] std::string convert_bytes(int64_t bytes) const {
    if (bytes < 1024ull)
      return std::to_string(bytes) + "b";
    else if (bytes < 1024ull * 1024ull)
      return std::to_string(bytes / static_cast<double>(1024)) + "kb";
    else if (bytes < 1024ull * 1024ull * 1024ull)
      return std::to_string(bytes / static_cast<double>(1024 * 1024)) + "mb";
    else if (bytes < 1024ull * 1024ull * 1024ull * 1024ull)
      return std::to_string(bytes / static_cast<double>(1024 * 1024 * 1024)) +
             "gb";
    else
      return std::to_string(
                 bytes / static_cast<double>(1024) / static_cast<double>(1024) /
                 static_cast<double>(1024) / static_cast<double>(1024)) +
             "tb";
  }
 
public:
  MemoryDebugger()
      : counter_(0), bytes_(0),
        init_(std::chrono::high_resolution_clock::now()) {}
 
  void clear() {
    counter_ = 0;
    bytes_ = 0;
    objects_.clear();
  }
 
  inline void pretty_print(std::ostream &os = Log(log::info)) const {
    using namespace std::literals;
 
    os << "Memory debugger (ID=" << std::to_string(id) << ")\n";
    for (const auto &obj : objects_)
      os << "[" << std::right << std::setw(10) << (obj.first - init_) / 1ns
         << "ns] " << std::right << std::setw(10) << obj.second.name_ << " "
         << std::right << std::setw(10) << convert_bytes(obj.second.bytes_)
         << "\n";
    os << "[     Total  ] " << std::right << std::setw(10) << counter_ << " "
       << std::right << std::setw(10) << convert_bytes(bytes_) << "\n";
  }
 
  template <typename T>
  void add(const std::string &name, [[maybe_unused]] const T &obj) {
    counter_++;
    bytes_ += sizeof(obj);
    objects_.insert(
        std::pair<std::chrono::high_resolution_clock::time_point, MemoryObject>(
            std::chrono::high_resolution_clock::now(),
            MemoryObject(name, sizeof(obj))));
  }
 
  void add(const std::string &name, const torch::Tensor &tensor) {
    counter_++;
    bytes_ += tensor.element_size() * tensor.numel();
    objects_.insert(
        std::pair<std::chrono::high_resolution_clock::time_point, MemoryObject>(
            std::chrono::high_resolution_clock::now(),
            MemoryObject(name, tensor.element_size() * tensor.numel())));
  }
 
  template <typename T, std::size_t N>
  void add(const std::string &name, const std::array<T, N> &array) {
    for (std::size_t i = 0; i < N; ++i)
      add(name + std::to_string(i), array[i]);
  }
};
 
template <std::size_t id>
inline std::ostream &operator<<(std::ostream &os,
                                const MemoryDebugger<id> &obj) {
  obj.pretty_print(os);
  return os;
}
 
static MemoryDebugger<std::numeric_limits<std::size_t>::max()>
    global_memory_debugger;
 
#define register_memory(obj) ::iganet::global_memory_debugger.add(#obj, obj)
 
} // namespace iganet