Paddle源碼之內存管理技術

前言

在深度學習模型訓練中，每次迭代過程中都涉及到Tensor的創建和銷燬，伴隨着的是內存的頻繁 malloc和free操作，可能對模型訓練帶來不必要的 overhead。

在主流的深度學習框架中，會藉助 chunk 機制的內存池管理技術來避免這一點。通過實事先統一申請不同 chunk size 的內存，並記錄到內存池中。創建一個Tensor時，若內存池中存在滿足需求的可用內存，則直接分配。銷燬一個Tensor時，並不馬上free掉還給系統，而是標記爲可用狀態，放在內存池供下個Tensor使用。

通過內存池管理技術，可以有效減少頻繁的malloc和free操作，避免不必要的overhead。

技術實現

chunk

每個chunk代表一段連續的存儲空間。不同的chunk按照地址升序組成雙向鏈表。每個chunk只有兩種狀態：空閒、已佔用。不存在部分使用的中間態。

在Paddle中，內存池統一通過 BuddyAllocator類來管理，下面逐一剖析相關實現。成員變量包括：

private:
    /*
     * 默認的內存分配器，支持CPUAllocator、GPUAllocator、CUDAPinnedAllocator。
     */
    std::unique_ptr<SystemAllocator> system_allocator_;

    // 用於表示一個內存段的信息
    using IndexSizeAddress = std::tuple<size_t, size_t, void*>;
    // 藉助有序的 set 存放可用的內存段
    using PoolSet = std::set<IndexSizeAddress>;
    
    PoolSet pool_; // 內存池，存放可用的不同 chunk size的內存信息
    PoolSet chunks_; // 內存池。存放從系統重新申請的內存塊

從BuddyAllocator的成員變量可以看出，不同BuddyAllocator對象可以管理不同類型的內存池，比如 CPU內存池、GPU內存池、CUDAPinned內存池。

構造函數顯式需要一個SystemAllocator來初始化：

public:
    BuddyAllocator(std::unqiue_ptr<SystemAllocator> system_allocator, size_t min_chunk_size, size_t max_chunk_size);

內存申請

BuddyAllocator如何避免內存頻繁的malloc和free操作呢？

申請內存時：

void* BuddyAllocator::Alloc(size_t unaligned_size){
  // step 1: 做內存對齊，保證申請的內存大小都是 min_chunk_size的整數倍
  size_t size = align(unaligned_size+sizeof(MemoryBlock::Desc), min_chunk_size_);
  
  // 加鎖
  std::lock_guard<std::mutex> lock(mutex_);
  
  // step 2: 如果申請內存超過 max_chunk_size_, 則交由system_allocator完成
  if(size > max_chunk_size_){
    return SystemAlloc(size);
  }
  
  // step 3: 否則，去內存池查找是否有滿足大小的可用內存塊
  auto it = FindExistChunk(size);
  
  // step 4: 若找不到，則向系統申請新內存塊，並記錄到內存池中
  if(it == pool_.end()){
    it = RefillPool(size);
    if(it == pool_.end()){
      return nullptr;
    }
  }else{
    VLOG(10)<<;
  }
  
  // step 5: 更新內存池 size 相關信息
  total_used_ += size;
  total_free_ -= size;
  
  // step 6: 若申請的size小於內存塊實際大小，則把多餘的部分切分掉，新建一個內存塊放到內存池中
  return reinterpret_cast<MemoryBlock*>(SplitToAlloc(it, size))->Data();
}

內存釋放

此處並非真正的將內存歸還給系統，而是將內存塊從佔用狀態標記爲可用狀態，並放到內存池中開放出去。

