AI编译器之后端优化(笔记)
日常感谢昇腾架构师ZOMI酱(我愿称之为肝帝!)的超级贡献,欢迎一键三连
ZOMI酱的个人空间-ZOMI酱个人主页-哔哩哔哩视频
他的github地址
GitHub - chenzomi12/DeepLearningSystem: Deep Learning System core principles introduction.
第一节-AI编译器后端优化介绍
在之前我们关注的是前端语言转化后的计算图前端优化,现在将目光更进一步,让我们看看后端优化是怎么玩的,那么前后端优化有什么区别呢?
比如我们看看其中这个卷积算法:
我们可以对输入输出和内部循环方式、内存访问方式进行各种各样的优化。
接下来看看后端优化的几个主要步骤:
首先把计算图高级ir转换为低级ir:
接着后端优化会把它变成更低级的ir,然后进行一些代码生成到具体的硬件上进行执行。
这里面的relay.add就是TVM IR的一个定义,最后由TVM的编译器生成底层内部的底层IR。
接下来再看看具体的算子相关问题:(共两类,访存密集和计算密集型)
假如我们有一个向量vec a以及向量vec b,都是从内存取出来的,执行了alu的乘法后再放入cache(内存),接着再从内存读出来再读下一次;这需要大量频繁的访存:
第一个表前面是卷积的参数,右边是计算量;不同卷积的入参计算量也是不同的,这里的计算量比我们的访存量大了很多,这就是计算密集型:
对于这两种算子优化还是有着很多挑战的:
访问和存储是天平的两端,这实际上是一个复杂的组合优化问题。。。
Cudnn其实早期只有100多个算子,现在又接近200多个算子。
那么问题就来了,AI算子迭代那么快,怎么赶得上更新的速度?
所以我们需要AI编译器,有人提出了自动kernel生成的方式,目标是给定算法自动生成目标平台上的高性能实现。
(下图上面部分提到了一些不错的文章可以去看看)
第二节-算子计算与调度
接下来看一个算子计算的具体例子:(注意什么叫做不同的调度方式与唯一定义)
但是高斯滤波这种图像模糊方法,有很多种的调度方式:
我可以先对横坐标进行迭代,再对纵坐标进行迭代;我们也可以先对纵坐标进行迭代再看横坐标,这里体现的就是调度方式的不同。
计算和调度解耦的好处是可以通过不同硬件特点好处对具体算法计算进行计算优化。
右边就是加速优化的高斯滤波,做了很多分片(没有大循环了)与数据展开(buffer),这就是调度的优化。
接下来看看神经网络算子的计算特征:
我们再来看一下根据算子计算特征抽象后的高斯滤波的代码:
大部分的算子调度都是由内存分配、循环、计算实现的。所以我们就可以根据这个去定义调度树:
我们可以根据形状不同去区别不同的阶段,我们可以用调度树很好的表达算子计算原语(类似画算子的状态机),真正loop结束后就是计算节点菱形。那么我们该如何理解调度树的语义?我们可以对调度树进行优先遍历DFS然后转换为对应的程序代码。AI编译器后端的目标就是在算子原语不变的情况下去优化调度树。
接下来看几种对调度树修改的方法(不影响算子原语):
这里的内联的意思是把f*x放到后续的计算,然后把gx,hx往前提。
那么基于调度树理解这些有啥用?—— search algorithm for schedules
左下角是我们调度树节点计算的耗时,我们需要根据整体,找到满足代价函数最优的情况,这时候就是最好的调度策略。
实际的调度策略可能是比上面复杂得多的多的。
具体的一些调度方法其实还可以去看看halide(http://halide-lang.org/),可以很方便的实现一些GPUAPI的调度方法。https://zhuanlan.zhihu.com/p/346468141
Halide是用C++作为宿主言语的一个图画处理相关的DSL(Domain Specified Language)言语,全程范畴专用言语。首要的效果为在软硬层面上(与算法自身的规划无关)对算法的底层加快。在OpenCV(传统图画处理库)中部分算法运用了Halide后端,而TVM(神经网络编译器)也是用了Halide的思维去优化神经网络算子。
具体也可以看openmmlab的一期开放麦对它进行了介绍:
【社区开放麦#27 | 部署神器Halide, 实现高性能算法】 https://www.bilibili.com/video/BV1Tg411B7ch/?share_source=copy_web&vd_source=b42f227a4f2d413fbde18499d83227cf
第三节-算子优化手工方式
我们来看看算子优化在不同框架的优化方式:
再展开具体这些东西是啥之前,我们先来看一个例子——基于源码的修改:
看到上面,我们可以看到先遍历一个10000再遍历200,外层循环遍历比较多,这里可能会导致内存消耗比较大。那如果我们先遍历200再遍历10000,每次在堆栈迭代下一个i,可以很好的利用内存空间。
另一个就是循环变量的实例化:
对于一些重复计算的表达式可以先构建好。这里是空间换时间的方法,之后在循环就可以通过读取内存或cache的方式直接快速读取到tmp里面,不用每次都进行一个计算。
另外我们可以通过把一些循环中用到的东西提前到外面实例化处理,避免非必要的函数调用开销。(不过其实我觉得很多情况下编译器也会对这些进行优化了。。)
另外,大家不要觉得AI编译器里面每一层我都要去实现自己的操作,其实每一层都有可替换方案和对应的开源项目。
另外,我们可以把算子调度优化方法分成三大类型:
为什么循环优化有那么多内容呢?那是因为我们的计算特征以多重循环为特点,多维张量计算为主要数据结构。所以算子优化大部分集中在循环loop的优化。
如果想要理解更深入可以看看参考文献的一些内容:
第四节-算子循环优化
4.1 循环展开
这部分看起来还是有点抽象,我咨询了一下我的朋友:
其实就是乱序的一种情况, 一个循环里面的有很多指令,后面的指令先执行,在循环里可以体现为下一个(甚至下几个)循环的某条指令先于当前指令执行,然后总体上来看就像循环里的指令重新排序了. 当然因为同一个循环里指令所指定的寄存器名字是一样的,为了解决这个冲突需要在寄存器组里使用虚拟化技术,硬件上叫寄存器重命名. 相当于不同的循环用了同样名字的寄存器,大家都叫xxx,但是实际上硬件会自动把他们处理成不同的寄存器. 还有一点不重要的,循环最后一条是跳转,本来要循环10次的,把两次loop合并为一次相当于跳转次数减半了。不过这个优化比较小,而且有延迟槽+转移预测(这个命中率特别高),所以我直接把他忽略了,面试还是要说的,但是工业上这个提升有没有都无所谓,当然如果是代码会在循环的不同层级和其他代码块之间频繁互相跳转就要另外说,但是我要强调这个情况是极其极其丑陋的代码习惯(说的就是goto)风险大,优化难。
接下来我们一起看一看什么叫做循环展开:(PS这其实是比较基础的,更详细的要看体系结构量化方法或者国科大的体系结构教材)
这里j在循环的时候跳两个位,里面执行了两个数据的操作。
4.2 循环分块
分块的原因是cache显存内存是有限的。了解这些特性是为了方便我们更好的写AI编译器自动调度的策略。分块与缓存块的大小应该一致,可以提高处理器整体的效率。(对齐)实现思路如下:
这个具体做法看起来有点抽象,一起来看个具体的例子:
本来是从0到m,我们把循环按照T大小进行了一个分块,每一次对分块后的j_0开始一直从j_0加到j_0+T。j_o是out loop,里面的j_i叫做inner loop。这个数据就可以塞到处理器里面,而不会导致在迭代的时候导致cache missing,需要重新对数据调入调出,可以最大程度利用芯片的内存空间。
其实实际问题比理想的情况更复杂:
4.3 循环重排
我们再来看看循环的排布策略:
假设这里的m非常非常大,n相对小。这时候需要把里面的数据全部塞入cache其实是不可能的,所以我们可以把m放到外面,比较小的n放到里面。假设这个n只有100,和cache的空间相同,提高了cache的命中率也就提高了芯片的满载负荷。
具体例子可以参考:
计算机系统基础学习笔记(4)-Cache友好代码https://blog.csdn.net/qq_43336390/article/details/106071998
4.4 循环融合
接下来再看看循环融合:
这里上面的循环的边界条件是不同的,所以我们要先进行一个对齐,让他们可以融合到下面的同样的从1到n-1。这种方法就可以加强软件的流水线并行效率。这里计算的数据还是要满足cache 友好的方式才可以让计算效率最大化。
4.5 循环拆分
并行的处理器是不擅长处理控制流的。 而下面的执行比较慢,我们就让他单独执行(15-17行),然后单独判断,这可以加快整体处理器的效率。
最后我们复习一下今天学习了什么:
第五节-指令和存储优化
接下来先看看指令优化部分(向量化与张量化)
5.1 向量化
我们以这个图为例,假设想要计算四个字节/单位的数据,首先从内存的低位开始读取,一次性读取四个数据;接着再从下一个低位读取四个数据,和之前的四个数据一起计算,不断循环这种+操作,被称为向量化操作——一次处理多个数据。
在没有执行向量化的时候,第2到5行的代码主要是对数组进行单纯的遍历求和。
下面的部分是用了向量化指令的结果。
(这里可能是伪代码,我们看看一个实际情况下的向量化)
现代编译器也已经有各种自动向量化的操作,例如:
详细可以参考:
向量化代码实践与思考:如何借助向量化技术给代码提速
https://blog.csdn.net/AlibabaTech1024/article/details/126359743
具体请以手册为准。
5.2 张量化
接下来看看张量化:
我们来看看左侧的图,这里的左边小部分是很多cuda core 右边是tensor core。
Tensor core主要就是用来做张量计算,考虑到神经网络的gemm等计算原理,实际上硬件设计也设计成了对应的形式。
接下来看看一些主流厂商是怎么做的:
这里给出对应的官方地址:
https://docs.nvidia.com/cuda/cublas/
https://developer.nvidia.com/cudnn
https://github.com/oneapi-src/oneDNN
cuBLAS是线性代数库,适合做矩阵乘;卷积可以用矩阵乘法的方式实现,这也是cuDNN计算卷积的方法之一,不过cudnn有更加深度的优化。
对于局限性,所以我们需要ai编译器提供增量化的指令或调度语言。
新的领域知识让我们的关注重点从向量化到张量化。自然语言里面甚至有些变长相关的计算。面对多元化的情况就有了TVM的说法。
5.3访存延迟
接下来我们再看看存储优化相关(分为访存延迟和存储分配):
在ai进行训练的时候有大量的核或者线程来执行我们的算子或计算,但是我们的计算是严重依赖memory的,计算结果中间结果都需要存储;这时候存储和计算就会重叠;我们可以尽量让他们重叠,最大提高硬件设备的利用率,这时候就需要访存延迟。(避免都在访存没人在计算)
接下来让我们看看gpu的访存延迟是怎么产生和解决的:
里面的cache分为三层,第一层是dram,第二个是 l2 最后一个是 l1,l1就离我们的计算单元cuda core很近了,具体计算单元到内存之间有一个 warp scheduler,专门管理各种线程。
下面就是warp scheduler 和我们具体指令之间的关系:warp scheduler会对我们的指令做好分发和预分配,然后给我们的instruction dispatch unit(IDU),接着IDU就会把具体的线程分发到不同的worker上执行。
假设我们有四个warp,每个warp都有一个指令,现在在一个读数据的过程当中,如果warp在读取数据,就会导致我们整个系统或者说线程的阻塞。这时候warp1和warp2不会因为warp0 的执行而阻塞;因为gpu 会根据 warp schedule去解决访存延迟的问题。
具体步骤也可以参考:
CUDA——SM中warp调度器调度机制&&访存延迟隐藏
https://blog.csdn.net/weixin_44444450/article/details/118058031
访存是有先后顺序的(SM、调度器是有限的)。
(1)应该让先访存完毕的warp去执行尽可能多的指令(不然运算单元空着也是浪费啊),去隐藏其它warp的访存时间。
(2)增加active warps的数量,让尽可能多的warp去隐藏访存延迟。(这个有局限性,warps是有限的)
接下来我们看看DAE,解耦访问和执行的架构:这个就是大部分NPU CPU所采用的方式
接下来看看软件部分实现:
TVM这么转换的好处是pipeline是明确的(下图的monolithicpipeline),可以通过软件去控制,但是这里有个问题,硬件的正确执行顺序需要软件控制而且需要低级的同步来实现。实现的算法好坏直接决定访存延迟性能。
5.4存储分配优化
接下来我们一起看看内存分配:(局部变量,全局变量,堆变量)
那么有同学问了,到底啥是堆变量?
接下来看一个简单的内存排布(详细直接去看程序的装载映像内存分配)
这个只是传统的cpu和编译器结合的方式。在神经网络跑的地方(AI加速器等)里面就复杂多了,这里不仅有L1L2 还有dram等。。。。面向不同的ai加速器不同的内存分配方式其实很复杂,人工控制cuda会很麻烦,所以交给AI编译器解决。
第六节-Auto-Tuning原理
其实传统编译器早就有了auto tuning的概念。以下是传统编译器对应的概念:
下面出自一个文章,对gemm矩阵乘进行优化;完美运用了现代处理器的四级缓存:(L3 L2 L1 还有alu旁边的一些寄存器进行了加速)
传统编译器发展后分为两个方法:优化选择 optimize selection 优化顺序 optimize sequence 。
现在我们从传统编译器回到AI 编译器,看看有如何特点。
大量相类似的算子计算模式比如不同的Normalization。。。pooling等。
这个问题其实非常难,需要大量的大神去研究。
现在一般来说总结为三个步骤:
接下来先看看参数化的内容:不管是循环优化、指令优化还是内存优化都可以组成一些相对固定的指令结构,变成可以调度的原语。TVM可以对调度模板进行建模。
下图中的意思:s是数据内容进行split。循环切分因子是64,意思是可能的参数空间、范围。Factor是可以作为搜索空间。
有了参数化之后就需要建立成本函数:
有了成本函数后就是开始搜索,对我们参数化的内容调度方式进行搜索,以便找到最好的参数化配置方案:
接下来我们看看TVM里面的一个栈是怎么实现的:
这个大哐哐里面基本就是 auto tuning的内容。这里的Declarative Tensor Expressions 是参数化,把tensor的表达方式抽离,类似halide对计算和调度逻辑进行解耦,右边是对一些硬件感知优化的原语。
有了参数化之后就可以很好的对我们的计算和硬件进行一个表示,有了这个表示后就建立了我们的成本模型,然后用ML basedautomated optimizer,之后就对loop program不断搜索让我们找到最优的成本函数。找到对应策略后就用编译器生成对应的指令代码给真正的硬件进行部署使用。
来看看ansor(TVM的最新版)
首先我们拿到AI框架表达计算图的子图, 然后对子图进行切分获得重要的子图,拿到重要的子图就需要去 4、5步骤;比如section 4:确定优化的结构,随机注释(主要是对应右边看到很多不同的循环和方式,这里会进行随机的初始化,产生了很多配置的方式);接下来就是启发式搜索与cost model等等。。。。最后迭代到比较好的策略后就可以把他部署到实际在硬件上可以执行的指令(黑盒测试),再迭代回来不停的优化。
完结撒花~ 大哥有ppt吗 zomi酱的那个github里面有全部 多谢,非常有用
页:
[1]