《基于推导的人工智能编译器-v3-翟季东.pdf》由会员分享,可在线阅读,更多相关《基于推导的人工智能编译器-v3-翟季东.pdf(37页珍藏版)》请在三个皮匠报告上搜索。
1、清华大学 Tsinghua UniversityPACMAN基于表达式推导的人工智能编译器基于表达式推导的人工智能编译器翟季冬翟季冬清华大学清华大学2023.3.26第十届开源操作系统年度技术会议第十届开源操作系统年度技术会议-OS2ATCOS2ATC清华大学 Tsinghua University背景和意义人工智能芯片人工智能芯片成为深度学习应用的主要计算力成为深度学习应用的主要计算力摩尔定律发展逐渐放缓深度学习对计算巨大需求为了降低编程复杂度,为了降低编程复杂度,深度学习编程框架深度学习编程框架及及人工智能编译器人工智能编译器应运而生应运而生人工智能芯片人工智能芯片Caffe编程框架编程框
2、架?人工智能编译器人工智能编译器2清华大学 Tsinghua University研究背景和意义深度学习模型的编译过程深度学习模型的编译过程深度学习代码深度学习代码人工智能芯片人工智能芯片体系结构无关无关的图层优化(图融合、替换、删冗)中间数据流图中间数据流图体系结构相关相关的算子优化数学算子库,数学算子库,张量编译等张量编译等为了提高人工智能应用的优化与部署效率,现有的张量编译器将优化过程解耦为高层次的图优化和低层次的算子优化3清华大学 Tsinghua University研究现状-1:图层优化传统的优化框架以基于规则的图优化为主,如 TensorFlow中的 Grappler,需要领域专
3、家对优化规则进行人工设计为了提高优化效率,以 TASO、PET 为代表的自动优化框架成为当前图优化领域发展的主流趋势优点可自动地对图优化空间进行搜索利用形式化验证保证正确性缺点优化结果只能包含算子库中的算子需尝试大量错误状态,优化时间长4清华大学 Tsinghua University研究现状-2:算子层优化以 cuDNN、cuBLAS、oneDNN 等为代表的人工智能算子库在特定硬件平台上提供极高的性能,但需要大量人力进行维护近年来以 TVM、Ansor 为代表的算子自动生成框架成为了发展趋势优点基于原语的编程接口相对简单具有一定的跨平台能力缺点很难达到算子库的性能缺少对图层信息的感知行优先
4、列优先并行分块向量化并行分块典型的调度原语a=tvm.nd.array(numpy.random.rand(M,K).astype(dtype),ctx)b=tvm.nd.array(numpy.random.rand(K,N).astype(dtype),ctx)k=te.reduce_axis(0,K),k)A=te.placeholder(M,K),name=A)B=te.placeholder(K,N),name=B)C=pute(M,N),lambda x,y:te.sum(Ax,k*Bk,y,axis=k),name=C)s=te.create_schedule(C.op)xo,y
5、o,xi,yi=sC.tile(C.op.axis0,C.op.axis1,bn,bn)(k,)=sC.op.reduce_axisko,ki=sC.split(k,factor=4)sC.reorder(xo,yo,ko,ki,xi,yi)sC.vectorize(yi)spackedB.parallel(x)用户手写的调度代码缺少一套可以融合图层与算子层优化的方法5清华大学 Tsinghua University现有图层优化编译器-TASOTASO:自动的等价计算图搜索和替换搜索由后端支持的算子所组成的计算图利用形式化验证形式化验证证明计算图等价性用等价的高效高效计算图替换低效低效计算图计
6、算图计算图搜索搜索输入子程序利用所有可用的算子,遍历所有算子个数不大于某个固定阈值的计算图后端支持的算子发现并利用等价的高效计算图6清华大学 Tsinghua University基于局部等价优化的人工智能应用编译优化Pet:Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections通过构造算子库中没有的内存布局重排算子,引入局部等价优化,为张量程序的优化带来新的机会7清华大学 Tsinghua University局部等价优化8.=ConvW1W2XConvAdd
7、ConvW1W2XYZAdd完全等价优化缺点:会错过一些优化机会优点:保证功能上的完全等价性张量中有成千上万个元张量中有成千上万个元素,是否可以在优化时素,是否可以在优化时只考虑“大部分”元素只考虑“大部分”元素的等价性?的等价性?局部等价优化局部等价优化清华大学 Tsinghua University局部等价优化9.=局部等价优化.XDilated ConvYWXConvZWConvW1W2XConvAddConvW1W2XYZAdd完全等价优化缺点:会错过一些优化机会优点:保证功能上的完全等价性优点:可以取得更好的性能 更高效的算子 更高效的张量内存布局 更多的系统优化机会缺点:可能造成精
8、度降低、功能不等价清华大学 Tsinghua University局部等价优化10.=局部等价优化.XDilated ConvYWXConvZWConvW1W2XConvAddConvW1W2XYZAdd完全等价优化缺点:会错过一些优化机会优点:保证功能上的完全等价性优点:可以取得更好的性能 更高效的算子 更高效的张量内存布局 更多的系统优化机会缺点:可能造成精度降低、功能不等价利用局部等价变形所带来的性能提升的同时利用局部等价变形所带来的性能提升的同时维持功能的等价性维持功能的等价性清华大学 Tsinghua University突变生成器11突变生成器突变纠错器程序优化器突变突变生成器生成
9、器输入子程序利用所有可用的算子,遍历所有算子个数不大于某个固定阈值的子程序可发现完全等价完全等价与局部等价局部等价的优化变形后端支持的所有算子内存布局重排算子清华大学 Tsinghua UniversityPET 整体架构12优化后的程序突变突变生成器生成器突变程序突变突变纠错器纠错器完成纠错的突变程序程序优化器优化器输入程序清华大学 Tsinghua University程序优化器13基于搜索的基于搜索的程序优化器程序优化器输入程序优化后的程序突变生成器与突变生成器与突变纠错器突变纠错器多重线性张量程序已完成纠错的突变利用非线性算子非线性算子将程序划分为子程序对每个子程序搜索其候选突变通过后
10、处理优化后处理优化来保证优化后程序的高效性突变生成器突变纠错器程序优化器清华大学 Tsinghua University挑战:如何检测和纠错?14利用多重线性张量程序良好的代数性质,快速定位不等价结果并生成纠错内核Program fProgram g突变生成器突变纠错器程序优化器具体见 PET:Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections.OSDI.2021.清华大学 Tsinghua University优化案例 dilated conv下图中的例子
11、通过将输入张量进行变形把 dilated conv 的计算变为普通conv 的计算15清华大学 Tsinghua University实验结果在 batch size 分别为 1 和 16 的情况下,在常用的 DNN 模型上可取得最高最高 2.5 2.5 倍倍的加速比16清华大学 Tsinghua University小结:基于内存布局重排的局部等价优化提出了基于局部等价变换的张量程序编译优化技术结合张量程序应用特征,创新性地提出了张量重排与局部等价的优化策略对常见的模型和算子可以有最高2.5倍的性能提升但该方法依然存在一定局限性优化空间受限:仅能利用用户指定的算子进行搜索和优化无效搜索较多:
12、遍历搜索需尝试大量错误状态,增加优化开销17清华大学 Tsinghua University基于表达式推导的图算融合优化EinNet:Optimizing Tensor Programs with Derivation-Based Transformations以表达式替代算子表示,利用代数运算规则对表达式进行自动推导,以发掘更多的图算融合优化潜力18清华大学 Tsinghua University基于张量表达式的形式表达张量表达式可表示遍历顺序、维度取值空间、内存布局等张量程序的重要特征人工智能应用中常用的算子均可表示,如卷积、矩阵乘等卷积算子的张量表达式示例引入了循环符号,显式表示出元素遍
13、历的顺序及取值空间显式标示出循环和求和变量的取值空间,以支持取值空间的映射变换以索引方式表示内存布局,以支持不同的内存布局变量在保证完备性的同时,准确表示了张量程序的特征19清华大学 Tsinghua University在表达式中融合图层信息通过表达式嵌套的方法将图层信息融合到表达式中AW1BConvW2ConvAdd 将不同算子以嵌套方式放入同一个表达式,不同子表达式间可以进行联合优化20清华大学 Tsinghua University利用表达式替代基于计算图的算子表示AKBConv缺点:优化受限于算子库 只能进行图层优化 表示方法简洁 底层硬件无关优点:缺点:缺少现有优化技术的支持 优化
14、的搜索空间变大 保留图层信息的同时包含算子计算逻辑 底层硬件无关的同时摆脱算子库束缚 可支撑图算融合优化优点:基于算子的计算图表示基于表达式的表示21清华大学 Tsinghua University表达式推导及算子映射的基本流程将已有算子映射到算子库将已有算子映射到算子库将算子库没有的算子进将算子库没有的算子进行自动生成行自动生成推导过程利用基本推导过程利用基本代数运算规则,如代数运算规则,如交换律、结合律等交换律、结合律等输入输入计算图计算图张量张量表达式表达式输出输出计算图计算图22清华大学 Tsinghua University优点:扩展优化空间相对现有自动图优化工作的优化空间局限性仅能
15、利用用户预定义的算子用户预定义的算子进行搜索和优化局限于POR变换(predefined operator representable)利用张量表达式将优化空间扩展至非预定义算子探索涵盖任意张量表达式的变换融合图层和算子层信息,发掘更多优化机会23清华大学 Tsinghua University推导规则汇总推导规则推导规则形式化描述形式化描述*描述描述算子间算子间规则规则表达式拆分将一个表达式拆分为多个无依赖的表达式表达式合并将多个无依赖的表达式合并为一个表达式表达式融合将多个有依赖的表达式合并为一个表达式算子内算子内规则规则求和拆分将一个求和作用域拆分为两个变量替换将遍历操作的迭代变量进行替
16、换遍历合并将两个作用域下的遍历操作合并为一个边界放松将迭代变量的取值范围扩大边界收紧将迭代变量的取值范围缩小推导规则可方便地进行扩展,以支持更多类型的优化*由于篇幅原因,此处只保留了推导规则的形式化描述的公式,省略了其限制条件24清华大学 Tsinghua University自动推导过程利用 TVM 进行代码生成 代码生成部分以访存代码生成部分以访存密集型算子为主,复密集型算子为主,复杂度较低,调度简单,杂度较低,调度简单,搜索时间较短搜索时间较短 本研究的形式化表示本研究的形式化表示方法与方法与TVM的表达式的表达式有兼容,可以方便地有兼容,可以方便地通过代码映射生成通过代码映射生成TVM
17、代码代码25清华大学 Tsinghua University新变换的机会单个卷积算子的部分推导方案Conv3x3计算可变换为Matmul、Conv1x1、Conv2x2、Group conv等计算转换为表达式(c)(b)().VSVS+TM+T0WMatmulOffsetConcat(i).()(g)(h)求和拆分变量替换SSVS推导规则遍历合并边界放松边界收紧TMBRBTeOpOp预定义算子张量表达式算子沿循环变量求和公式推导表达式实例化()Conv2x2T0WOffsetReducer2s2Pad+Split(n)BRConv2x2T0WOffsetAddSplitConv2x1 Conv
18、1x2 Conv1x1(p)(22)(11)(11)(m)(22)(11)(11)(o)T0WGroup Conv3x3Duplicate(k)BRVS+SS+(c)(l)VS+TM+()(j)fMatmulT0WOffsetReducers(e)可行变换方案T0W(a)Conv3x3输入程序Conv1x1T0WOffsetReducekTranspose(f)SSVS+BR+()()()(c)()()()(d)表达式推导过程Im2col变换26图算融合的搜索空间带来了大量新的变换机会图算融合的搜索空间带来了大量新的变换机会清华大学 Tsinghua University新优化空间的挑战探索图
19、算融合优化空间面临的挑战1.自动算子生成开销高利用TVM在GPU上生成典型矩阵乘算子需要30分钟无法对所有表达式都进行自动算子生成2.图算融合搜索空间大1.预定义算子数量较少,但存在无穷多无穷多张量表达式EinNet搜索空间搜索空间27清华大学 Tsinghua University高效探索图算融合空间 算子库映射将计算密集部分映射至算子库,访存密集部分映射至算子生成框架同时利用算子库的高效性和算子生成框架的灵活性基于循环变量映射表,实现表达式至算子库算子的匹配根据预定义算子的张量表达式,自动生成相关信息循环变量映射表28清华大学 Tsinghua University高效探索图算融合空间 搜
20、索方法两阶段搜索算法探索式推导阶段探索式推导阶段:遍历所有推导规则,探索有效的变化方案利用搜索深度限制搜索范围收敛式推导阶段收敛式推导阶段:选择性执行推导规则,减小搜索开销比较当前表达式和目标算子(如Matmul)表达式,计算表达式距离表达式距离执行使距离减小的推导规则,加速实现到目标算子的匹配29清华大学 Tsinghua University基于表达式的图算整合优化框架整体设计计算图表达式自动推导模块代码生成模块基于代数运算的推导规则表达式合并与拆分表达式计算转换算子库匹配基于表达式的代码生成输入程序优化后的程序基于表达式的张量程序-翻译代数运算搜索计算融合搜索计算密集型访存密集型人工智能
21、芯片人工智能芯片Caffe现有编程框架现有编程框架通过代数表达式建立了图层与算子层的桥梁,为图算融合图算融合的优化技术提供了有力支撑30清华大学 Tsinghua University典型优化案例1 im2col复制输入张量,将卷积计算变换为矩阵乘计算31清华大学 Tsinghua University典型优化案例2 im2col 镜像优化累加中间结果,将卷积计算变换为矩阵乘计算32清华大学 Tsinghua University典型优化案例3 从Conv5x5到Conv3x3增加冗余计算,将5x5卷积计算变为6x6卷积,并最终变为3x3卷积33清华大学 Tsinghua University
22、实验结果 模型加速比在batch size分别为1和16的情况下,在常用的DNN模型和A100 GPU上可取得最高最高2.32.3倍倍的加速比34清华大学 Tsinghua University实验结果 不同后端加速效果在专家手写的算子库和自动算子生成框架上均能取得显著加速算子库:cuBLAS、cuDNN自动算子生成框架:AutoTVM、Ansor35清华大学 Tsinghua University基于表达式推导的图算融合优化技术通过代数表达式建立了图层与算子层的桥梁,为图算融合的优化技术提供了有力支撑,显著扩大了优化空间提出了基于代数表达式推导的张量程序自动优化技术,可以发掘出未被现有优化工作找到的优化策略提出了面向表达式的代码生成方法,同时保证优化的高效性与灵活性将计算密集型的部分映射到现有算子库以保证其高效性将访存密集型的部分通过其表达式进行自动代码生成来保证其功能上的灵活性36清华大学 Tsinghua University谢谢!PACMAN