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());
}
}
