TensorFlow内存管理bfc算法实例

2023-08-02 11:22:05 73

1.基本介绍

tensorflow设备内存管理模块实现了一个best-fitwithcoalescing算法（后文简称bfc算法）。

bfc算法是DoungLea'smalloc(dlmalloc)的一个非常简单的版本。

它具有内存分配、释放、碎片管理等基本功能。

2.bfc基本算法思想

1.数据结构

整个内存空间由一个按基址升序排列的Chunk双向链表来表示，它们的直接前趋和后继必须在地址连续的内存空间。Chunk结构体里含有实际大小、请求大小、是否被占用、基址、直接前趋、直接后继、Bin索引等信息。

2.申请

用户申请一个内存块（malloc)。根据chunk双链表找到一个合适的内存块，如果该内存块的大小是用户申请的大小的二倍以上，那么就将该内存块切分成两块，这就是split操作。

返回其中一块给用户，并将该内存块标识为占用

Spilt操作会新增一个chunk，所以需要修改chunk双链表以维持前驱和后继关系

如果用户申请512的空间，正好有一块1024的chunk2是空闲的，由于1024/512=2，所以chunk2被split为2块：chunk2_1和chunk2_2。返回chunk2_1给用户并将其标志位占用状态。

3.释放

用户释放一个内存块（free）。先将该块标记为空闲。然后根据chunk数据结构中的信息找到其前驱和后继内存块。如果前驱和后继块中有空闲的块，那么将刚释放的块和空闲的块合并成一个更大的chunk（这就是merge操作，合并当前块和其前后的空闲块）。再修改双链表结构以维持前驱后继关系。这就做到了内存碎片的回收。

如果用户要freechunk3，由于chunk3的前驱chunk2也是空闲的，所以将chunk2和chunk3合并得到一个新的chunk2'，大小为chunk2和chunk3之和。

3.bins

1.bins数据结构

bfc算法采取的是被动分块的策略。最开始整个内存是一个chunk，随着用户申请空间的次数增加，最开始的大chunk会被不断的split开来，从而产生越来越多的小chunk。当chunk数量很大时，为了寻找一个合适的内存块而遍历双链表无疑是一笔巨大的开销。为了实现对空闲块的高效管理，bfc算法设计了bin这个抽象数据结构。

每个bin都有一个size属性，一个bin是一个拥有chunksize>=binsize的空闲chunk的集合。集合中的chunk按照chunksize的升序组织成单链表。bfc算法维护了一个bin的集合：bins。它由多个bin以及从属于每个bin的chunks组成。内存中所有的空闲chunk都由bins管理。

图中每一列表示一个bin，列首方格中的数字表示bin的size。binsize的大小都是256的2^n的倍。每个bin下面挂载了一系列的空闲chunk，每个chunk的chunksize都大于等于所属的bin的binsize，按照chunksize的升序挂载成单链表。

2.bins操作

bfc算法针对bins这个集合设计了三个操作：search、insert、delete。

search

给定一个chunksize，从bins中找到大于等于该chunksize的最小的那个空闲chunk。Search操作具体流程如下。如果bin以数组的形式组织，那么可以从index=chunksize/256>>2的那个bin开始查找。最好的情况是开始查找的那个bin的chunk链表非空，那么直接返回链表头即可。这种情况时间复杂度是常数级的。最坏的情况是遍历bins数组中所有的bin。对于一般大小的内存来说，bins数组元素非常少，比如4G空间只需要23个bin就足够了（256*2^23>4G），因此也很快能返回结果。总体来说search操作是非常高效的。对于固定大小内存来说，查找时间是常数量级的。

insert

将一个空闲的chunk插入到一个bin所挂载的chunk链表中，同时需要维持chunk链表的升序关系。具体流程是直接将chunk插入到index=chunksize/256>>2的那个bin中即可。

delete

将一个空闲的chunk从bins中移除。

4.总结

将内存分块管理，按块进行空间分配和释放。

通过split操作将大内存块分解成用户需要的小内存块。

通过merge操作合并小的内存块，做到内存碎片回收

