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 框架源码