树模型编译器
Treelite是用于有效部署决策树集合的模型编译器。
- 决策树Ensemble的编译优化:https://zhuanlan.zhihu.com/p/597511551
深度学习编译器
深度学习编译原理
AI 编译器前端优化
前端优化作为AI编译器的整体架构主要模块,主要优化的对象是计算图,而计算图是通过AI框架产生的,值得注意的是并不是所有的AI框架都会生成计算图,有了计算图就可以结合深度学习的原理知识进行图的优化。
- 计算图层(Graph IR)
- 算子融合(OP Fusion)
- 布局转换(Layout Transform)
- 内存分配(Memory Allocation)
- 常量折叠(Constant Fold)
- 公共子表达式消除(CSE)
- 死代码消除(DCE)
- 代数简化(ARM)
AI 编译器后端优化
后端优化作为AI编译器跟硬件之间的相连接的模块,更多的是算子或者Kernel进行优化,而优化之前需要把计算图转换称为调度树等IR格式,然后针对每一个算子/Kernel进行循环优化、指令优化和内存优化等技术。
- 算子循环优化
- 指令和内存优化
深度学习编译工具
TVM
XLA
Glow
深度学习编译优化
计算图优化:XLA、TVM、nGraph
算子生成:AutoTVM、TC、Halide
图优化
目标:通过图的等价变换化简计算图,从而降低计算复杂度或内存开销。
数据流图作为深度学习框架中的高层中间表示,可以允许任何等价图优化Pass去化简计算流图或提高执行效率
图优化 (1):算术表达式化简
通过代数运算等价变换化简计算图,如:
a * 0 -> 0
a * broadcast(0) -> broadcast(0)
a * 1 -> a
a * broadcast(1)-> a
a + 0 -> a
a + broadcast(O)-> a
log(exp(x)/y)->x-log(y)
图优化 (2):公共子表达式消除
目的是将相同输入的表达式进行消除,由一个节点来代替,复用计算结果
图优化 (3):常数传播
如果一个算子的所有输入张量都是常数的话,那么该算子的结果也为常数张量
在编译器计算并化简
例子:假设参数0和参数1为常数张量,最终的图可以化简为什么?
注意:常数传播可能会引起内存的扩张,如:Broadcast
图优化 (4): GEMM自动融合
Batch same-type operators to leverage GPU massive parallelism
图优化 (4): GEMM自动融合
通过将输入张量合并成一个大的张量来实现将相同的算子合并成一个更大的算子,从而更好的利用硬件并行度
图优化(5):算子融合
向量化的多个算子的操作可以合并成一个向量化操作
减少内核启动开销
减少内存的读取,提高计算密度
图优化 (6):子图替换
利用子图匹配识别出可替换的复杂子图,替换为更高效的合并算子
图优化 (6):随机子图替换
TASO: Optimizing Deep Learning Computation with Automatic Generation of Graph Substitutions
小结: 图优化的总结
计算图作为深度学习编译框架的第一层中间表示
基于计算图的优化算法:
算术表达式化简
公共子表达式消除
常数传播
矩阵合并
算子融合
子图替换/随机子图替换
思考:
你可以想到哪些其它的计算图上的优化?
计算图还有哪些其它的好处
内存优化
内存优化
目标:通过对计算图的变化以及张量的合理分配来降低内存使用的总量。
内存优化 (1):基于拓扑序的最小内存分配
将计算流图按照某种拓扑序进行排序。如BFS,ReverseDFS等
按照节点的拓扑顺序依次分配其使用到的输出张量。
当一个张量后面没有其它算子使用时, 则回收到内存池。
当所有张量分配完成后,内存池的最大分配空间就是该计算图需要的最小内存
拓扑序的选择会同时影响模型的计算时间和最大内存占用
内存优化 (2):根据整数线性规划求解最优内存放置
目标:给定任意的计算图最小化执行时间
约束:有限的快速内存,如GPU内存
变量:每一个张量是否放置在快速内存中,还是较慢的外存中,如CPU内存
过程:最优化张量的移动
需求:每个内核计算的测量时间
方法:将上述的优化问题建模为整数线性规划问题
内存优化 (3):张量换入换出与张量重计算
方法:
DRAM存储量相对GPU显存来说比较大,可以将数据在GPU与DRAM之间进行转移,或者直接重新计算
发现:
在训练过程中张量的访问模式比较规律
核心思想
根据运行时的张量访问来动态管理内存
替换+预取
重新计算
小结: 内存优化的总结
·基于计算图的内存优化算法: 基于拓扑序的最小内存分配 根据整数线性规划求解最优内存放置张量换入换出与张量重计算 ·思考: 你可以想到哪些其它的内存优化方法?
内核优化
问题: 每个后端平台针对每个算子都需要单独实现至少一个内核
考虑:编程模型、数据排布、线程模型、缓存大小等等
张量运算编译
核心思想:分离计算逻辑与调度逻辑
通过张量运算表达式表示每个算子的通用计算逻辑
通过调度语言描述算子在映射到具体硬件上时的调度空间
相关工作:
TVM、 Halide、 TACO、 Tensor Comprehension, FlexTensor等
张量运算表达式:例:TVMIR
C=A*B
C= tvm.compute((m,n),lambda i,j:tvm.sum(A[i,k] * B[k,j], axis=k)
其它张量运算表达式
Affine Transformation
out = tvm.compute((n,m),lambda i, j: tvm.sum(data[i, k] * w[j, k], k))
out = tvm.compute((n,m),lambda i, j: out[i, j]十bias[i])
Convolution
out=tvm.compute(c, h, w), lambda i, x, y: tvm.sum(data[kc,x+kx,y+ky] * w[i, kx, ky], [kx, ky, kc]))
ReLu
out = tvm.compute(shape, lambda *i: tvm.max(0, out(*i))
其它算子调度优化
每一种优化都可能产生出多个内核代码的实现
利用自动机器学习 Auto Tuner 在给定时间内搜索出最高效的实现
小结: 内核优化的总结
内核优化与内核生成
算子表达式 算子表示与调度逻辑的分离 自动调度搜索与代码生成
调度优化
NNFusion:全局计算调度优化
目标:通过将多个算子进行协同调度以及精确映射每一个算子到硬件计算单元来充分利用硬件并行度
中间表示:数据流图+细粒度算子并行单元
结果:将每个子图编译成一个硬件计算内核
充分减少上层调度的开销
高效利用硬件并行度
通过引出新的调度原语来支持任务级调度
APEEND: 将一个算子中的一个任务调度到硬件的一个计算单元上
GROUP_SYNC: 维护任务间的依赖关系
挑战:调度算子的任务到GPU上的挑战
简单的任务级调度可能引起正确性问题
- 错误依赖
- 死锁
依赖关系的映射
通过将硬件计算单元抽象到软件可控计算单元,并引入细粒度任务级同步支持来保证计算正确性。
并行度的映射
任务级调度可以支持任意算子之间的并行调度,从而最大化硬件利用率。
本次课程总结
深度神经网络编译器的概念与架构
中间表达、前端、后端、优化过程
计算图优化
算术表达式化简、 公共子表达式消除、 常数传播、矩阵合并、 算子融合、 子图替换/随机子图替换
内存优化
基于拓扑序的最小内存分配、 根据整数线性规划求解最优内存放置、 张量换入换出与张量重计算
内核优化
算子表达式、 算子表示与调度逻辑的分离、 自动调度搜索与代码生成
调度优化