通过bin这个抽象数据结构实现对空闲块高效管理。

5.代码分析

1.代码地址

https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/common_runtime

2.数据结构

Chunk

staticconstintkInvalidChunkHandle=-1;
...
structChunk{
size_tsize=0;//Fullsizeofbuffer.

//Wesometimesgivechunksthatarelargerthanneededtoreduce
//fragmentation.requested_sizekeepstrackofwhattheclient
//actuallywantedsowecanunderstandwhetheroursplitting
//strategyisefficient.
size_trequested_size=0;

//allocation_idissetto-1whenthechunkisnotinuse.Itisassigneda
//valuegreaterthanzerobeforethechunkisreturnedfrom
//AllocateRaw,andthisvalueisuniqueamongvaluesassignedby
//theparentallocator.
int64allocation_id=-1;
void*ptr=nullptr;//pointertograntedsubbuffer.

//IfnotkInvalidChunkHandle,thememoryreferredtoby'prev'isdirectly
//precedingthememoryusedbythischunk.E.g.,Itshouldstart
//at'ptr-prev->size'
ChunkHandleprev=kInvalidChunkHandle;

//IfnotkInvalidChunkHandle,thememoryreferredtoby'next'isdirectly
//followingthememoryusedbythischunk.E.g.,Itshouldbeat
//'ptr+size'
ChunkHandlenext=kInvalidChunkHandle;

//Whatbinarewein?
BinNumbin_num=kInvalidBinNum;

boolin_use()const{returnallocation_id!=-1;}
};

Bin

//ABinisacollectionofsimilar-sizedfreechunks.
structBin{
//Allchunksinthisbinhave>=bin_sizememory.
size_tbin_size=0;

structChunkComparator{
explicitChunkComparator(BFCAllocator*allocator)
:allocator_(allocator){}
//Sortfirstbysizeandthenusepointeraddressasatiebreaker.
booloperator()(constChunkHandleha,
constChunkHandlehb)constNO_THREAD_SAFETY_ANALYSIS{
constChunk*a=allocator_->ChunkFromHandle(ha);
constChunk*b=allocator_->ChunkFromHandle(hb);
if(a->size!=b->size){
returna->sizesize;
}
returna->ptrptr;
}

private:
BFCAllocator*allocator_;//Theparentallocator
};

typedefstd::setFreeChunkSet;
//Listoffreechunkswithinthebin,sortedbychunksize.
//Chunk*notowned.
FreeChunkSetfree_chunks;
Bin(BFCAllocator*allocator,size_tbs)
:bin_size(bs),free_chunks(ChunkComparator(allocator)){}
};

AllocationRegion

AllocationRegion给一个连续的内存区域做指针到ChunkHandle的映射。

RegionManager

RegionManager聚集了一个或多个AllocationRegion，并提供一个从指针到基础ChunkHandle的间接层，这个间接层可在多个不连续的内存区域进行分配。

3.分配大小

将每次分配的内存大小调整为kMinAllocationSize的N倍，这样所有内存地址都是很好地按字节对齐了。

