微软亚洲研究院推出AI编译器界“工业重金属四部曲”

（本文阅读时间：14分钟）

编译是程序开发的一个重要步骤——把用高级语言书写的源代码翻译成在计算机硬件可执行的机器码，而编译器就是实现这一功能的特殊应用程序。如今，人工智能技术和大模型无疑是当今计算机领域的 C 位担当，其自身的特性也对编译器提出了新的挑战。

从最初的 RNN、CNN 到 Transformer，人工智能的主流模型架构在不断变化，这意味着上层的应用程序也在随之改变。同时，底层加速器硬件如 GPU、NPU 等，也在快速迭代更新，有些新的硬件设计甚至颠覆了之前的架构。那么，要想让新的人工智能模型更好地运行在新的芯片等计算机硬件上，就需要全新的 AI 编译器。

对此，微软亚洲研究院的研究员们和国内外合作者围绕着 AI 编译器的核心问题展开了一系列研究工作，并陆续推出了 AI 编译界的“工业重金属四部曲”：Rammer、Roller、Welder、Grinder，为当前主流的人工智能模型和硬件编译提供了系统性的创新解决方法。

基于统一切块（tile）抽象的四个核心 AI 编译技术

深度神经网络（DNN）是当前图像分类、自然语言处理和许多其他人工智能任务中广泛采用的方法。由于其重要性，许多计算设备，如 CPU、GPU、FPGA 和专门设计的 DNN 加速器被用来执行 DNN 计算。其中影响 DNN 计算效率的关键因素之一是调度，即决定在目标硬件上执行各种计算任务的顺序。现有的 DNN 框架和编译器通常将数据流图（data flow graph）中的 DNN 算子视为不透明的库函数，并将它们调度到加速器上单独执行。同时，这一过程还依赖于另一层调度器（通常在硬件中实现）来利用算子中可用的并行特性。这样的两层方法就导致了显著的调度开销，并且通常不能充分利用可用的硬件资源。

为此，研究员们提出了一种新的 DNN 编译器 Rammer，它可以优化 DNN 工作负载在大规模并行加速器上的执行。事实上，我们可以将进行 AI 编译时的调度空间想象成一个二维空间，并将计算任务看作是可以被拆分成不同大小和形状的“砖块”，调度的目的就是在二维空间的计算单元上将这些“砖块”像垒墙一样紧密排列起来，最大程度地利用计算单元不留空隙，因为一旦出现空隙，不仅已有空间得不到有效利用，而且还会影响“垒墙”的速度。而 Rammer 就是这个二维空间中的一台“夯土机”，在将 DNN 程序翻译成“砖块”后，可放置在芯片的不同计算单元上，将其压实。

Rammer 技术框架

换言之，Rammer 在编译时为 DNN 生成了有效的静态时空调度，最大限度地减少了调度开销。同时，通过为计算任务和硬件加速器提出的几个新的、与硬件无关的抽象，使 Rammer 获得了更丰富的调度空间，实现了算子间和算子内的协同调度，从而可以全面利用并行性。这些新颖且具有启发式的方法，让 Rammer 可以更好地探索空间并找到有效的调度，大幅提高硬件利用率。

研究员们在 NVIDIA GPU、AMD GPU 和 Graphcore IPU 等多个硬件后端对 Rammer 进行了测试。实验表明，Rammer 在 NVIDIA 和 AMD GPU 上的性能显著优于 XLA 和 TVM 等最先进的编译器，加速比高达20.1倍。与 NVIDIA 的专有 DNN 推理库 TensorRT 相比，加速比达3.1倍。

在计算机芯片上不仅有并行计算单元，还有多层内存，一个大的计算任务需要一层一层地向上传递，并在这个过程中将任务逐层切分成更小的“砖块”，最终交给最上层的处理器进行计算。这其中的难点在于如何把大的“砖块”铺满内存空间，进而更好地利用内存并提升效率。目前已有的方法是通过机器学习进行搜索，寻找更好的“砖块”切分策略，但是这通常需要数千个搜索步骤，每个步骤都要在加速器中进行评估，以找到合理的解决方案，所以这会花费大量的时间，如编译一个完整模型甚至需要几天或几周。

研究员们认为，在了解了计算逻辑和各内存层的参数，也就是在已知软件和硬件信息的情况下，其实完全可以估算出“砖块”切割的最佳方法和大小，从而实现更快的编译。这也是 Roller 的设计思路，它相当于一台压路机，在考虑内存特性的前提下，像铺地板一样把高维的张量数据平铺到二维的内存中，找到最优的切块（tile）大小。同时，它还封装了与底层加速器的硬件特性一致的张量形状，通过限制形状选择来实现高效编译。

Roller 技术框架

通过对6种主流 DNN 模型和119种流行的 DNN 算子的评估表明，Roller 可以在几秒内生成高度优化的内核，尤其是对于大型昂贵的自定义算子。最终，Roller 在编译时间上比现有的编译器实现了三个数量级的改进。Roller 生成的内核性能与包括 DNN 库在内的最先进的张量编译器的性能相当，有些算子甚至表现更好。与此同时，Roller 也已被用于微软内部开发的自定义 DNN 内核上，在实际开发中验证了 Roller 可以显著加快开发周期的优越性能。

