memory_usage.cc

Go to the documentation of this file.

/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */
 
#include <algorithm>
#include <atomic>
#include <cassert>
#include <iostream>
#include <memory>
#include <mutex>
#include <vector>
 
#include "MEM_guardedalloc.h"
#include "mallocn_intern.hh"
 
#include "../../source/blender/blenlib/BLI_strict_flags.h"
 
namespace {
 
struct Local;
struct Global;
 
struct alignas(128) Local {
  std::shared_ptr<Global> global;
 
  bool destructed = false;
  bool is_main = false;
  std::atomic<int64_t> mem_in_use = 0;
  std::atomic<int64_t> blocks_num = 0;
  std::atomic<int64_t> mem_in_use_during_peak_update = 0;
 
  Local();
  ~Local();
};
 
struct Global {
  std::mutex locals_mutex;
  std::vector<Local *> locals;
  std::atomic<int64_t> mem_in_use_outside_locals = 0;
  std::atomic<int64_t> blocks_num_outside_locals = 0;
  std::atomic<size_t> peak = 0;
};
 
}  // namespace
 
static std::atomic<bool> use_local_counters = true;
static constexpr int64_t peak_update_threshold = 1024 * 1024;
 
static std::shared_ptr<Global> &get_global_ptr()
{
  static std::shared_ptr<Global> global = std::make_shared<Global>();
  return global;
}
 
static Global &get_global()
{
  return *get_global_ptr();
}
 
static Local &get_local_data()
{
  static thread_local Local local;
  assert(!local.destructed);
  return local;
}
 
Local::Local()
{
  this->global = get_global_ptr();
 
  std::lock_guard lock{this->global->locals_mutex};
  if (this->global->locals.empty()) {
    /* This is the first thread creating #Local, it is therefore the main thread because it's
     * created through #memory_usage_init. */
    this->is_main = true;
  }
  /* Register self in the global list. */
  this->global->locals.push_back(this);
}
 
Local::~Local()
{
  std::lock_guard lock{this->global->locals_mutex};
 
  /* Unregister self from the global list. */
  this->global->locals.erase(
      std::find(this->global->locals.begin(), this->global->locals.end(), this));
  /* Don't forget the memory counts stored locally. */
  this->global->blocks_num_outside_locals.fetch_add(this->blocks_num, std::memory_order_relaxed);
  this->global->mem_in_use_outside_locals.fetch_add(this->mem_in_use, std::memory_order_relaxed);
 
  if (this->is_main) {
    /* The main thread started shutting down. Use global counters from now on to avoid accessing
     * thread-locals after they have been destructed. */
    use_local_counters.store(false, std::memory_order_relaxed);
  }
  /* Helps to detect when thread locals are accidentally accessed after destruction. */
  this->destructed = true;
}
 
static void update_global_peak()
{
  Global &global = get_global();
  /* Update peak. */
  global.peak = std::max<size_t>(global.peak, memory_usage_current());
 
  std::lock_guard lock{global.locals_mutex};
 
  for (Local *local : global.locals) {
    assert(!local->destructed);
    /* Updating this makes sure that the peak is not updated too often, which would degrade
     * performance. */
    local->mem_in_use_during_peak_update = local->mem_in_use.load(std::memory_order_relaxed);
  }
}
 
void memory_usage_init()
{
  /* Makes sure that the static and thread-local variables on the main thread are initialized. */
  get_local_data();
}
 
void memory_usage_block_alloc(const size_t size)
{
  if (LIKELY(use_local_counters.load(std::memory_order_relaxed))) {
    Local &local = get_local_data();
    /* Increase local memory counts. This does not cause thread synchronization in the majority of
     * cases, because each thread has these counters on a separate cache line. It may only cause
     * synchronization if another thread is computing the total current memory usage at the same
     * time, which is very rare compared to doing allocations. */
    local.blocks_num.fetch_add(1, std::memory_order_relaxed);
    local.mem_in_use.fetch_add(int64_t(size), std::memory_order_relaxed);
 
    /* If a certain amount of new memory has been allocated, update the peak. */
    if (local.mem_in_use - local.mem_in_use_during_peak_update > peak_update_threshold) {
      update_global_peak();
    }
  }
  else {
    Global &global = get_global();
    /* Increase global memory counts. */
    global.blocks_num_outside_locals.fetch_add(1, std::memory_order_relaxed);
    global.mem_in_use_outside_locals.fetch_add(int64_t(size), std::memory_order_relaxed);
  }
}
 
void memory_usage_block_free(const size_t size)
{
  if (LIKELY(use_local_counters)) {
    /* Decrease local memory counts. See comment in #memory_usage_block_alloc for details regarding
     * thread synchronization. */
    Local &local = get_local_data();
    local.mem_in_use.fetch_sub(int64_t(size), std::memory_order_relaxed);
    local.blocks_num.fetch_sub(1, std::memory_order_relaxed);
  }
  else {
    Global &global = get_global();
    /* Decrease global memory counts. */
    global.blocks_num_outside_locals.fetch_sub(1, std::memory_order_relaxed);
    global.mem_in_use_outside_locals.fetch_sub(int64_t(size), std::memory_order_relaxed);
  }
}
 
size_t memory_usage_block_num()
{
  Global &global = get_global();
  std::lock_guard lock{global.locals_mutex};
 
  /* Count the number of active blocks. */
  int64_t blocks_num = global.blocks_num_outside_locals;
  for (const Local *local : global.locals) {
    blocks_num += local->blocks_num;
  }
  return size_t(blocks_num);
}
 
size_t memory_usage_current()
{
  Global &global = get_global();
  std::lock_guard lock{global.locals_mutex};
 
  /* Count the memory that's currently in use. */
  int64_t mem_in_use = global.mem_in_use_outside_locals;
  for (const Local *local : global.locals) {
    mem_in_use += local->mem_in_use;
  }
  return size_t(mem_in_use);
}
 
size_t memory_usage_peak()
{
  update_global_peak();
  Global &global = get_global();
  return global.peak;
}
 
void memory_usage_peak_reset()
{
  Global &global = get_global();
  global.peak = memory_usage_current();
}