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框架中最重要的两个概念,如果一定要做一个类比的话 ...
随机推荐
- STL priority_queue 优先队列 小记
今天做题发现一个很有趣的地方,竟然还是头一次发现,唉,还是太菜了. 做图论用STL里的priority_queue去优化prim,由于特殊需求,我需要记录生成树中是用的哪些边. 于是,我定义的优先队列 ...
- ByteArrayOutputStream内存流
简介 ByteArrayOutputStream 对byte类型数据进行写入的类 相当于一个中间缓冲层,创建ByteArrayOutputStream类实例时,内存中会创建一个byte数组类型的缓冲区 ...
- MySQL 1093 - You can't specify target table 'sc' for update in FROM clause
错误代码如下: #(8) 把"邓维杰"同学的成绩全部删除. SELECT * FROM sc WHERE EXISTS(SELECT * FROM student WHERE st ...
- php7.3编译安装 支持微擎2.0
再次整理 //一下配置在命令粘贴时注意句尾加 \ , 在 \ 后不能有空格,不然会自动执行,相当于回车./configure --prefix=/usr/local/php \ --with-co ...
- java中的成员变量和局部变量的区别
成员变量: 在类体里面定义的变量叫做成员变量: 如果在变量有static关键字修饰,就叫作静态变量或类变量: 如果该变量没有static关键字修饰,就叫作非静态变量或实例变量: 局部变量: 方法内定义 ...
- 帆软报表(finereport)JS实现cpt中详细单元格刷新
1.刷新固定单元格 setInterval(function(){ //获取第二行第 5 列 E2 单元格对象 var _changeCell = $("tr[tridx=1]" ...
- 深入MySQL(四):MySQL的SQL查询语句性能优化概述
关于SQL查询语句的优化,有一些一般的优化步骤,本节就介绍一下通用的优化步骤. 一条查询语句是如何执行的 首先,我们如果要明白一条查询语句所运行的过程,这样我们才能针对过程去进行优化. 参考我之前画的 ...
- python-利用faker模块生成测试数据
Python-利用faker模块生成测试数据 1.前言: Faker模块是一个生成伪数据的第三方模块,他提供了一系列方法,使用非常方便,在做自动化测试时,注册信息,用这个模块生成测试数据就体现了它的好 ...
- RFC3918组播转发时延测试——网络测试仪实操
一.简介 1.RFC3918简介 历史 · 在1999年3月成为正式标准 功能 · 评测网络互连设备或网络系统的性能 · 网络设备: 交换机,路由器- 内容 · 定义了一整套测试方法,为不同厂家的设备 ...
- 项目报错:/uploads: Read-only file system(解决办法)
项目报错:/uploads: Read-only file system(解决办法) 本来以为是service层没加注解,翻到最后才发现问题 原因是项目根目录没有对应的文件夹,在项目根目录创建uplo ...