基于推导的人工智能编译器-v3-翟季东.pdf

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