Tensorflow中的内存分配

Tensorflow中的内存分配

Tensorflow系统复杂，支持平台多，各类数据结构也多，所以设计一个统一的内存管理分配接口很重要。本文主要探讨tensorflow中的内存分配的相关机制，会重点研究其中实现的两种分配算法。

内存分配是系统中非常重要的一环，大家平常接触最多的就是malloc(new)和free(delete)。也就是分配和施放，在tensorflow中也是这样。Tensorflow提供了一个公共的接口类，Allocator类，该类提供了两个重要的方法：

AllocateRaw分配内存，DeallocateRaw施放内存。Tensorflow中提供的内存分配方法都是继承该类去实现。

PoolAllocator是一种实现了LRU策略的内存池分配算法。

        

PtrRecord是分配的数据块，在poolAlloctor包含一个双向链表和一个multimap类型成员pool_,pool_的key是分配的内存大小。这两个成员都保存了当前系统未分配的内存块。

poolAlloctor在施放内存时，接口只传入了地址，并没有传递需要施放内存的大小，所以在poolalloctor实现时，采用了cookie技术，即采用了下面的数据结构，来存放分配的

Poolalloctor每次向系统申请的内存时如上图所示，但是前面会多一个sizeof(chunkprefix)字节的数据，真正返回给用户使用的内存起始地址是user_ptr指向的地址。

再来看分配代码，第104行可以看出，分配的字节数numbytes已经加上了ChunkPrefix结构图的大小。

110行，tensorflow使用numbytes做为key,从pool寻找适合当前请求的PtrRecord。可以看出为了提高查找效率，tensorflow使用mutilmap数据结构。

117行表示找到了符合当前请求的内存块，119行表示需要从双向链表内删除这个块，120行表示从pool_中删除这个块。

125行表示找到了内存，127行删除的是pool_中找到的元素。128行对内存进行处理，加上ccookie信息。

130表示没找到需要的内存块，需要向系统分配一款内存。131行也是对刚从系统分配的内存进行处理。

再看施放函数，137行表示根据ptr，查找cookie 信息，得到真正分配的内存地址

139-140行表示，如果分配的内存无限制且不自动resize，直接对内存进行施放。

142-151行，将内存加入pool_内，以备下次申请使用。144行表示当前池的大小已经到达上限，需要施放。EviceOne函数会施放双向链表的尾部元素。这样符合lru策略的逻辑。

147-149行，申请PtrRecord节点。150行，将节点插入链表头，151行，将元素插入pool_中。

         接下来我们来看BFCAllocator内存管理算法。它包含了21个Bin，每个Bin下面包含一个chunk的集合，每一类bin下面的chunk的大小是一样的，bin的排序也是根据chunk的大小进行升序排列的。也就是说chunk的大小也是固定有21类。其中最小的chunk是256个字节，之后每一个chunk大小是256向右移一位，即512,1024,2048等。

BFCAllocator的数据结构类图如上所示，BFCAlloctor内部包使用vector管理chunk。Bin中的chunks的元素使用的是chunk在vector中的索引。

该函数主要是将用户需求的内存大小对齐到256整数倍。

BFCAllocator最大的特点内存块的分裂和合并。当查找到符合满足用户需求大小的chunk大小是需求的2倍时，会执行分裂操作，系统也会根据当前的内存碎片会执行合并操作。

先看分配函数，主要看232行实现的AllocateRawInternal函数。

函数的主体部分已经截全，367-370行，主要计算对齐大小，并根据该大小定位要查找的chunk所在的bin。

373-376行尝试合并带时间戳的chunk。

377-380行表示在bin中去查找满足分配要求的chunk。

383-388行表示当没有满足要求的chunk时，需要从系统中分配一块满足要求的新chunk。

390-400行表示会从带时间戳的chunk，尚未被合并的列表中去合并一块满足当前要求的chunk。

407-411表示向系统未被使用的内存，由系统继续合并，然后再重新申请并尝试分配。

FindChunkPtr查找chunk函数

432行表示从满足要求的chunk 大小所在的bin进行查找。

444行表示查找到了满足要求的chunk，447行需要将该chunk从bin中删除。

453-458对chunk进行分裂，分裂的条件是chunk的大小是需求大小的2倍或者chunk大小和需求大小的差值大于128mb。

457行返回新chunk。

SplitChunk分裂函数。

495行建立新的chunk。注意新chunk的起始地址是原 chunk地址加上需求的大小

497行保存新chunk的地址和索引，这里是为了将来合并chunk时使用。注意region_manager成员变量。他主要保存chunk中的内存块地址和chunk索引的映射关系。

512-519行，对分配后的chunk进行串联处理。这一步也是为了将来合并chunk时使用。

TryToCoalesce对chunk块进行合并，合并的是之前分裂过的块。

650-656行查找chunk 的next块，如果next是空闲，则对齐进行合并。

660-667行查找chunk的prev块，如果prev是空闲，则对齐进行合并。