//kMinAllocationSize=256
staticconstsize_tkMinAllocationBits=8;
staticconstsize_tkMinAllocationSize=1<
4.初始化bin
typedefintBinNum;
staticconstintkInvalidBinNum=-1;
staticconstintkNumBins=21;
...
//二进制2^8往左移0，1，2位
//(static_cast(256)<<0)=256
//(static_cast(256)<<1)=512
//(static_cast(256)<<2)=1024
size_tBinNumToSize(BinNumindex){
returnstatic_cast(256)<(&(bins_space_[index*sizeof(Bin)]));
}
...
//Wecreatebinstofitallpossiblerangesthatcoverthe
//memory_limit_startingfromallocationsupto256bytesto
//allocationsupto(andincluding)thememorylimit.
for(BinNumb=0;b
5.查找bin
//求属于第几个bin
BinNumBinNumForSize(size_tbytes){
uint64v=std::max(bytes,256)>>kMinAllocationBits;
intb=std::min(kNumBins-1,Log2FloorNonZero(v));
returnb;
}
//最高位非零的二进制位数，eg:00010101B为5
inlineintLog2FloorNonZero(uint64n){
#ifdefined(__GNUC__)
return63^__builtin_clzll(n);
#elifdefined(PLATFORM_WINDOWS)
unsignedlongindex;
_BitScanReverse64(&index,n);
returnindex;
#else
intr=0;
while(n>0){
r++;
n>>=1;
}
returnr;
#endif
}
6.查找Chunk
//先加锁
mutex_lockl(lock_);
void*ptr=FindChunkPtr(bin_num,rounded_bytes,num_bytes);
if(ptr!=nullptr){
returnptr;
}
//FindChunkPtr函数内部
void*BFCAllocator::FindChunkPtr(BinNumbin_num,size_trounded_bytes,
size_tnum_bytes){
//Firstidentifythefirstbinthatcouldsatisfyrounded_bytes.
for(;bin_numfree_chunks.begin();citer!=b->free_chunks.end();
++citer){
//从之前得到的Bin索引开始，查找合适的空闲Chunk：
constBFCAllocator::ChunkHandleh=(*citer);
BFCAllocator::Chunk*chunk=ChunkFromHandle(h);
DCHECK(!chunk->in_use());
if(chunk->size>=rounded_bytes){
//Wefoundanexistingchunkthatfitsusthatwasn'tinuse,soremove
//itfromthefreebinstructurepriortousing.
RemoveFreeChunkIterFromBin(&b->free_chunks,citer);

//Ifwecanbreakthesizeofthechunkintotworeasonably
//largepieces,doso.
//
//TODO(vrv):Whatshouldbethecriteriawhendecidingwhen
//tosplit?
//具体实现后面会分析
if(chunk->size>=rounded_bytes*2){
SplitChunk(h,rounded_bytes);
chunk=ChunkFromHandle(h);//Updatechunkpointerincaseitmoved
}

//Therequestedsizeofthereturnedchunkiswhattheuser
//hasallocated.
chunk->requested_size=num_bytes;
//Assignauniqueidandincrementtheidcounter,markingthe
//chunkasbeinginuse.
chunk->allocation_id=next_allocation_id_++;

//Updatestats.
++stats_.num_allocs;
stats_.bytes_in_use+=chunk->size;
stats_.max_bytes_in_use=
std::max(stats_.max_bytes_in_use,stats_.bytes_in_use);
stats_.max_alloc_size=
std::max(stats_.max_alloc_size,chunk->size);

VLOG(4)<<"Returning:"<ptr;
if(VLOG_IS_ON(4)){
LOG(INFO)<<"A:"<ptr;
}
}
}
returnnullptr;
}

7.拆分Chunk
如果Chunk的大小大于等于申请内存大小的2倍，那么将该Chunk拆分成2个：第一个Chunk的大小等于申请内存大小，第二个Chunk作为它的直接后继。
if(chunk->size>=rounded_bytes*2){
SplitChunk(h,rounded_bytes);
chunk=ChunkFromHandle(h);//Updatechunkpointerincaseitmoved
}

voidBFCAllocator::SplitChunk(BFCAllocator::ChunkHandleh,size_tnum_bytes){
//AllocatethenewchunkbeforewedoanyChunkFromHandle
ChunkHandleh_new_chunk=AllocateChunk();

Chunk*c=ChunkFromHandle(h);
CHECK(!c->in_use()&&(c->bin_num==kInvalidBinNum));

//Createanewchunkstartingnum_bytesafterc
BFCAllocator::Chunk*new_chunk=ChunkFromHandle(h_new_chunk);
new_chunk->ptr=static_cast(static_cast(c->ptr)+num_bytes);
region_manager_.set_handle(new_chunk->ptr,h_new_chunk);

//Setthenewsizesofthechunks.
new_chunk->size=c->size-num_bytes;
c->size=num_bytes;

//Thenewchunkisnotinuse.
new_chunk->allocation_id=-1;

//Maintainthepointers.
//c<->c_neighborbecomes
//c<->new_chunk<->c_neighbor
BFCAllocator::ChunkHandleh_neighbor=c->next;
new_chunk->prev=h;
new_chunk->next=h_neighbor;
c->next=h_new_chunk;
if(h_neighbor!=kInvalidChunkHandle){
Chunk*c_neighbor=ChunkFromHandle(h_neighbor);
c_neighbor->prev=h_new_chunk;
}

//Addthenewlyfreechunktothefreebin.
InsertFreeChunkIntoBin(h_new_chunk);
}

8.回收chunk
加锁，获得ChunkHandle
mutex_lockl(lock_);
BFCAllocator::ChunkHandleh=region_manager_.get_handle(ptr);
FreeAndMaybeCoalesce(h);
FreeAndMaybeCoalesce
voidBFCAllocator::FreeAndMaybeCoalesce(BFCAllocator::ChunkHandleh){
Chunk*c=ChunkFromHandle(h);
CHECK(c->in_use()&&(c->bin_num==kInvalidBinNum));

//Markthechunkasnolongerinuse
c->allocation_id=-1;

//Updatesthestats.
stats_.bytes_in_use-=c->size;

//Thischunkisnolongerin-use,considercoalescingthechunk
//withadjacentchunks.
ChunkHandlechunk_to_reassign=h;

//Ifthenextchunkisfree,coalescethetwo
if(c->next!=kInvalidChunkHandle){
Chunk*cnext=ChunkFromHandle(c->next);
if(!cnext->in_use()){
//VLOG(8)<<"Chunkat"<ptr<<"mergingwithc"<<
//c->ptr;

chunk_to_reassign=h;

//Deletesc->next
RemoveFreeChunkFromBin(c->next);
Merge(h,ChunkFromHandle(h)->next);
}
}

//Ifthepreviouschunkisfree,coalescethetwo
c=ChunkFromHandle(h);
if(c->prev!=kInvalidChunkHandle){
Chunk*cprev=ChunkFromHandle(c->prev);
if(!cprev->in_use()){
//VLOG(8)<<"Chunkat"<ptr<<"mergingintoc->prev"
//<ptr;

chunk_to_reassign=c->prev;

//Deletesc
RemoveFreeChunkFromBin(c->prev);
Merge(ChunkFromHandle(h)->prev,h);
c=ChunkFromHandle(h);
}
}

InsertFreeChunkIntoBin(chunk_to_reassign);
}

Merge
//Mergesh1andh2whenChunk(h1)->nextish2andChunk(h2)->previsc1.
//WemergeChunk(h2)intoChunk(h1).
voidBFCAllocator::Merge(BFCAllocator::ChunkHandleh1,
BFCAllocator::ChunkHandleh2){
Chunk*c1=ChunkFromHandle(h1);
Chunk*c2=ChunkFromHandle(h2);
//Wecanonlymergechunksthatarenotinuse.
CHECK(!c1->in_use()&&!c2->in_use());

//c1'sprevdoesn'tchange,stillpointstothesameptr,andis
//stillnotinuse.

//Fixupneighborpointers
//
//c1<->c2<->c3shouldbecome
//c1<->c3

BFCAllocator::ChunkHandleh3=c2->next;
c1->next=h3;
CHECK(c2->prev==h1);
if(h3!=kInvalidChunkHandle){
BFCAllocator::Chunk*c3=ChunkFromHandle(h3);
c3->prev=h1;
}

//Setthenewsize
c1->size+=c2->size;

DeleteChunk(h2);
}

以上这篇TensorFlow内存管理bfc算法实例就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。
声明：本文内容来源于网络，版权归原作者所有，内容由互联网用户自发贡献自行上传，本网站不拥有所有权，未作人工编辑处理，也不承担相关法律责任。如果您发现有涉嫌版权的内容，欢迎发送邮件至：czq8825#qq.com（发邮件时，请将#更换为@）进行举报，并提供相关证据，一经查实，本站将立刻删除涉嫌侵权内容。

TensorFlow内存管理bfc算法实例

热门推荐

随机推荐