text: /**
   * \file src/core/impl/graph/var_node_mem_mgr/static_mem_alloc/impl.cpp
   * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
   *
   * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   */
  
  #include "./impl.h"
  #include "./interval_move.h"
  #include "./best_fit.h"
  #include "./pushdown.h"
  
  #include <map>
  
  #if MGB_ENABLE_DEBUG_UTIL
  #include "megbrain/graph/var_node.h"
  #include "megbrain/graph/operator_node.h"
  #include <cstdio>
  #include <fstream>
  #endif
  
  using namespace mgb;
  using namespace cg;
  
  constexpr size_t StaticMemAllocImplHelper::INVALID;
  
  StaticMemAllocImplHelper::Interval*
  StaticMemAllocImplHelper::Interval::overwrite_dest_root_path_compression() {
      auto &&ptr = m_overwrite_dest_root;
      if (!ptr)
          return this;
      auto root = ptr->overwrite_dest_root_path_compression();
      if (root != ptr) {
          m_offset_in_overwrite_dest_root += ptr->m_offset_in_overwrite_dest_root;
          ptr = root;
      }
      return root;
  }
  
  void StaticMemAllocImplHelper::init_overwrite_dest() {
      for (auto &&spec: m_overwrite_spec) {
          auto src = m_interval_storage.data() + std::get<0>(spec),
               dest = m_interval_storage.data() + std::get<1>(spec);
          // src overwrites a part in dest
          size_t offset = std::get<2>(spec);
          mgb_assert(src->time_begin < dest->time_end);
  
          auto orig_src = dest->m_overwrite_src;
  
          // each interval could only be overwritten by one interval, and we
          // prefer the interval with largest size to be overwritter
          if (src->time_begin == dest->time_end - 1 && !src->m_overwrite_dest &&
                  (!orig_src || src->size > orig_src->size)) {
              if (orig_src) {
                  orig_src->m_overwrite_dest = nullptr;
                  orig_src->m_offset_in_overwrite_dest = 0;
              }
              dest->m_overwrite_src = src;
              src->m_overwrite_dest = dest;
              src->m_offset_in_overwrite_dest = offset;
          }
      }
  
      for (auto &&i: m_interval_storage) {
          if (i.m_overwrite_dest) {
              i.m_overwrite_dest_root = i.m_overwrite_dest;
              i.m_offset_in_overwrite_dest_root = i.m_offset_in_overwrite_dest;
              mgb_assert(i.m_overwrite_dest->m_overwrite_src == &i);
          }
          if (i.m_overwrite_src)
              mgb_assert(i.m_overwrite_src->m_overwrite_dest == &i);
      }
  
      for (auto &&i: m_interval_storage)
          i.overwrite_dest_root_path_compression();
  }
  
  size_t StaticMemAllocImplHelper::add(size_t begin, size_t end, size_t size,
          UserKeyType key) {
  
      mgb_assert(begin < end);
      auto id = m_interval_storage.size();
      m_interval_storage.push_back({begin, end, size + m_padding, key, id});
      return id;
  }
  
  StaticMemAlloc& StaticMemAllocImplHelper::add_overwrite_spec(
          size_t iid_src, size_t iid_dest, size_t offset) {
      auto &&src = m_interval_storage.at(iid_src),
           &&dest = m_interval_storage.at(iid_dest);
      mgb_assert(iid_src != iid_dest);
      mgb_assert(offset + src.size <= dest.size);
      m_overwrite_spec.emplace_back(iid_src, iid_dest, offset);
      return *this;
  }
  
  size_t StaticMemAllocImplHelper::get_start_addr(UserKeyType key) const {
      return m_userkey2itrv.at(key)->addr_begin;
  }
  
  StaticMemAlloc& StaticMemAllocImplHelper::solve() {
      dbg_dump_interval_list();
      dbg_load_interval_list();
      m_interval.clear();
      m_interval.reserve(m_interval_storage.size());
      m_userkey2itrv.clear();
      for (auto &&i: m_interval_storage) {
          m_interval.push_back(&i);
          auto ist = m_userkey2itrv.insert({i.key, &i});
          mgb_assert(ist.second, "duplicated user key");
      }
  
      init_overwrite_dest();
  
      do_solve();
  
      check_result_and_calc_lower_bound();
  
      return *this;
  }
  
  void StaticMemAllocImplHelper::dbg_dump_interval_list() {
  #if MGB_ENABLE_DEBUG_UTIL
      const char *fdir = MGB_GETENV("MGB_DUMP_INTERVAL_LIST_DIR");
      if (!fdir)
          return;
      static int run_id = 0;
      auto fpath = ssprintf("%s/mgb-interval-%d.txt", fdir, run_id ++);
      mgb_log_warn("dump static mem alloc interval list to %s", fpath.c_str());
      FILE *fout = fopen(fpath.c_str(), "w");
      mgb_assert(fout, "failed to open %s", fpath.c_str());
  
      fprintf(fout, "%zu\n"_fmt, m_interval_storage.size());
      for (auto &&i: m_interval_storage)
          fprintf(fout, "%zu %zu %zu\n", i.time_begin_orig, i.time_end_orig,
                  i.size_orig);
  
      fprintf(fout, "%zu\n"_fmt, m_overwrite_spec.size());
      for (auto &&i: m_overwrite_spec)
          fprintf(fout, "%zu %zu %zu\n"_fmt, std::get<0>(i),
                  std::get<1>(i), std::get<2>(i));
  
      fclose(fout);
  #endif
  }
  
  void StaticMemAllocImplHelper::dbg_load_interval_list() {
  #if MGB_ENABLE_DEBUG_UTIL
      const char *fpath = MGB_GETENV("MGB_LOAD_INTERVAL");
      if (!fpath)
          return;
      unsetenv("MGB_DUMP_INTERVAL_LIST_DIR");
      unsetenv("MGB_LOAD_INTERVAL");
      mgb_log_warn("load interval from %s for debug", fpath);
      std::ifstream fin(fpath);
      mgb_assert(fin.good(), "failed to open %s", fpath);
  
      m_interval_storage.clear();
      m_overwrite_spec.clear();
      size_t nr_interval;
      fin >> nr_interval;
      for (size_t i = 0; i < nr_interval; ++ i) {
          size_t begin, end, size;
          fin >> begin >> end >> size;
          add(begin, end, size, reinterpret_cast<UserKeyType>(i));
      }
  
      size_t nr_overwrite;
      fin >> nr_overwrite;
      for (size_t i = 0; i < nr_overwrite; ++ i) {
          size_t s, d, o;
          fin >> s >> d >> o;
          add_overwrite_spec(s, d, o);
      }
  
      solve();
  
      printf("allocation result tot_alloc=%zu(%.2fMiB):\n",
              tot_alloc(), tot_alloc() / 1024.0 / 1024.0);
      for (auto &&i: m_interval_storage) {
          printf("id=%zu size=%zu(%.2fMiB) time=[%zu, %zu) addr=[%zu, %zu)\n",
                  i.id, i.size_orig, i.size_orig / 1024.0 / 1024,
                  i.time_begin_orig, i.time_end_orig,
                  i.addr_begin, i.addr_end());
      }
      fflush(stdout);
  
      mgb_trap();
  #endif
  }
  
  void StaticMemAllocImplHelper::check_result_and_calc_lower_bound() {
      size_t peak = 0;
  
      // time => pair(alloc, free)
      using TimeEvent = std::pair<IntervalPtrArray, IntervalPtrArray>;
      std::map<size_t, TimeEvent> time2event;
  
      for (auto &&i: m_interval_storage) {
          mgb_assert(i.addr_begin != INVALID);
          time2event[i.time_begin_orig].first.push_back(&i);
          time2event[i.time_end_orig].second.push_back(&i);
          update_max(peak, i.addr_end());
          if (i.is_overwrite_root()) {
              // modify size for calc lower bound
              i.size = align(i.size_orig);
              mgb_assert(i.addr_begin == align(i.addr_begin));
          } else {
              auto offset = i.offset_in_overwrite_dest_root();
              i.size = align(offset + i.size_orig) - (
                      offset - (offset & (m_alignment - 1)));
          }
      }
      mgb_assert(peak <= tot_alloc() && align(peak) == align(tot_alloc()));
  
      // get lower bound
      {
          m_peak_lower_bound = 0;
          size_t usage = 0;
          for (auto &&tpair: time2event) {
              for (auto i: tpair.second.first) {
                  if (i->is_overwrite_root())
                      usage += i->size;
              }
              for (auto &&i: tpair.second.second) {
                  usage -= i->size;
                  if (i->m_overwrite_src) {
                      // this interval is overwritten by another one, so count its
                      // size in current usage
                      usage += i->m_overwrite_src->size;
                  }
              }
              update_max(m_peak_lower_bound, usage);
          }
          mgb_assert(!usage);
      }
  
      print_bottleneck_oprs(time2event);
  
      // restore time and size; check overwrite addr
      for (auto &&i: m_interval_storage) {
          i.time_begin = i.time_begin_orig;
          i.time_end = i.time_end_orig;
          i.size = i.size_orig;
  
          if (!i.is_overwrite_root()) {
              mgb_assert(i.overwrite_dest()->addr_begin +
                      i.offset_in_overwrite_dest() == i.addr_begin);
          }
      }
  
      std::map<size_t, Interval*> cur_allocated;
      IntervalPtrArray id_overwriter;
  
      auto remove_alloc = [&](Interval *i) {
          auto iter = cur_allocated.find(i->addr_begin);
          mgb_assert(iter != cur_allocated.end() && iter->second == i);
          cur_allocated.erase(iter);
  
          if (auto s = i->overwrite_src()) {
              auto ins = cur_allocated.insert({s->addr_begin, s});
              mgb_assert(ins.second);
          }
      };
  
      // check for conflicts
      for (auto &&tpair: time2event) {
  
          // free and set overwriter addr
          id_overwriter.clear();
          for (auto i: tpair.second.second) {
              if (!i->is_overwrite_root() &&
                      i->time_end_orig == i->overwrite_dest()->time_end_orig &&
                      !i->offset_in_overwrite_dest()) {
                  // a overwrites b, a and b share same time end, zero offset
                  mgb_assert(i->addr_begin == i->overwrite_dest()->addr_begin);
                  id_overwriter.push_back(i);
                  continue;
              }
              remove_alloc(i);
          }
          for (auto i: id_overwriter)
              remove_alloc(i);
  
          // alloc
          for (auto i: tpair.second.first) {
              auto iter = cur_allocated.lower_bound(i->addr_begin);
  
              if (i->is_overwrite_root()) {
                  if (iter != cur_allocated.end()) {
                      mgb_assert(i->addr_end() <= iter->first);
                  }
                  if (!cur_allocated.empty() && iter != cur_allocated.begin()) {
                      -- iter;
                      mgb_assert(iter->second->addr_end() <= i->addr_begin);
                  }
                  cur_allocated[i->addr_begin] = i;
              }
          }
      }
  
      mgb_assert(cur_allocated.empty());
  }
  
  template<typename T>
  void StaticMemAllocImplHelper::print_bottleneck_oprs(const T& time2event) {
  #if MGB_ENABLE_DEBUG_UTIL
      if (!MGB_GETENV("MGB_PRINT_STATIC_ALLOC_BOTTLENECK"))
          return;
      mgb_assert(dbg_key2varnode);
  
      size_t peak = 0, usage = 0;
      std::unordered_set<UserKeyType> alive, peak_alive;
  
      for (auto &&tpair: time2event) {
          for (Interval* i: tpair.second.first) {
              if (i->is_overwrite_root()) {
                  usage += i->size;
                  alive.insert(i->key);
              }
          }
          for (Interval *i: tpair.second.second) {
              usage -= i->size;
              if (i->m_overwrite_src) {
                  usage += i->m_overwrite_src->size;
                  alive.insert(i->m_overwrite_src->key);
              }
          }
          for (Interval *i: tpair.second.second) {
              alive.erase(i->key);
          }
          if (usage > peak) {
              peak = usage;
              peak_alive = alive;
          }
      }
      mgb_assert(!usage && alive.empty());
  
      printf("mgb static alloc bottleneck: size=%.3fMiB {\n",
              peak / 1024.0 / 1024);
      using SizeVar = std::tuple<size_t, size_t, VarNode*>;
      std::vector<SizeVar> vars;
      for (auto i: peak_alive) {
          auto var = dbg_key2varnode(i);
          vars.emplace_back(var->mem_plan().chunk().size(), var->id(), var);
      }
      auto cmp = [](const SizeVar &a, const SizeVar &b) {
          auto sza = std::get<0>(a), szb = std::get<0>(b);
          return sza > szb || (sza == szb && std::get<1>(a) < std::get<1>(b));
      };
      std::sort(vars.begin(), vars.end(), cmp);
      for (auto &&i: vars) {
          auto size = std::get<0>(i);
          VarNode* v = std::get<2>(i);
          OperatorNodeBase* o = v->owner_opr();
          printf("  var%zu %s owner=%s{%s} shape=%s alloc_size=%zu\n",
                  v->id(), v->cname(), o->cname(), o->dyn_typeinfo()->name,
                  v->shape().to_string().c_str(),
                  size);
      }
      printf("}\n");
  #endif
  }
  
  StaticMemAllocImplHelper::~StaticMemAllocImplHelper() noexcept = default;
  
  std::unique_ptr<StaticMemAlloc> StaticMemAlloc::make(AllocatorAlgo algo) {
      switch (algo) {
  #if !MGB_BUILD_SLIM_SERVING
          case AllocatorAlgo::INTERVAL_MOVE:
              return std::make_unique<StaticMemAllocIntervalMove>();
          case AllocatorAlgo::BEST_FIT:
              return std::make_unique<StaticMemAllocBestFit>();
  #endif
          case AllocatorAlgo::PUSHDOWN:
              return std::make_unique<StaticMemAllocPushdown>();
          default:
              mgb_assert(0, "unknown mem allocator algorithm");
      }
  }
  
  // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}
  
  