摘要：深度学习编译器可以作为框架和硬件之间的公共组件和桥梁，最终希望实现的目标是我们只用开发一次，就能够为自动为任何设备生成最优代码。
本文分享自华为云社区《深度学习编译器简介》，原文作者：luchangli 。
最近的十几年深度学习发展十分迅速，业界出现了很多深度学习算法开发框架。同时，由于深度学习具有广泛应用场景和对算力的巨大需求，我们需要将深度学习算法运行在各种通用和专用的硬件上，比如各种类型的 CPU，GPU，TPU，NPU 等。那么这就出现了框架和硬件之间的组合爆炸，如图 1 所示。比如说 TensorFlow 要支持 GPU 计算，就要把 tensorflow 里面的所有算子开发一个 GPU 版本，如果又要支持 D 芯片，又需要把每个算子开发一个 D 芯片的版本。这个过程无疑非常耗时耗力。
图 1
于此同时，我们现在有非常多的算法网络，比如说 YOLO, BERT, GPT 等等。而这些算法网络是是由不同类型、不同 shape，不同连接关系的算子组成的。最终它们又运行在不同种类和型号的硬件上面。这就导致人工去为每个场景开发和实现最优算子成本很高。这里举了两个例子，如图 2 所示，算子融合是一个常见的性能优化方法，在融合之前，每个算子计算前后都需要把数据从内存读到缓存，再从缓存写回到内存。而融合之后，可以避免算子之间内存读写从而提高性能。传统的做法就是人工去根据算子连接关系开发融合算子，但是不同网络不同类别算子连接关系几乎不可能完全枚举。另一个例子就是算子调优，算子实现过程有很多参数会影响性能，但是传统人工算子开发方式很难去表达和维护这些参数，并且对这些参数进行调优从而实现不同 shape 和硬件的最优性能。
图 2
深度学习编译器正是为了解决上面一系列问题而诞生的，它可以作为框架和硬件之间的公共组件和桥梁，最终希望实现的目标是我们只用开发一次，就能够为自动为任何设备生成最优代码。比如为 CPU 开发的算子可以几乎原封不同的用于 GPU 和 D 芯片，从而显著降低成本。
这里简单介绍一下深度学习编译器的组成部分和功能，如图 3 所示。首先它的前端是从不同的框架拿到计算图，并且使用这个 Highlevel 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 代码并进行推理。
这里再简单介绍下 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 等，大家可能已经用过。
总之，深度学习编译器具有很多优势。比如易于支持新硬件，避免重复开发，采用一系列自动优化代替人工优化，可以实现极致性价比等。目前深度学习编译器也有一些不足，仍然出于一个快速发展的状态。例如调优时间长，对于复杂的算子无法有效生成，一个模型中深度学习编译器生成的算子能超过库调用的算子比例较低等，仍然需要大家持续投入和优化。
