// cudamatrix/cu-allocator.cc
  
  // Copyright      2015-2018  Johns Hopkins University (author: Daniel Povey)
  
  // See ../../COPYING for clarification regarding multiple authors
  //
  // 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
  //
  // THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  // WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  // MERCHANTABLITY OR NON-INFRINGEMENT.
  // See the Apache 2 License for the specific language governing permissions and
  // limitations under the License.
  
  
  
  #include "cudamatrix/cu-allocator.h"
  
  #if HAVE_CUDA == 1
  
  #include <cublas_v2.h>
  #include <cuda.h>
  #include <cuda_runtime_api.h>
  
  #include <string>
  #include <vector>
  #include <algorithm>
  #ifndef _MSC_VER
  #include <dlfcn.h>
  #endif
  
  #include "cudamatrix/cu-common.h"
  #include "cudamatrix/cu-device.h"
  #include "cudamatrix/cu-matrix.h"
  #include "base/kaldi-error.h"
  #include "base/kaldi-utils.h"
  #include "util/common-utils.h"
  
  namespace kaldi {
  
  
  void* CuMemoryAllocator::Malloc(size_t size) {
    Timer tim;
    if (!opts_.cache_memory) {
      void *ans;
      CU_SAFE_CALL(cudaMalloc(&ans, size));
      double elapsed = tim.Elapsed();
      tot_time_taken_ += elapsed;
      malloc_time_taken_ += elapsed;
      t_++;
      return ans;
    }
  
    // We could perhaps change this to KALDI_PARANOID_ASSERT to save time.
    KALDI_ASSERT(size != 0);
  
    // Round up 'size' to a multiple of 256; this ensures the right kind of
    // memory alignment.
    size = (size + 255) & ~((size_t)255);
    void *ans = MallocInternal(size);
    tot_time_taken_ += tim.Elapsed();
    return ans;
  }
  
  
  CuMemoryAllocator::MemoryBlock *CuMemoryAllocator::SplitBlock(
      MemoryBlock *block, size_t size) {
    SubRegion *subregion = block->subregion;
    // new_block will become the right-most part of 'block', and 'block' will
    // be the left-most part.
    MemoryBlock *new_block = new MemoryBlock;
    bool return_new_block;
    char *new_begin;
  
    // We now decide whether to make the left part of 'block' be of size ('size')
    // and return it (the 'if' branch of the if-else block below), or the right
    // part (the 'else' branch).  We decide this based on heuristics.  Basically,
    // we want to allocate the sub-block that's either next to the edge of the
    // MemoryRegion, or next to something that was allocated long ago (and which,
    // we assume won't be deallocated for a relatively long time).  That is: we
    // want to leave the un-allocated memory next to a memory block that was
    // recently allocated (and thus is likely to be freed sooner), so that when
    // that block is freed we can merge it with the still-unallocated piece into a
    // larger block; this will reduce fragmentation.  But if this block spans
    // multiple sub-regions we don't want to do that, as that would be against our
    // heuristic of, where possible, allocating memory from lower-numbered
    // sub-regions.
    //
    // Bear in mind that we can assume block->next and block->prev, if they are
    // non-NULL, are both currently allocated, since 'block' is un-allocated and
    // we would have merged any adjacent un-allocated sub-regions.
    if (block->next != NULL && block->prev != NULL &&
        block->prev->t < block->next->t &&
        block->next->subregion == subregion) {
      // We'll allocate the right part of the block, since the left side is next
      // to a relatively recently-allocated block.
      return_new_block = true;
      new_begin = block->end - size;
    } else {
      // We'll allocate the left part of the block.
      return_new_block = false;
      new_begin = block->begin + size;
    }
  
    // The following code makes sure the SubRegion for 'new_block' is correct,
    // i.e. its 'begin' is >= the 'begin' of the subregion and < the 'end' of the
    // subregion.  If the following loop segfaults, it indicates a bug somewhere
    // else.
    while (new_begin >= subregion->end)
      subregion = subregion->next;
    MemoryBlock *next_block = block->next;
    new_block->begin = new_begin;
    new_block->end = block->end;
    new_block->subregion = subregion;
    new_block->allocated = false;
    new_block->thread_id = block->thread_id;
    new_block->t = block->t;
    new_block->next = next_block;
    new_block->prev = block;
    if (next_block)
      next_block->prev = new_block;
    block->next = new_block;
    block->end = new_begin;
  
    // Add the split-up piece that we won't be allocating, to the
    // 'free_blocks' member of its subregion.
    if (return_new_block) {
      AddToFreeBlocks(block);
      return new_block;
    } else {
      AddToFreeBlocks(new_block);
      return block;
    }
  }
  
  
  void CuMemoryAllocator::RemoveFromFreeBlocks(MemoryBlock *block) {
    SubRegion *subregion = block->subregion;
    size_t block_size = block->end - block->begin;
    std::pair<size_t, MemoryBlock*> p(block_size, block);
    size_t num_removed = subregion->free_blocks.erase(p);
    KALDI_ASSERT(num_removed != 0);
    // Update largest_free_block_, if needed.
    size_t subregion_index = subregion->subregion_index;
    if (block_size == largest_free_block_[subregion_index]) {
      if (subregion->free_blocks.empty())
        largest_free_block_[subregion_index] = 0;
      else
        largest_free_block_[subregion_index] =
            subregion->free_blocks.begin()->first;
    }
  }
  
  void CuMemoryAllocator::AddToFreeBlocks(MemoryBlock *block) {
    SubRegion *subregion = block->subregion;
    KALDI_PARANOID_ASSERT(block->begin >= subregion->begin &&
                          block->begin < subregion->end);
    size_t block_size = block->end - block->begin,
         subregion_index = subregion->subregion_index;
    // Update largest_free_block_, if needed.
    if (block_size > largest_free_block_[subregion_index]) {
      largest_free_block_[subregion_index] = block_size;
    }
    subregion->free_blocks.insert(std::pair<size_t, MemoryBlock*>(block_size, block));
  }
  
  
  void* CuMemoryAllocator::MallocFromSubregion(SubRegion *subregion,
                                               size_t size) {
    // NULL is implementation defined and doesn't have to be zero so we can't
    // guarantee that NULL will be <= a valid pointer-- so we cast to a pointer
    // from zero instead of using NULL.
    std::pair<size_t, MemoryBlock*> p(size, (MemoryBlock*)0);
  
    std::set<std::pair<size_t, MemoryBlock*> >::iterator iter =
        subregion->free_blocks.lower_bound(p);
    // so now 'iter' is the first member of free_blocks whose size_t value is >=
    // size.  If 'iter' was equal to the end() of that multi_map, it would be a
    // bug because the calling code checked that the largest free block in this
    // region was sufficiently large.  We don't check this; if it segfaults, we'll
    // debug.
  
    // search for a block that we don't have to synchronize on
    int max_iters = 20;
    auto search_iter = iter;
    for (int32 i = 0;
         search_iter != subregion->free_blocks.end() && i < max_iters;
         ++i, ++search_iter) {
      if (search_iter->second->thread_id == std::this_thread::get_id() ||
          search_iter->second->t <= synchronize_gpu_t_) {
        iter = search_iter;
        break;
      }
    }
  
    MemoryBlock *block = iter->second;
    // Erase 'block' from its subregion's free blocks list... the next lines are
    // similar to RemoveFromFreeBlocks(), but we code it directly as we have the
    // iterator here, and it would be wasteful to do another lookup.
    subregion->free_blocks.erase(iter);
    // Update largest_free_block_, if needed.  The following few lines of code also appear
    // in RemoveFromFreeBlocks().
    size_t block_size = block->end - block->begin,
        subregion_index = subregion->subregion_index;
    if (block_size == largest_free_block_[subregion_index]) {
      if (subregion->free_blocks.empty())
        largest_free_block_[subregion_index] = 0;
      else
        largest_free_block_[subregion_index] =
            subregion->free_blocks.begin()->first;
    }
  
    KALDI_PARANOID_ASSERT(block_size >= size && block->allocated == false);
  
    // the most memory we allow to be 'wasted' by failing to split a block, is the
    // smaller of: 1/16 of the size we're allocating, or half a megabyte.
    size_t allowed_extra_size = std::min<size_t>(size >> 4, 524288);
    if (block_size > size + allowed_extra_size) {
      // If the requested block is substantially larger than what was requested,
      // split it so we don't waste memory.
      block = SplitBlock(block, size);
    }
  
    if (std::this_thread::get_id() != block->thread_id &&
        block->t > synchronize_gpu_t_) {
      // see NOTE ON SYNCHRONIZATION in the header.
      SynchronizeGpu();
      synchronize_gpu_t_ = t_;
      num_synchronizations_++;
    }
    block->allocated = true;
    block->t = t_;
    allocated_block_map_[block->begin] = block;
    allocated_memory_ += (block->end - block->begin);
    if (allocated_memory_ > max_allocated_memory_) 
      max_allocated_memory_ = allocated_memory_;
    return block->begin;
  }
  
  // By the time MallocInternal is called, we will have ensured that 'size' is
  // a nonzero multiple of 256 (for memory aligment reasons).
  // inline
  void* CuMemoryAllocator::MallocInternal(size_t size) {
  start:
    std::vector<size_t>::const_iterator iter = largest_free_block_.begin(),
        end = largest_free_block_.end();
    size_t subregion_index = 0;
    for (; iter != end; ++iter, ++subregion_index) {
      if (*iter > size) {
        return MallocFromSubregion(subregions_[subregion_index], size);
      }
    }
    // We dropped off the loop without finding a subregion with enough memory
    // to satisfy the request -> allocate a new region.
    AllocateNewRegion(size);
    // An infinite loop shouldn't be possible because after calling
    // AllocateNewRegion(size), there should always be a SubRegion
    // with that size available.
    goto start;
  }
  
  // Returns max(0, floor(log_2(i))).   Not tested independently.
  static inline size_t IntegerLog2(size_t i) {
    size_t ans = 0;
    while (i > 256) {
      i >>= 8;
      ans += 8;
    }
    while (i > 16) {
      i >>= 4;
      ans += 4;
    }
    while (i > 1) {
      i >>= 1;
      ans++;
    }
    return ans;
  }
  
  std::string GetFreeGpuMemory(int64* free, int64* total) {
  #ifdef _MSC_VER
    size_t mem_free, mem_total;
    cuMemGetInfo_v2(&mem_free, &mem_total);
  #else
    // define the function signature type
    size_t mem_free, mem_total;
    {
      // we will load cuMemGetInfo_v2 dynamically from libcuda.so
      // pre-fill ``safe'' values that will not cause problems
      mem_free = 1; mem_total = 1;
      // open libcuda.so
      void* libcuda = dlopen("libcuda.so", RTLD_LAZY);
      if (NULL == libcuda) {
        KALDI_WARN << "cannot open libcuda.so";
      } else {
        // define the function signature type
        // and get the symbol
        typedef CUresult (*cu_fun_ptr)(size_t*, size_t*);
        cu_fun_ptr dl_cuMemGetInfo = (cu_fun_ptr)dlsym(libcuda,"cuMemGetInfo_v2");
        if (NULL == dl_cuMemGetInfo) {
          KALDI_WARN << "cannot load cuMemGetInfo from libcuda.so";
        } else {
          // call the function
          dl_cuMemGetInfo(&mem_free, &mem_total);
        }
        // close the library
        dlclose(libcuda);
      }
    }
  #endif
    // copy the output values outside
    if (NULL != free) *free = mem_free;
    if (NULL != total) *total = mem_total;
    // prepare the text output
    std::ostringstream os;
    os << "free:" << mem_free/(1024*1024) << "M, "
       << "used:" << (mem_total-mem_free)/(1024*1024) << "M, "
       << "total:" << mem_total/(1024*1024) << "M, "
       << "free/total:" << mem_free/(float)mem_total;
    return os.str();
  }
  
  void CuMemoryAllocator::PrintMemoryUsage() const {
    if (!opts_.cache_memory) {
      KALDI_LOG << "Not caching allocations; time taken in "
                << "malloc/free is " << malloc_time_taken_
                << "/" << (tot_time_taken_ - malloc_time_taken_)
                << ", num operations is " << t_
                << "; device memory info: "
                << GetFreeGpuMemory(NULL, NULL);
      return;
    }
  
    size_t num_blocks_allocated = 0, num_blocks_free = 0,
        memory_allocated = 0, memory_held = 0,
        largest_free_block = 0, largest_allocated_block = 0;
  
    for (size_t i = 0; i < memory_regions_.size(); i++) {
      MemoryBlock *m = memory_regions_[i].block_begin;
      KALDI_ASSERT(m->begin == memory_regions_[i].begin);
      for (; m != NULL; m = m->next) {
        size_t size = m->end - m->begin;
        if (m->allocated) {
          num_blocks_allocated++;
          memory_allocated += size;
          if (size > largest_allocated_block)
            largest_allocated_block = size;
        } else {
          num_blocks_free++;
          if (size > largest_free_block)
            largest_free_block = size;
        }
        memory_held += size;
        // The following is just some sanity checks; this code is rarely called so
        // it's a reasonable place to put them.
        if (m->next) {
          KALDI_ASSERT(m->next->prev == m && m->end == m->next->begin);
        } else {
          KALDI_ASSERT(m->end == memory_regions_[m->subregion->memory_region].end);
        }
      }
    }
    KALDI_LOG << "Memory usage: " << memory_allocated << "/"
              << memory_held << " bytes currently allocated/total-held; "
              << num_blocks_allocated << "/" << num_blocks_free
              << " blocks currently allocated/free; largest "
              << "free/allocated block sizes are "
              << largest_allocated_block << "/" << largest_free_block
              << "; time taken total/cudaMalloc is "
              << tot_time_taken_ << "/" << malloc_time_taken_
              << ", synchronized the GPU " << num_synchronizations_
              << " times out of " << (t_/2) << " frees; "
              << "device memory info: " << GetFreeGpuMemory(NULL, NULL)
              << "maximum allocated: " << max_allocated_memory_  
              << "current allocated: " << allocated_memory_; 
  }
  
  // Note: we just initialize with the default options, but we can change it later
  // (as long as it's before we first use the class) by calling SetOptions().
  CuMemoryAllocator::CuMemoryAllocator():
      opts_(CuAllocatorOptions()),
      t_(0),
      synchronize_gpu_t_(0),
      num_synchronizations_(0),
      tot_time_taken_(0.0),
      malloc_time_taken_(0.0),
      max_allocated_memory_(0),
      allocated_memory_(0) {
    // Note: we don't allocate any memory regions at the start; we wait for the user
    // to call Malloc() or MallocPitch(), and then allocate one when needed.
  }
  
  
  void* CuMemoryAllocator::MallocPitch(size_t row_bytes,
                                       size_t num_rows,
                                       size_t *pitch) {
    Timer tim;
    if (!opts_.cache_memory) {
      void *ans;
      CU_SAFE_CALL(cudaMallocPitch(&ans, pitch, row_bytes, num_rows));
      double elapsed = tim.Elapsed();
      tot_time_taken_ += elapsed;
      malloc_time_taken_ += elapsed;
      return ans;
    }
  
    // Round up row_bytes to a multiple of 256.
    row_bytes = (row_bytes + 255) & ~((size_t)255);
    *pitch = row_bytes;
    void *ans = MallocInternal(row_bytes * num_rows);
    tot_time_taken_ += tim.Elapsed();
    return ans;
  }
  
  void CuMemoryAllocator::Free(void *ptr) {
    Timer tim;
    if (!opts_.cache_memory) {
      CU_SAFE_CALL(cudaFree(ptr));
      tot_time_taken_ += tim.Elapsed();
      t_++;
      return;
    }
    t_++;
    unordered_map<void*, MemoryBlock*>::iterator iter =
        allocated_block_map_.find(ptr);
    if (iter == allocated_block_map_.end()) {
      KALDI_ERR << "Attempt to free CUDA memory pointer that was not allocated: "
                << ptr;
    }
    MemoryBlock *block = iter->second;
    allocated_memory_ -= (block->end - block->begin);
    allocated_block_map_.erase(iter);
    block->t = t_;
    block->thread_id = std::this_thread::get_id();
    block->allocated = false;
  
    // If this is not the first block of the memory region and the previous block
    // is not allocated, merge this block into the previous block.
    MemoryBlock *prev_block = block->prev;
    if (prev_block != NULL && !prev_block->allocated) {
      RemoveFromFreeBlocks(prev_block);
      prev_block->end = block->end;
      if (prev_block->thread_id != block->thread_id) {
        // the two blocks we're merging were freed by different threads, so we
        // give the 'nonexistent thread' as their thread, which means that
        // whichever thread requests that block, we force synchronization.  We can
        // assume that prev_block was previously allocated (prev_block->t > 0)
        // because we always start from the left when allocating blocks, and we
        // know that this block was previously allocated.
        prev_block->thread_id = std::thread::id();
      }
      prev_block->t = t_;
      prev_block->next = block->next;
      if (block->next)
        block->next->prev = prev_block;
      delete block;
      block = prev_block;
    }
  
    // If this is not the last block of the memory region and the next block is
    // not allocated, merge the next block into this block.
    MemoryBlock *next_block = block->next;
    if (next_block != NULL && !next_block->allocated) {
      // merge next_block into 'block', deleting 'next_block'.  Note: at this
      // point, if we merged with the previous block, the variable 'block' may now
      // be pointing to that previous block, so it would be a 3-way merge.
      RemoveFromFreeBlocks(next_block);
      block->end = next_block->end;
      if (next_block->thread_id != block->thread_id && next_block->t > 0) {
        // the two blocks we're merging were freed by different threads, so we
        // give the 'nonexistent thread' as their thread, which means that
        // whichever thread requests that block, we force synchronization.  there
        // is no need to do this if next_block->t == 0, which would mean it had
        // never been allocated.
        block->thread_id = std::thread::id();
      }
      // We don't need to inspect the 't' value of next_block; it can't be
      // larger than t_ because t_ is now.
      block->next = next_block->next;
      if (block->next)
        block->next->prev = block;
      delete next_block;
    }
    AddToFreeBlocks(block);
    tot_time_taken_ += tim.Elapsed();
  }
  
  void CuMemoryAllocator::AllocateNewRegion(size_t size) {
    int64 free_memory, total_memory;
    std::string mem_info = GetFreeGpuMemory(&free_memory, &total_memory);
    opts_.Check();
    size_t region_size = static_cast<size_t>(free_memory * opts_.memory_proportion);
    if (region_size < size)
      region_size = size;
    // Round up region_size to an exact multiple of 1M (note: we expect it will
    // be much larger than that).  1048575 is 2^20 - 1.
    region_size = (region_size + 1048575) & ~((size_t)1048575);
  
    if (!memory_regions_.empty()) {
      // If this is not the first region allocated, print some information.
      KALDI_LOG << "About to allocate new memory region of " << region_size
                << " bytes; current memory info is: " << mem_info;
    }
    void *memory_region;
    cudaError_t e;
    {
      Timer tim;
      e = cudaMalloc(&memory_region, region_size);
      malloc_time_taken_ += tim.Elapsed();
    }
    if (e != cudaSuccess) {
      PrintMemoryUsage();
      if (!CuDevice::Instantiate().IsComputeExclusive()) {
        KALDI_ERR << "Failed to allocate a memory region of " << region_size
                  << " bytes.  Possibly this is due to sharing the GPU.  Try "
                  << "switching the GPUs to exclusive mode (nvidia-smi -c 3) and using "
                  << "the option --use-gpu=wait to scripts like "
                  << "steps/nnet3/chain/train.py.  Memory info: "
                  << mem_info
                  << " CUDA error: '" << cudaGetErrorString(e) << "'";
      } else {
        KALDI_ERR << "Failed to allocate a memory region of " << region_size
                  << " bytes.  Possibly smaller minibatch size would help.  "
                  << "Memory info: " << mem_info
                  << " CUDA error: '" << cudaGetErrorString(e) << "'";
      }
    }
    // this_num_subregions would be approximately 'opts_.num_subregions' if
    // 'region_size' was all the device's memory.  (We add one to round up).
    // We're aiming to get a number of sub-regions approximately equal to
    // opts_.num_subregions by the time we allocate all the device's memory.
    size_t this_num_subregions = 1 +
        (region_size * opts_.num_subregions) / total_memory;
  
    size_t memory_region_index = memory_regions_.size();
    memory_regions_.resize(memory_region_index + 1);
    MemoryRegion &this_region = memory_regions_.back();
  
    this_region.begin = static_cast<char*>(memory_region);
    this_region.end = this_region.begin + region_size;
    // subregion_size will be hundreds of megabytes.
    size_t subregion_size = region_size / this_num_subregions;
  
    std::vector<SubRegion*> new_subregions;
    char* subregion_begin = static_cast<char*>(memory_region);
    for (size_t i = 0; i < this_num_subregions; i++) {
      SubRegion *subregion = new SubRegion();
      subregion->memory_region = memory_region_index;
      subregion->begin = subregion_begin;
      if (i + 1 == this_num_subregions) {
        subregion->end = this_region.end;
        KALDI_ASSERT(subregion->end > subregion->begin);
      } else {
        subregion->end = subregion_begin + subregion_size;
        subregion_begin = subregion->end;
      }
      subregion->next = NULL;
      if (i > 0) {
        new_subregions.back()->next = subregion;
      }
      new_subregions.push_back(subregion);
    }
    // Initially the memory is in a single block, owned by
    // the first subregion.  It will be split up gradually.
    MemoryBlock *block = new MemoryBlock();
    block->begin = this_region.begin;
    block->end = this_region.end;
    block->subregion = new_subregions.front();
    block->allocated = false;
    block->t = 0; // was never allocated.
    block->next = NULL;
    block->prev = NULL;
    for (size_t i = 0; i < this_num_subregions; i++)
      subregions_.push_back(new_subregions[i]);
    SortSubregions();
    this_region.block_begin = block;
  
    AddToFreeBlocks(block);
  }
  
  // We sort the sub-regions according to the distance between the start of the
  // MemoryRegion of which they are a part, and the start of the SubRegion.  This
  // will generally mean that the highest-numbered SubRegion-- the one we keep
  // free at all costs-- will be the end of the first block which we allocated
  // (which under most situations will be the largest block).
  void CuMemoryAllocator::SortSubregions() {
    largest_free_block_.resize(subregions_.size());
  
    std::vector<std::pair<size_t, SubRegion*> > pairs;
    for (size_t i = 0; i < subregions_.size(); i++) {
      SubRegion *subregion = subregions_[i];
      MemoryRegion &memory_region = memory_regions_[subregion->memory_region];
      size_t distance = subregion->begin - memory_region.begin;
      pairs.push_back(std::pair<size_t, SubRegion*>(distance, subregion));
    }
    std::sort(pairs.begin(), pairs.end());
    for (size_t i = 0; i < subregions_.size(); i++) {
      subregions_[i] = pairs[i].second;
      subregions_[i]->subregion_index = i;
      if (subregions_[i]->free_blocks.empty())
        largest_free_block_[i] = 0;
      else
        largest_free_block_[i] = subregions_[i]->free_blocks.begin()->first;
    }
  }
  
  CuMemoryAllocator::~CuMemoryAllocator() {
    // We mainly free these blocks of memory so that cuda-memcheck doesn't report
    // spurious errors.
    for (size_t i = 0; i < memory_regions_.size(); i++) {
      // No need to check the return status here-- the program is exiting anyway.
      cudaFree(memory_regions_[i].begin);
    }
    for (size_t i = 0; i < subregions_.size(); i++) {
      SubRegion *subregion = subregions_[i];
      for (auto iter = subregion->free_blocks.begin();
           iter != subregion->free_blocks.end(); ++iter)
        delete iter->second;
      delete subregion;
    }
  }
  
  
  CuMemoryAllocator g_cuda_allocator;
  
  
  }  // namespace kaldi
  
  
  #endif // HAVE_CUDA
  
  
  namespace kaldi {
  
  // Define/initialize this global variable.  It was declared in cu-allocator.h.
  // This has to be done outside of the ifdef, because we register the options
  // whether or not CUDA is compiled in (so that the binaries accept the same
  // options).
  CuAllocatorOptions g_allocator_options;
  
  }