当前位置: 首页 > 科技观察

关于深度学习编译器,这些知识你需要知道

时间:2023-03-13 12:34:04 科技观察

深度学习近十几年来发展迅速,业界涌现出很多深度学习算法开发框架。同时,由于深度学习应用场景广泛,对算力需求巨大,我们需要在各种通用和专用硬件上运行深度学习算法,比如各类CPU、GPU、TPUs、NPUs等。然后是框架和硬件之间的组合爆炸,如图1所示。比如TensorFlow要支持GPU计算,就需要开发一个GPU版本的tensorflow中的所有算子。如果要支持D芯片,就需要为每个运营商开发一个D芯片版本。这个过程无疑是费时费力的。图1同时,我们现在有很多算法网络,比如YOLO、BERT、GPT等等。这些算法网络由不同类型、形状和连接的运算符组成。它们最终运行在不同种类和型号的硬件上。这导致为每个场景手动开发和实施最佳运算符的成本很高。这里给出两个例子。如图2所示,算子融合是一种常见的性能优化方法。在fusion之前,每个算子在计算前后都需要从内存中读取数据到缓存中,再从缓存中写回。记忆。融合后可以避免算子之间的内存读写,提高性能。传统的做法是根据运营商连接关系人工开发融合运营商,但几乎不可能完全枚举出不同网络中不同类型运营商的连接关系。另一个例子是运营商调整。算子实现过程中有很多参数会影响性能,但是传统的手工算子开发方式很难表达和维护这些参数,也很难调优这些参数来实现不同的形态和硬件。最佳性能。图2深度学习编译器就是为解决上述一系列问题而诞生的。它可以作为框架和硬件之间的公共组件和桥梁。最终的目标是我们只需要开发一次,就可以自动为任何开发设备生成最优代码。例如,为CPU开发的算子几乎可以原封不动地用于GPU和D芯片,从而显着降低成本。这里简单介绍一下深度学习编译器的组成部分和功能,如图3所示。首先,它的前端从不同的框架中获取计算图,并用这个HighlevelIR数据结构来表示,以及然后在这个阶段进行一系列的图优化,比如常量折叠、算子融合、等价替换等。这是等效替换的示例。原来的计算图是这样的。如果我们把它换成一种计算方法,结果不会改变,但性能可能会更好。接下来,针对计算图中的每一个算子,使用领域特定语言DSL来描述算子的计算过程,并对算子进行优化。比如优化tiling、multi-core、double-buffer等算子。由于一个算子的计算过程通常是通过多次循环实现的,例如矩阵乘法就是三次循环。深度学习编译器可以很容易地在循环上进行各种变换,并对这些变换的参数进行调优,从而针对不同的形状和硬件获得最佳的算子实现。最后,基于低级IR为不同的硬件生成特定代码。图3最后介绍了业界现有的编译器项目。TVM是目前最完整、开源、框架无关的项目,已经被很多公司采用。TVM过程如图3a所示。TVM可以导入各种框架的模型,如TensorFlowpb、onnx、TorchScript等模型,统一由TVM称为Relay的HighlevelIR表示。IR中的每个算子都使用Tensor表达式的DSL进行计算描述和调度。这个DSL使用爱因斯坦的记法来描述算子的计算,而算子的计算一般体现在多个??for循环中。然后根据Halide的思路,利用schedule对这多个for循环进行各种变换,比如循环合并、拆分、序列变换等。最后由lowertolow-levelIR生成具体的设备代码并进行推理。这里简单介绍一下TVM是如何生成最优算子代码的。上面提到需要用compute来描述operator,然后需要对compute对应的多个for循环进行调度转换,也就是schedule。TVM的算子生成和调优经历了三代的发展。第一代的TVM/AutoTVM,这一代需要用户编写算子的计算和算子的调度。AutoTVM和TVM之间的区别在于,可以在调度中定义一些可变参数,然后可以使用遗传算法等方式调整参数。例如,如果一个循环被分成两段,那么可以优化分到哪里。第二代AutoScheduler(Ansor),这一代只需要用户自己开发operatorompute,Ansor就会自动按照一些规则进行调度转换。由于调度开发需要同时熟悉TVM的表达机制和底层硬件原理,调度开发往往难度很大,因此Ansor可以显着降低开发者的工作量和开发难度。缺点是Ansor调优时间较长,往往需要1小时才能调优一个操作员。以卷积网络为例,Ansor在某些场景下可以超越TensorFlow算子的性能,而TensorRT与TensorRT有一定的差距。第三代MetaSchedule(AutoTensorIR)刚刚起步,有望优化调优速度和性能,但目前还没有,我们拭目以待。TVM的实现包括华为D芯片TBE算子开发工具,在TVM的基础上增加了对D芯片的代码生成支持。TVM走的是Halide计算+调度的路线,另外还有一个编译器走的是多面体算法路线,比如华为自研的TensorComprehensions、Tiramisu、AKG等。这种方式和Ansor一样,只需要用户自己开发operatorcompute,而不需要schedule,更加人性化。其中,AKG已经应用于MindSpore的图像计算融合。其他深度学习编译器还有TensorFlow的XLA、TensorRT等,你可能用过。总之,深度学习编译器有很多优点。例如,易于支持新硬件,避免重复开发,使用一系列自动优化代替手动优化,达到极致的性价比。目前,深度学习编译器也存在一些不足,仍处于快速发展的状态。比如调优时间长,不能有效生成复杂的算子,深度学习编译器生成的算子在模型中能超过库调用算子的比例比较低等,这些仍然需要持续投资和优化。