前言
在深度學習模型訓練中，每次迭代過程中都涉及到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());
}
}