摘要:深度学习编译器能够做为框架和硬件之间的公共组件和桥梁,最终但愿实现的目标是咱们只用开发一次,就可以为自动为任何设备生成最优代码。
本文分享自华为云社区《深度学习编译器简介》,原文做者:luchangli 。前端
最近的十几年深度学习发展十分迅速,业界出现了不少深度学习算法开发框架。同时,因为深度学习具备普遍应用场景和对算力的巨大需求,咱们须要将深度学习算法运行在各类通用和专用的硬件上,好比各类类型的CPU,GPU,TPU,NPU等。那么这就出现了框架和硬件之间的组合爆炸,如图 1所示。好比说TensorFlow要支持GPU计算,就要把tensorflow里面的全部算子开发一个GPU版本,若是又要支持D芯片,又须要把每一个算子开发一个D芯片的版本。这个过程无疑很是耗时耗力。算法
图 1express
于此同时,咱们如今有很是多的算法网络,好比说YOLO, BERT, GPT等等。而这些算法网络是是由不一样类型、不一样shape,不一样链接关系的算子组成的。最终它们又运行在不一样种类和型号的硬件上面。这就致使人工去为每一个场景开发和实现最优算子成本很高。这里举了两个例子,如图 2所示,算子融合是一个常见的性能优化方法,在融合以前,每一个算子计算先后都须要把数据从内存读到缓存,再从缓存写回到内存。而融合以后,能够避免算子之间内存读写从而提升性能。传统的作法就是人工去根据算子链接关系开发融合算子,可是不一样网络不一样类别算子链接关系几乎不可能彻底枚举。另外一个例子就是算子调优,算子实现过程有不少参数会影响性能,可是传统人工算子开发方式很难去表达和维护这些参数,而且对这些参数进行调优从而实现不一样shape和硬件的最优性能。缓存
图 2性能优化
深度学习编译器正是为了解决上面一系列问题而诞生的,它能够做为框架和硬件之间的公共组件和桥梁,最终但愿实现的目标是咱们只用开发一次,就可以为自动为任何设备生成最优代码。好比为CPU开发的算子能够几乎原封不一样的用于GPU和D芯片,从而显著下降成本。网络
这里简单介绍一下深度学习编译器的组成部分和功能,如图 3所示。首先它的前端是从不一样的框架拿到计算图,而且使用这个High level IR的数据结构来表示,而后在这个阶段进行一系列图优化,好比常量折叠,算子融合,等价替换等。这里展现了一个等价替换的例子,原来计算图是这样的,咱们给它换一个计算方式,结果不变,可是性能可能更优。接着,对于计算图里面的每个算子,采用DSL一种领域特定的语言来描述算子的计算过程和对算子进行优化。好比对算子进行tiling,多核,double-buffer等优化。因为算子的计算过程一般是用多重循环来实现的,好比说矩阵乘法是一个三重的循环。深度学习编译器能够很方便的对循环进行各类变换,而且对这些变换的参数进行调优,从而获得不一样shape和硬件的最佳算子实现。最后,基于low level IR为不一样硬件生成具体的代码。数据结构
图 3框架
最后介绍下业界已有的编译器项目。目前生态最完善,开源的,框架不依赖的项目首推TVM,已经被不少公司所采用。TVM流程如如图 3a所示,TVM能够导入各个框架的模型,例如TensorFlow pb,onnx,TorchScript等模型,统一用TVM称为Relay的High level IR进行表示。IR中每一个算子采用了Tensor expression的DSL来进行计算描述和调度。这个DSL采用Einstein’s notation的方式进行算子的compute描述,算子compute通常体现为多重for循环。而后基于Halide思想使用schedule对这个多重for循环进行各类变换,例如循环合并,split,顺序变换等等。最后,lower到low-level IR生成具体的device代码并进行推理。ide
这里再简单介绍下TVM具体如何生成最优的算子代码。上面介绍了算子须要进行compute描述,而后须要对compute对应的多重for循环进行调度变换,即schedule。TVM的算子生成和调优经历了3代发展。第一代TVM/AutoTVM,这一代须要用户编写算子的compute和算子的schedule,AutoTVM与TVM的区别在于能够在schedule定义一些可变的参数,而后采用例如遗传算法进行参数调优。例如把一个loop切分为2段,那么在哪里进行切分是能够进行优化的。第二代AutoScheduler (Ansor),这一代只须要用户开发算子ompute,Ansor内部自动根据一些规则进行调度变换。因为调度开发须要同时熟悉TVM的表达机制和底层硬件原理,schedule开发每每具备很高的难度,所以Ansor能够显著下降开发人员工做量和开发难度,缺点就是Ansor调优时间很长,每每须要1小时才能调优1个算子。以卷积网络为例,Ansor在部分场景能超过TensorFlow算子性能,距离TensorRT实现有必定差距。第三代Meta Schedule (AutoTensorIR)才处于起步阶段,预期会对调优速度和性能进行优化,暂时还不可用,咱们拭目以待。工具
TVM的落地包括华为D芯片TBE算子开发工具,在TVM的基础上增长了D芯片的代码生成支持。TVM采用了Halide计算+调度的路线,还有另一种采用polyhedral算法路线的编译器,好比Tensor Comprehensions,Tiramisu,华为自研的AKG等。这种方法跟Ansor同样,也只须要用户开发算子compute,无需开发schedule,所以对用户也较为友好。其中AKG已经用在了MindSpore的图算融合里面。其余的深度学习编译器还有TensorFlow的XLA、TensorRT等,你们可能已经用过。
总之,深度学习编译器具备不少优点。好比易于支持新硬件,避免重复开发,采用一系列自动优化代替人工优化,能够实现极致性价比等。目前深度学习编译器也有一些不足,仍然出于一个快速发展的状态。例如调优时间长,对于复杂的算子没法有效生成,一个模型中深度学习编译器生成的算子能超过库调用的算子比例较低等,仍然须要你们持续投入和优化。