现代 DNN 模型对高速内存的要求变得越来越高，在分析了一些最新的 DNN 模型后，研究员们发现当前大部分 DNN 计算的瓶颈主要在于 GPU 的访存，如这些模型对内存带宽利用率高达96.7%，但计算核的平均利用率只有51.6%，而且随着硬件与 DNN 模型的不断发展，这两者之间的差距还会持续增大。尤其是当前的人工智能模型需要处理高保真度的数据，如更大的图像、更长的句子、更高清的图形，这些数据在计算中都占用了更多的内存带宽。同时，更高效的专有计算核（如TensorCore）也进一步加大了内存压力。

为了解决内存问题，研究员们提出了 Welder 深度学习编译器，全面优化由通用算子组成的端到端 DNN 模型的内存访问效率。其实，DNN 模型可以看作是由多个算子连成的一张图，整个计算过程涉及多个阶段，即数据需要流过不同的算子，在每个阶段都需要将张量切分成块，先搬运到处理器上进行计算，然后再搬运回内存，这就会造成很大的搬运开销。由于整个计算过程包含多个流程，所以还可以将这一过程想象成逐层向上搬运“砖块”的场景，其中第一个“工人”将“砖块”拿上去加工然后再放回去，第二个“工人”再拿上来雕刻一下再放回去，然后是第三个、第四个……，反复搬运，可以预想其中的开销不言而喻。那么是否可以让第一个“工人”在顶层完成一部分子任务后直接交给下一个“工人”继续处理，然后再将多项任务“焊接”起来，实现流水化作业呢？Welder 正是扮演了电焊机这个角色，通过链接不同的算子，它可以让数据块以流水线的方式处理，大大降低了访存量，在近几年人工智能模型对访存效率要求越来越高的情况下，可以大幅提升计算效率。

Welder 技术框架

在对10个主流的 DNN 模型（包括用于各种任务的经典和最新模型结构，如视觉、自然语言处理、3D 图形等）进行评估后可以表明，Welder 在 NVIDIA 和 AMD 的 GPU 上都显著超过了现有的主流框架和编译器，如 PyTorch、ONNXRuntime 和 Ansor，速度提升分别到达21.4倍、8.7倍和2.8倍。Welder 的自动优化甚至超过了 TensorRT 和 Faster Transformer，最高可实现3.0倍和1.7倍的加速。此外，当在 TensorCore 等具有更快计算核心的硬件上运行这些模型时，其性能有了更大的提高，突显了内存优化对未来人工智能加速器的重要性。

在计算程序中，对数据块的搬运过程有时候需要引入一些更复杂的控制逻辑，这就是人工智能程序数据流之外的控制流，如循环地遍历一个句子中的每个单词，或者根据输入动态决定执行哪一部分程序。当前的编译器大多都是在解决数据流问题，对控制流的支持并不高效，因此导致了控制流较多的模型无法高效利用加速器性能。研究员们认为，可以将控制流和数据流切分重组以此来进行更高效的优化，并推出了 Grinder[1]。Grinder 好像一个便携的研磨切割机，它在把数据流切分成不同规模的并行计算块后，会再把控制流融入数据流，让控制流也能在加速器上高效执行。

Grinder 技术框架

Grinder 可以在硬件加速器上共同优化控制流和数据流的执行，并通过一种新的抽象来统一包括控制流和数据流的人工智能模型的表示，这就允许 Grinder 向较低级别的硬件并行性暴露用于重新调度控制流的整体调度空间。Grinder 使用启发式策略找到了有效的调度方案，且能够自动将控制流移动到设备内核中，进而实现了跨控制流边界的优化。实验表明，Grinder 可以对控制流密集的 DNN 模型加速8.2倍，是目前针对控制流的 DNN 框架和编译器中速度最快的一个。

基于同一套抽象和统一的中间表示层（Intermediate Representation，IR），这四款 AI 编译器解决了当前 AI 编译器中的不同问题——并行、编译效率、内存、控制流，构成了一套完整的编译解决方案。在推进研究的进程中，微软亚洲研究院的编译原型系统已经为 Office、Bing、Xbox 等微软产品的部署和模型优化提供了帮助，同时也在微软新型算子的定制和优化中发挥了作用。

“在大模型成为主流的今天，人工智能模型对效率、算力有了更高的要求。一方面，AI 编译器需要针对硬件资源做出极致的算子融合、定制和优化；另一方面，也需要对新型大规模硬件架构进行系统编译支持，如片上网络互联（NoC）的芯片、混合内存架构等，甚至通过白盒编译方法指导硬件定制。我们提出的这套 AI 编译器已被证明能够大幅提升 AI 编译的效率，可以更好地助力人工智能模型的训练和部署。同时，大模型的发展也为 AI 编译带来了机遇，未来大模型本身或许就可以帮助我们实现优化和编译。”微软亚洲研究院首席研究员薛继龙表示。

你也许还想看：