XLA的中间表达分析和调度内存分配分析(一)

本章是系列文章的案例学习,不属于正篇,主要介绍了TensorFlow引入的XLA的优化算法。XLA也有很多局限性,XLA更多的是进行合并,但有时候如果参数特别多的场景下,也需要进行分割。XLA没有数据切分的功能。当前最主流的编译器领域的编译优化功能还是mlir。

本文中的所有内容来自学习DCC888的学习笔记或者自己理解的整理,如需转载请注明出处。周荣华@燧原科技

9.1什么是XLA

第一次看到Accelerated被简称为X的时候,有点奇怪,因为Accelerated里面可没有一个字母是X,但Accelerated的发音和X相同,这样简化之后可以避免一个简写中存在多个A的不协调,XLA读起来确实比ALA朗朗上口一点:)

在2017年XLA诞生的时候,那时给出的帧处理加速数据如下:

带来相应加速效果的主要因素是通过分析和调度内存使用,删除了一些中间表达的存储缓存,其中一个主要的方法就是缓冲区指派算法,也就是本文主要准备描述的。

XLA的设计理念是一种近似SSA的中间表达:

XLA: Optimizing Compiler for Machine Learning | TensorFlow中有个油管视频详细讲解了XLA的原理,通过这个也可以理解一下TensorFlow的原理:

tf.function→ tf2xla桥→ 优化前的xla hlo→ xla的一些列优化→ 优化后的hlo→ 可执行binary→tf2xla桥→ tf runtime→ target arch上执行

9.2静态内存分配分析9.2.1为什么可以做分析9.2.2静态内存分析的优势9.2.3静态内存分析的局限性9.3缓冲区管理的目标

缓冲区定义:每个算子定义一个缓冲区

缓冲区申请、支配原则:

9.4缓冲区分析(有可以称为别名分析)

缓冲区分析的过程和指针分析的过程有很多类似的地方,所以很多地方又称为别名分析。

一个IR需要定义≥1个逻辑缓冲区

用{def, {}} 来定义一个缓冲区

缓冲区{b, {}} 和{b, {1}}可以相互覆盖

来自不同IR的逻辑缓冲区可以复用同一块内存

例如对下面的伪代码,可以知道d和b是别名关系,因为它们指向同一片内存:

9.4.1定义所有指令的所有逻辑缓冲区

按拓扑顺序遍历(选择什么顺序?逆后根排序)计算图,为每个指令分配缓冲区,例如上面的伪代码,生成缓冲区如下:

1 Buffer(a, {}) : [ (a, {}) ]
2 Buffer(b, {}) : [ (b, {}) ]
3 Buffer(c, {}) : [ (c, {})]

9.4.2HLO内部的别名分析后的结果

1 Buffer(a, {}) : [ (c, {0}), (a, {}) ]
2 Buffer(b, {}) : [ (c, {1}), (b, {}) , (d, {}) ]
3 Buffer(c, {}) : [ (c, {})]

9.4.3跨HLO的别名分析

基于近似SSA的HLO语法定义,编译过程变得简单了很多(SSA化是很多编译中的主要工作)

9.4.4基于上面伪代码的生命周期干扰分析

从下面生成的图来看,a和b互相干扰,不能公用缓冲区,e理论上是d的拷贝别名,所以和b也是别名关系。和寄存器分配不同的是,考虑到多线程执行场景,不同流中要用到的缓冲区不能分配到同一个组,所以a/b虽然和e在下面的计算图中没有干扰,但由于e是后面HLO的输入,所有e不能和当前计算图中的任意一个缓冲形成别名关系。

9.5buffer指派的功能9.5.1将能够重用的buffer尽可能重用

没有生命期干扰的缓冲区都可以分配到同样的内存

9.5.2缓冲区分配复合

将着色相同的缓冲区复合到一起申请(可能不同缓冲区占用某个实际缓冲区的不同部分,但大家相互之间的关系决定了它们可以相邻申请)

9.5.3从全局分析去掉内存碎片9.5.4峰值内存压力预测9.5.5内存分配统计

© 版权声明
THE END
喜欢就支持一下吧
点赞210 分享
评论 抢沙发
头像
欢迎您留下宝贵的见解!
提交
头像

昵称

取消
昵称表情代码图片