void BuddyAllocator::Free(void* p){
  // step 1: 將指針轉換爲內存塊指針
  auto block = static_cast<MemoryBlock*>(p)->MetaData();
  
  std::lock_guard<std::mutex> lock(mutex_);
  // step 2: 獲取內存塊的詳細元信息，釋放內存需要
  auto* desc = cache_.LoadDesc(block);
  if(desc->get_type() == MemoryBlock::HUGE_CHUNK){
    // 在前面申請大內存時，也是交由system_allocator完成的，解鈴還須繫鈴人
    system_allocator_->Free(block, desc->get_totoal_size(), desc->get_index());
    // 刪除內存塊對應的元信息
    cache_.Invalidate(block);
    return;
  }
  
  // step 3: 若待釋放內存塊大小在[min_chunk_size_, max_chunk_size_]之間
  block->MarkAsFree(&cache_);  // 修改元信息，標記爲 可用 狀態
  
  // step 4: 更新總內存信息
  total_used_ -= desc->get_total_size();
  total_free += desc->get_total_size();
  
  // step 5: 看是否可以將此內存塊與左右空閒的內存塊合併，避免內存碎片
  MemoryBlock* right_buddy = block->GetRightBuddy(&cache_);
  if(right_buddy){
    auto rb_desc = cache_.LoadDesc(right_buddy);
    if(rb_desc->get_type() == MemoryBlock::FREE_CHUNK){
      pool_.erase(IndexSizedAddress(rb_desc->get_index(), rb_desc->get_total_size(), right_buddy));
      block->Merge(&cache_, right_buddy);
    }
  }
   
  MemoryBlock* left_buddy = block->GetLeftBuddy(&cache_);
  // .... (省略對前序內存塊的合併操作)
  
  // step 6: 將合併後的內存塊放入到可用內存池中
  pool_.insert(IndexSizeAddress(desc->get_index(), desc->get_total_size(), block));
}

內存歸還

此階段纔是真正的將內存歸還給操作系統，此過程分爲兩個步驟：

• 把後來的、通過system_allocator_申請的內存 free掉（調用Release函數）
• 析構BuddyAllocator對象時，對內存池剩餘的內存 free掉（調用析構函數）

我們先看第一階段 Release邏輯：

uint64_t BuddyAllocator::Release(){
  // 先加鎖
  std::lock_guard<std::mutex> lock(mutex_);
  int num = 0; // 標記後來新增申請的內存塊
  uint64_t bytes = 0; // 統計總共可釋放的內存
  bool del_flag = false;
  // step 1: 有序遍歷可用內存池中的每個內存塊
  for(auto iter = pool_.begin(); iter != pool_.end()){
    auto remain_size = std::get<1>(*iter);
    auto remain_ptr = std::get<2>(*iter);
    
    for(auto& chunk : chunks_){
      auto init_size = std::get<1>(chunk);
      auto init_ptr = std::get<2>(chunk);
      // step 2: 若在之前的chunks_記錄中找到地址一樣，空間一樣的chunk
      if(init_size = remain_size && init_ptr == remain_ptr){
        ++num;
        bytes += init_size;
        total_free_ -= init_size;
        auto block = static_cast<MemoryBlock*>(init_ptr);
        // step 3: 則歸還內存給系統，標記爲此內存塊爲可回收狀態
        system_allocator_->Free(init_ptr, init_size, std::get<0>(chunk));
        cache_.Invalidate(block);
        del_flag = true;
        break;
      }
    }
    // step 4: 對於標記爲可回收狀態的內存塊，從內存池中移除
    if(del_flag){
      iter = pool_.erase(iter);
    }else{
      iter++;
    }
  }
  return bytes;
}

Release支持被顯式調用，以歸還未用到的內存給操作系統。

當BuddyAllocator對象在模型訓練結束後，會被析構掉。析構時需要保證之前申請的內存必須正確的歸還給操作系統，否則會導致內存泄露。

BuddyAllocator::~BuddyAllocator(){
  while(!pool.empty()){
    // step 1: 遍歷內存池中所有的內存塊
    auto block = static_cast<MemoryBlock*>(std::get<2>(pool_.begin()));
    auto desc = cache_.LoadDesc(block);
    // step 2: Free掉，歸還給系統
    system_allocator_->Free(block, desc->get_total_size(), desc->get_index());
    // step 3: 刪除元信息
    cache_.Invalidata(block);
    pool_.erase(pool_.begin());
  }
}

參考資料

1. Paddle 框架源碼