1. 基本介绍
- tensorflow 设备内存管理模块实现了一个 best-fit with coalescing 算法（后文简称 bfc 算法）。
- bfc 算法是 Doung Lea’s malloc(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 操作，合并当前块和其前后的空闲块）。再修改双链表结构以维持前驱后继关系。这就做到了内存碎片的回收。
- 如果用户要 free chunk3，由于 chunk3 的前驱 chunk2 也是空闲的，所以将 chunk2 和 chunk3 合并得到一个新的 chunk2’，大小为 chunk2 和 chunk3 之和。
3. bins
1. bins 数据结构
- bfc 算法采取的是被动分块的策略。最开始整个内存是一个 chunk，随着用户申请空间的次数增加，最开始的大 chunk 会被不断的 split 开来，从而产生越来越多的小 chunk。当 chunk 数量很大时，为了寻找一个合适的内存块而遍历双链表无疑是一笔巨大的开销。为了实现对空闲块的高效管理，bfc 算法设计了 bin 这个抽象数据结构。
- 每个 bin 都有一个 size 属性，一个 bin 是一个拥有 chunk size >= binsize 的空闲 chunk 的集合。集合中的 chunk 按照 chunk size 的升序组织成单链表。bfc 算法维护了一个 bin 的集合：bins。它由多个 bin 以及从属于每个 bin 的 chunks 组成。内存中所有的空闲 chunk 都由 bins 管理。
- 图中每一列表示一个 bin，列首方格中的数字表示 bin 的 size。bin size 的大小都是 256 的 2^n 的倍。每个 bin 下面挂载了一系列的空闲 chunk，每个 chunk 的 chunk size 都大于等于所属的 bin 的 bin size，按照 chunk size 的升序挂载成单链表。
2. bins 操作
- bfc 算法针对 bins 这个集合设计了三个操作：search、insert、delete。
- search
给定一个 chunk size，从 bins 中找到大于等于该 chunksize 的最小的那个空闲 chunk。Search 操作具体流程如下。如果 bin 以数组的形式组织，那么可以从 index = chunk size /256 >>2 的那个 bin 开始查找。最好的情况是开始查找的那个 bin 的 chunk 链表非空，那么直接返回链表头即可。这种情况时间复杂度是常数级的。最坏的情况是遍历 bins 数组中所有的 bin。对于一般大小的内存来说，bins 数组元素非常少，比如 4G 空间只需要 23 个 bin 就足够了（256 * 2 ^ 23 > 4G），因此也很快能返回结果。总体来说 search 操作是非常高效的。对于固定大小内存来说，查找时间是常数量级的。 - insert
将一个空闲的 chunk 插入到一个 bin 所挂载的 chunk 链表中，同时需要维持 chunk 链表的升序关系。具体流程是直接将 chunk 插入到 index = chunk size /256 >>2 的那个 bin 中即可。 - delete
将一个空闲的 chunk 从 bins 中移除。
- search
4. 总结
- 将内存分块管理，按块进行空间分配和释放。
- 通过 split 操作将大内存块分解成用户需要的小内存块。
- 通过 merge 操作合并小的内存块，做到内存碎片回收
- 通过 bin 这个抽象数据结构实现对空闲块高效管理。
5. 代码分析
1. 代码地址
2. 数据结构
- Chunk
static const int kInvalidChunkHandle = -1;...struct Chunk {size_t size = 0; // Full size of buffer.// We sometimes give chunks that are larger than needed to reduce// fragmentation. requested_size keeps track of what the client// actually wanted so we can understand whether our splitting// strategy is efficient.size_t requested_size = 0;// allocation_id is set to -1 when the chunk is not in use. It is assigned a// value greater than zero before the chunk is returned from// AllocateRaw, and this value is unique among values assigned by// the parent allocator.int64 allocation_id = -1;void* ptr = nullptr; // pointer to granted subbuffer.// If not kInvalidChunkHandle, the memory referred to by 'prev' is directly// preceding the memory used by this chunk. E.g., It should start// at 'ptr - prev->size'ChunkHandle prev = kInvalidChunkHandle;// If not kInvalidChunkHandle, the memory referred to by 'next' is directly// following the memory used by this chunk. E.g., It should be at// 'ptr + size'ChunkHandle next = kInvalidChunkHandle;// What bin are we in?BinNum bin_num = kInvalidBinNum;bool in_use() const { return allocation_id != -1; }};
- Bin
// A Bin is a collection of similar-sized free chunks.struct Bin {// All chunks in this bin have >= bin_size memory.size_t bin_size = 0;struct ChunkComparator {explicit ChunkComparator(BFCAllocator* allocator): allocator_(allocator) {}// Sort first by size and then use pointer address as a tie breaker.bool operator()(const ChunkHandle ha,const ChunkHandle hb) const NO_THREAD_SAFETY_ANALYSIS {const Chunk* a = allocator_->ChunkFromHandle(ha);const Chunk* b = allocator_->ChunkFromHandle(hb);if (a->size != b->size) {return a->size < b->size;}return a->ptr < b->ptr;}private:BFCAllocator* allocator_; // The parent allocator};typedef std::set<ChunkHandle, ChunkComparator> FreeChunkSet;// List of free chunks within the bin, sorted by chunk size.// Chunk * not owned.FreeChunkSet free_chunks;Bin(BFCAllocator* allocator, size_t bs): bin_size(bs), free_chunks(ChunkComparator(allocator)) {}};
- AllocationRegion
- AllocationRegion 给一个连续的内存区域做指针到 ChunkHandle 的映射。
- RegionManager
- RegionManager 聚集了一个或多个 AllocationRegion，并提供一个从指针到基础 ChunkHandle 的间接层，这个间接层可在多个不连续的内存区域进行分配。
3. 分配大小
- 将每次分配的内存大小调整为 kMinAllocationSize 的 N 倍，这样所有内存地址都是很好地按字节对齐了。
// kMinAllocationSize = 256static const size_t kMinAllocationBits = 8;static const size_t kMinAllocationSize = 1 << kMinAllocationBits;...size_t BFCAllocator::RoundedBytes(size_t bytes) {size_t rounded_bytes =(kMinAllocationSize *((bytes + kMinAllocationSize - 1) / kMinAllocationSize));DCHECK_EQ(size_t{0}, rounded_bytes % kMinAllocationSize);return rounded_bytes;}
4. 初始化 bin
typedef int BinNum;static const int kInvalidBinNum = -1;static const int kNumBins = 21;...// 二进制2^8往左移0，1，2位// (static_cast<size_t>(256) << 0) = 256// (static_cast<size_t>(256) << 1) = 512// (static_cast<size_t>(256) << 2) = 1024size_t BinNumToSize(BinNum index) {return static_cast<size_t>(256) << index;}...char bins_space_[sizeof(Bin) * kNumBins];// Map from bin size to BinBin* BinFromIndex(BinNum index) {return reinterpret_cast<Bin*>(&(bins_space_[index * sizeof(Bin)]));}...// We create bins to fit all possible ranges that cover the// memory_limit_ starting from allocations up to 256 bytes to// allocations up to (and including) the memory limit.for (BinNum b = 0; b < kNumBins; b++) {size_t bin_size = BinNumToSize(b);VLOG(1) << "Creating bin of max chunk size "<< strings::HumanReadableNumBytes(bin_size);new (BinFromIndex(b)) Bin(this, bin_size);CHECK_EQ(BinForSize(bin_size), BinFromIndex(b));CHECK_EQ(BinForSize(bin_size + 255), BinFromIndex(b));CHECK_EQ(BinForSize(bin_size * 2 - 1), BinFromIndex(b));if (b + 1 < kNumBins) {CHECK_NE(BinForSize(bin_size * 2), BinFromIndex(b));}}
5. 查找 bin
// 求属于第几个binBinNum BinNumForSize(size_t bytes) {uint64 v = std::max<size_t>(bytes, 256) >> kMinAllocationBits;int b = std::min(kNumBins - 1, Log2FloorNonZero(v));return b;}// 最高位非零的二进制位数，eg: 0001 0101B 为5inline int Log2FloorNonZero(uint64 n) {#if defined(__GNUC__)return 63 ^ __builtin_clzll(n);#elif defined(PLATFORM_WINDOWS)unsigned long index;_BitScanReverse64(&index, n);return index;#elseint r = 0;while (n > 0) {r++;n >>= 1;}return r;#endif}
6. 查找 Chunk
// 先加锁mutex_lock l(lock_);void* ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes);if (ptr != nullptr) {return ptr;}// FindChunkPtr函数内部void* BFCAllocator::FindChunkPtr(BinNum bin_num, size_t rounded_bytes,size_t num_bytes) {// First identify the first bin that could satisfy rounded_bytes.for (; bin_num < kNumBins; bin_num++) {// Start searching from the first bin for the smallest chunk that fits// rounded_bytes.Bin* b = BinFromIndex(bin_num);for (auto citer = b->free_chunks.begin(); citer != b->free_chunks.end();++citer) {// 从之前得到的Bin索引开始，查找合适的空闲Chunk：const BFCAllocator::ChunkHandle h = (*citer);BFCAllocator::Chunk* chunk = ChunkFromHandle(h);DCHECK(!chunk->in_use());if (chunk->size >= rounded_bytes) {// We found an existing chunk that fits us that wasn't in use, so remove// it from the free bin structure prior to using.RemoveFreeChunkIterFromBin(&b->free_chunks, citer);// If we can break the size of the chunk into two reasonably// large pieces, do so.//// TODO(vrv): What should be the criteria when deciding when// to split?// 具体实现后面会分析if (chunk->size >= rounded_bytes * 2) {SplitChunk(h, rounded_bytes);chunk = ChunkFromHandle(h); // Update chunk pointer in case it moved}// The requested size of the returned chunk is what the user// has allocated.chunk->requested_size = num_bytes;// Assign a unique id and increment the id counter, marking the// chunk as being in use.chunk->allocation_id = next_allocation_id_++;// Update stats.++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<std::size_t>(stats_.max_alloc_size, chunk->size);VLOG(4) << "Returning: " << chunk->ptr;if (VLOG_IS_ON(4)) {LOG(INFO) << "A: " << RenderOccupancy();}return chunk->ptr;}}}return nullptr;}
7. 拆分 Chunk
- 如果 Chunk 的大小大于等于申请内存大小的 2 倍，那么将该 Chunk 拆分成 2 个：第一个 Chunk 的大小等于申请内存大小，第二个 Chunk 作为它的直接后继。
if (chunk->size >= rounded_bytes * 2) {SplitChunk(h, rounded_bytes);chunk = ChunkFromHandle(h); // Update chunk pointer in case it moved}void BFCAllocator::SplitChunk(BFCAllocator::ChunkHandle h, size_t num_bytes) {// Allocate the new chunk before we do any ChunkFromHandleChunkHandle h_new_chunk = AllocateChunk();Chunk* c = ChunkFromHandle(h);CHECK(!c->in_use() && (c->bin_num == kInvalidBinNum));// Create a new chunk starting num_bytes after cBFCAllocator::Chunk* new_chunk = ChunkFromHandle(h_new_chunk);new_chunk->ptr = static_cast<void*>(static_cast<char*>(c->ptr) + num_bytes);region_manager_.set_handle(new_chunk->ptr, h_new_chunk);// Set the new sizes of the chunks.new_chunk->size = c->size - num_bytes;c->size = num_bytes;// The new chunk is not in use.new_chunk->allocation_id = -1;// Maintain the pointers.// c <-> c_neighbor becomes// c <-> new_chunk <-> c_neighborBFCAllocator::ChunkHandle h_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;}// Add the newly free chunk to the free bin.InsertFreeChunkIntoBin(h_new_chunk);}
8. 回收 chunk
加锁，获得 ChunkHandle
mutex_lock l(lock_);BFCAllocator::ChunkHandle h = region_manager_.get_handle(ptr);FreeAndMaybeCoalesce(h);
FreeAndMaybeCoalesce
void BFCAllocator::FreeAndMaybeCoalesce(BFCAllocator::ChunkHandle h) {Chunk* c = ChunkFromHandle(h);CHECK(c->in_use() && (c->bin_num == kInvalidBinNum));// Mark the chunk as no longer in usec->allocation_id = -1;// Updates the stats.stats_.bytes_in_use -= c->size;// This chunk is no longer in-use, consider coalescing the chunk// with adjacent chunks.ChunkHandle chunk_to_reassign = h;// If the next chunk is free, coalesce the twoif (c->next != kInvalidChunkHandle) {Chunk* cnext = ChunkFromHandle(c->next);if (!cnext->in_use()) {// VLOG(8) << "Chunk at " << cnext->ptr << " merging with c " <<// c->ptr;chunk_to_reassign = h;// Deletes c->nextRemoveFreeChunkFromBin(c->next);Merge(h, ChunkFromHandle(h)->next);}}// If the previous chunk is free, coalesce the twoc = ChunkFromHandle(h);if (c->prev != kInvalidChunkHandle) {Chunk* cprev = ChunkFromHandle(c->prev);if (!cprev->in_use()) {// VLOG(8) << "Chunk at " << c->ptr << " merging into c->prev "// << cprev->ptr;chunk_to_reassign = c->prev;// Deletes cRemoveFreeChunkFromBin(c->prev);Merge(ChunkFromHandle(h)->prev, h);c = ChunkFromHandle(h);}}InsertFreeChunkIntoBin(chunk_to_reassign);}
- Merge
// Merges h1 and h2 when Chunk(h1)->next is h2 and Chunk(h2)->prev is c1.// We merge Chunk(h2) into Chunk(h1).void BFCAllocator::Merge(BFCAllocator::ChunkHandle h1,BFCAllocator::ChunkHandle h2) {Chunk* c1 = ChunkFromHandle(h1);Chunk* c2 = ChunkFromHandle(h2);// We can only merge chunks that are not in use.CHECK(!c1->in_use() && !c2->in_use());// c1's prev doesn't change, still points to the same ptr, and is// still not in use.// Fix up neighbor pointers//// c1 <-> c2 <-> c3 should become// c1 <-> c3BFCAllocator::ChunkHandle h3 = c2->next;c1->next = h3;CHECK(c2->prev == h1);if (h3 != kInvalidChunkHandle) {BFCAllocator::Chunk* c3 = ChunkFromHandle(h3);c3->prev = h1;}// Set the new sizec1->size += c2->size;DeleteChunk(h2);}
参考资料：
- https://github.com/tensorflow/tensorflow
- http://www.cnblogs.com/yao62995/p/5773166.html
- https://deepreve.github.io/articles/20161007.html
原文：https://blog.csdn.net/qq_33096883/article/details/77479647
若有收获，就点个赞吧
0 人点赞
