tensorflow源码解析之framework-allocator
目录
- 什么是allocator
- 内存分配器的管理
- 内存分配追踪
- 其它结构
- 关系图
- 涉及的文件
- 迭代记录
1. 什么是allocator
Allocator是所有内存分配器的基类,它定义了内存分配器需要实现的接口。
class Allocator {
public:
//内存分配与返还
virtual void* AllocateRaw(size_t alignment, size_t num_bytes) = 0;
virtual void DeallocateRaw(void* ptr) = 0;
T* Allocate(size_t num_elements);
T* Allocate(size_t num_elements, const AllocationAttributes& allocation_attr);
void Deallocate(T* ptr, size_t num_elements);
//追踪内存分配信息
virtual bool TracksAllocationSizes();
virtual bool ShouldAllocateEmptyTensors();
virtual size_t RequestedSize(void* ptr);
virtual size_t AllocatedSize(void* ptr);
virtual int64 AllocationId(void* ptr);//本次内存分配的编号
virtual size_t AllocatedSizeSlow(void* ptr);
virtual void GetStats(AllocatorStats* stats);
};
这些API可以被分为两类,一是内存分配与返还,二是内存分配信息追踪。前者是内存分配器的本职工作,后者提供了对内存分配器所分配的内存进行追踪和管理的功能。若想实现后者的功能,需要提供一定的数据结构支持,下文中的“内存分配追踪”会详细讲到。
另外还遇到了一个新的结构AllocatorStats,这个类是对内存分配器当前分配内存量的一个宏观统计,它的定义如下:
struct AllocatorStats {
int64 num_allocs;//内存分配次数
int64 bytes_in_use;//分配的内存中,当前正在使用的大小
int64 max_bytes_in_use;//使用中的内存大小的峰值
int64 max_alloc_size;//最大的单次内存分配大小
int64 bytes_limit;//当前内存分配器能分配的最大内存量,如果申请内存大小超过这个阈值,返回0
//...
}
Allocator除了提供内存申请的接口之外,还提供了为申请好的内存调用默认构造和析构函数的接口。如果在申请的时候指定了对象的类型,就可以选择调用对象所属类的构造和析构方法。Allocator提供了针对三种常用类的构造方法,分别是String,ResourceHandle和Variant。
class Allocator {
public:
//...
private:
void RunCtor(T* p, size_t n);
virtual void RunStringCtor(string* p, size_t n);
virtual void RunStringDtor(string* p, size_t n);
virtual void RunResourceCtor(ResourceHandle* p, size_t n);
virtual void RunResourceDtor(ResourceHandle* p, size_t n);
virtual void RunVariantCtor(Variant* p, size_t n);
virtual void RunVariantDtor(Variant* p, size_t n);
};
除了抽象的内存分配器接口之外,TF还为最常用的CPU内存分配器,提供了一个默认实现:
class CPUAllocator : public Allocator {
public:
//...
void GetStats(AllocatorStats* stats) override {
mutex_lock l(mu_);
*stats = stats_;
}
size_t AllocatedSizeSlow(void *ptr) override {
return port::MallocExtension_GetAllocatedSize(ptr);
}
private:
mutex mu_;
AllocatorStats stats_ GUARDED_BY(mu_);
};
需要注意两点,第一,它为内存分配宏观统计加入了一个数据成员stats_,在需要时直接将其返回,第二,为追踪内存分配大小提供了一个“慢”版本的实现,这个实现的意义在于,如果需要知道某个指针对应的分配内存的大小,而我们又没有专门为其准备这样的数据记录时,可以直接调用操作系统层面的函数,来获取指针对应的分配内存的大小。这个操作会相对耗时,但聊胜于无。
2. 内存分配器的管理
不同类型的设备,可能需要不同的内存分配器,即便对于相同类型的设备,考虑到效率问题,也可能会提供不同的内存分配器版本。因此,需要一个对内存分配器进行集中管理的地方,TF为我们提供了类AllocatorRegistry:
class AllocatorRegistry {
public:
//内存分配器注册
void Register(const string& name, int priority, Allocator* allocator);
//返回最高优先级的内存分配器
Allocator* GetAllocator();
//返回一个全局的内存分配器的注册器
static AllocatorRegistry* Global();
private:
//内存分配器存储位置
std::vector<AllocatorRegistryEntry> allocators_;
//...
}
所有的内存分配器都被保存在allocators_这个向量里,但向量中保存的并不是内存分配器本身,而是对它的一个封装,我们看下这个封装的结构:
typedef struct {
string name;
int priority;
Allocator* allocator;
} AllocatorRegistryEntry;
除了内存分配器之外,这个entry里还存放了内存分配器的名称和优先级。当向AllocatorRegistry请求一个内存分配器时,它返回的是具有最高优先级的分配器,如果多个分配器有相同的优先级,就返回其中的一个。
AllocatorRegistry实际上是一个单例对象,它的Global接口返回一个全局静态的注册器对象。为了方便进行注册,TF还设计了一个统一的注册入口类:
class AllocatorRegistration {
public:
AllocatorRegistration(const string& name, int priority, Allocator* allocator){
AllocatorRegistry::Global()->Register(name,priority,allocator);
}
};
另外,TF还设计了一个宏来简化注册过程,感兴趣的读者可以参考源代码。
3. 内存分配追踪
刚才提到,在内存分配器的公共API中,有一类专门用于追踪内存分配,需要有一些专用的数据结构,来保存每一次内存分配的信息。这些工作被TrackingAllocator类实现:
class TrackingAllocator : public Allocator {
public:
std::tuple<size_t, size_t, size_t> GetSizeAndUnRef();
//...
private:
bool UnRef() EXCLUSIZE_LOCKS_REQUIRED(mu_);
Allocator* allocator_;
mutex mu_;
int ref_ GUARDED_BY(mu_);
//当前仍在使用的分配内存大小,如果allocator_不支持内存追踪,则为0
size_t allocated_ GUARDED_BY(mu_);
//allocated_的峰值
size_t high_watermark_ GUARDED_BY(mu_);
//当前内存分配器总共分配的内存大小
size_t total_bytes_ GUARDED_BY(mu_);
const bool track_sizes_locally_;
struct Chunk {
size_t requested_size;
size_t allocated_size;
int64 allocation_id;
};
std::unordered_map<void*, Chunk> in_use_ GUARDED_BY(mu_);
};
除了提供allocated_,high_watermark_,total_bytes_三个数据成员记录内存分配的统计信息之外,更重要的是加入了in_use_这个数据成员。它从一个指针映射到一个Chunk,而这个Chunk中保存了每次内存分配需要的内存大小、实际分配的内存大小、本次内存分配的唯一标识。用这样一个结构保存了每次内存分配的详细信息。
再细说一下ref_成员的作用。这个数据成员存在的意义在于,保存当前内存分配器分配内存的次数,当所有分配的内存全部回收之后,就删除掉当前的内存分配器。但你会发现,这个ref_成员在TrackingAllocator对象初始化的时候,本身已经赋值为1了,那即便我们每次返回内存时都调用UnRef函数将它减一,最终不还是不能为0吗?原因在于,我们希望对于GetSizeAndUnRef这个函数在对象的生命周期内只调用一次,因此这个函数调用后会将ref_减一。而不论是GetSizeAndUnRef函数还是调用UnRef,只要ref_值减到0,就删除这个对象。这要求我们必须仅调用GetSizeAndUnRef函数一次,否则就会出现内存泄漏。
4. 其它结构
在TF中,计算是发生在节点上的,而节点被分配在具体的设备上。对于一个GPU设备,在它上面运行的节点,是不是就不需要CPU内存了呢?显然不是,比如,为了使用DMA给某些设备传送数据,运行在GPU上的节点仍然需要申请CPU内存。因此,当节点向一个设备索要内存分配器时,需要给它提供一些信息,告诉设备我们想要申请哪种类型的内存,这些信息就存储在AllocatorAttributes类中。
struct AllocatorAttributes {
void set_on_host(bool v);
bool on_host() const;
void set_nic_compatible(bool v);
bool nic_compatible() const;
void set_gpu_compatible(bool v);
bool gpu_compatible() const;
void set_track_sizes(bool v);
bool track_sizes() const;
void Merge(AllocatorAttributes other);
bool IsEqualOrLessRestrictiveThan(const AllocatorAttributes& other);
uint32 value = 0;//这个数值的高8位被保留为设备相关的设置。各设备的实现可以根据需要自行解析,为这些设备实现的操作也需要正确的解析它
}
AllocatorAttributes很容易与另外一个类混淆,那就是AllocationAttributes。后者记录的是为内存分配器的某一次具体的内存分配的属性信息,使用时机完全不一样。
class AllocationAttributes {
bool no_retry_on_failure = false; //如果首次内存分配失败了,不再尝试。
bool allocation_will_be_logged = false;//本次内存分配是否会被记录
}
另外,有时候我们想对某个内存分配器进行封装,以便在某个API上实现定制化。TF为此准备了类AllocatorWrapper类,它本质上就是对Allocator类的直接封装,感兴趣的读者可以去看下源码。
5. 关系图
A(Allocator)-->|派生|B(CPUAllocator)
A(Allocator)-->|派生|C(AllocatorWrapper)
A(Allocator)-->|派生|D(TrackingAllocaotr)
E(AllocatorRegistration)-.封装.->F(AllocatorRegistry)
A(Allocator)-.注册.->F(AllocatorRegistration)
6. 涉及的文件
- allocator
- allocator_registry
- tracking_allocator
7. 迭代记录
- v1.0 2018-08-25 文档创建
- v2.0 2018-09-08 文档重构
tensorflow源码解析之framework-allocator的更多相关文章
- tensorflow源码解析之framework拾遗
把framework中剩余的内容,按照文件名进行了简单解析.时间原因写的很仓促,算是占个坑,后面有了新的理解再来补充. allocation_description.proto 一个对单次内存分配结果 ...
- tensorflow源码解析系列文章索引
文章索引 framework解析 resource allocator tensor op node kernel graph device function shape_inference 拾遗 c ...
- Tensorflow源码解析1 -- 内核架构和源码结构
1 主流深度学习框架对比 当今的软件开发基本都是分层化和模块化的,应用层开发会基于框架层.比如开发Linux Driver会基于Linux kernel,开发Android app会基于Android ...
- tensorflow源码解析之common_runtime-executor-上
目录 核心概念 executor.h Executor NewLocalExecutor ExecutorBarrier executor.cc structs GraphView ExecutorI ...
- tensorflow源码解析之common_runtime-executor-下
目录 核心概念 executor.h Executor NewLocalExecutor ExecutorBarrier executor.cc structs GraphView ExecutorI ...
- tensorflow源码解析之common_runtime拾遗
把common_runtime中剩余的内容,按照文件名排序进行了简单的解析,时间原因写的很仓促,算是占个坑,后续有了新的理解再来补充. allocator_retry 有时候内存分配不可能一次完成,为 ...
- tensorflow源码解析之distributed_runtime
本篇主要介绍TF的分布式运行时的基本概念.为了对TF的分布式运行机制有一个大致的了解,我们先结合/tensorflow/core/protobuf中的文件给出对TF分布式集群的初步理解,然后介绍/te ...
- tensorflow源码解析之framework-tensor
目录 什么是tensor tensor继承体系 与Eigen3库的关系 什么是tensor_reference tensor_shape tensor_slice 其它结构 关系图 涉及的文件 迭代记 ...
- tensorflow源码解析之framework-op
目录 什么是op op_def定义 op注册 op构建与注册辅助结构 op重写 关系图 涉及的文件 迭代记录 1. 什么是op op和kernel是TF框架中最重要的两个概念,如果一定要做一个类比的话 ...
随机推荐
- SqlServer基础语法
历史 有很多软件公司开发了数据库产品,其中微软公司的数据库产品命名为 SQL Server,也称 MS SQL Server. 1989年 Ashton-Tate/Microsoft SQL Serv ...
- git 初始化本地项目并推送到远程
有一个新项目,开发了一些代码之后想推送到远程,具体的操作方式和命令如下: (使用 git bash) 1.切到项目目录中,例如 E:\git\smart-open 2.初始化git仓库并在本地提交 / ...
- 学习jsp篇:jsp简单实例之一注册
编程环境:IDEA,Tomcat ,JavaEE 实例一.注册 1.先在IDEA建一个web工程(不懂的可以在网上搜,一大堆..)ServletTest,在工程目录下的web目录建一个文件夹regis ...
- Docker-Compose 容器集群的编排
Docker-Compose 容器集群的编排 1.Docker-Compose 概述 2.Docker Compose 环境安装 Docker-Compose功能: 可以实现单机上的容器集群编排 1 ...
- xss标签和属性爆破
当网站过滤了大部分的HTML标签和属性,我们就尝试爆破一下,看哪些标签和属性没有没过滤. 爆破的步骤: 1. 首先在测试点输入我们正常的exp,并抓包发送到Intruder模块. 2. 将exp改为 ...
- 测试提高路线图_tester+
https://mp.weixin.qq.com/s/30ZT0w164Q3iLdPg4R8org
- 04 前端之BOM与DOM
目录 前端之BOM与DOM BOM与DOM操作 BOM操作 前端之BOM与DOM BOM与DOM操作 BOM 浏览器对象模型>>>:使用js操作浏览器 DOM 文档对象模型>& ...
- Linux基础:操作系统的启动
Centos6: # 1.加电自检(BIOS)# 2.MBR引导(512k)dd </dev/zero >/dev/sda bs=1k count=400 # 3.GRUB菜单(选择系统) ...
- [源码解析] NVIDIA HugeCTR,GPU版本参数服务器---(3)
[源码解析] NVIDIA HugeCTR,GPU版本参数服务器---(3) 目录 [源码解析] NVIDIA HugeCTR,GPU版本参数服务器---(3) 0x00 摘要 0x01 回顾 0x0 ...
- Latex公式导出word,Latex转换MathML使用POI导出公式可编辑的Word文件
背景 之前在 使用spire.doc导出支持编辑Latex公式的标准格式word 博客中写过,使用spire.doc来生成word,不得不说spire.doc的api操作起来还是比较方便,但是使用的过 ...