[图解tensorflow源码] [转载] tensorflow设备内存分配算法解析 （BFC算法）

转载自 http://weibo.com/p/1001603980563068394770   @ICT_吴林阳

tensorflow设备内存管理模块实现了一个best-fit with coalescing算法（后文简称bfc算法）。bfc算法是Doung Lea’s malloc(dlmalloc)的一个非常简单的版本。它具有内存分配、释放、碎片管理等基本功能。

关于dlmalloc算法，参考下面链接：

http://gee.cs.oswego.edu/

Bfc算法思想：

将内存分成一系列内存块，每个内存块由一个chunk数据结构管理。从chunk结构中可以获取到内存块的使用状态、大小、数据的基址、前驱和后继chunk等信息。整个内存可以通过一个chunk的双链表结构来表示。内存分块结构参考下图：

用户申请一个内存块（malloc）。根据建立的chunk双链表找到一个合适的内存块（后面会说明什么是合适的内存块），如果该内存块的大小是用户申请大小的两倍以上，那么将该内存块切分成两块，这就是split操作。返回其中一块给用户，并将该内存块标识为占用。Spilt操作会新增一个chunk，所以需要修改chunk双链表以维持前驱和后继关系。下面给出一个例子。如下图所示，用户申请512的空间，正好有一块1024的chunk2是空闲的，由于1024/512 =2，所以chunk2 被split为2块：chunk2_1和chunk2_2。返回chunk2_1给用户并将其标志位占用状态。

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

到这里bfc的基本思想介绍的差不多了。其核心思想是：

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

2.      通过split操作将大内存块分解成小内存块。

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

但是还留下许多疑问。比如说申请内存空间时，什么样的块算合适的内存块？如何快速管理这种块？下面继续解释。

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

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

图中每一列表示一个bin，列首方格中的数字表示bin的size。bin size的大小都是256的2^n的倍。每个bin下面挂载了一系列的空闲chunk，每个chunk的chunk size都大于等于所属的bin的bin size，按照chunk size的升序挂载成单链表。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中移除。

有了bin这个抽象数据结构，我们再来看看下面两个问题：

1.      申请空间时，什么是合适大小的内存块？

2.      假设用户需要Malloc一块大小为x的内存空间时，如何从bins中找到一个合适大小的内存块给用户？

我们首先看看search操作和malloc之间的关系。bfc算法本身的设计决定了chunksize通常都是256的2^n倍，而x是用户随意指定的，所以通常情况下不会找到一个与x正好匹配的chunk size。在以块为单位返回申请空间的前提下，我们希望实际分配的空间大于等于x（当然最好是等于）。如果实际分配的空间大于x，我们进一步希望多分配的那部分空间不会太大，以减少内存浪费。所以bfc选择合适内存块的原则是：找到chunk size大于等于x的最小的那个空闲内存块，这就是合适大小的内存块。我们再回头看看search操作的介绍，可以发现通过search操作就能从bins中找到一个合适大小的内存块给用户。这就回答了上面两个问题。

我们再来看看insert操作和split之间的联系。当内存中的空闲块都很大时，即使是chunk size大于等于x的最小的那个空闲内存块的大小也可能远远大于x。这时我们就需要一些更小的chunk了。于是bfc算法设计了split这一操作。假设x落在某两个邻近的bin size区间[a，b)上（a和b能够很容易通过x计算出来，这里b = 2 * a）。如果通过search找到的内存块大小大于等于b（在 x = a的情况下是a）的两倍，将其split成两份。其中一份大小为b（在 x = a的情况下是a），这一份返回给用户。另外一份作为一个新的空闲chunk插入到bins中合适的位置上，这里就会用到Insert操作。

至此bfc算法的整体思路解析完毕。回过头来总结一下其核心思想如下：

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

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

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

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