深度学习编译中间件之NNVM(五)TVM论文阅读

参考文档

1. https://mp.weixin.qq.com/s/irvBbPKENiZX9G_6wh5c-Q 陈天奇等人提出TVM：深度学习自动优化代码生成器
2. https://arxiv.org/abs/1802.04799v1 TVM: End-to-End Optimization Stack for Deep Learning

摘要：现今，像Tensorflow，MXNet，Caffe和Pytorch这样的可扩展框架共同驱动了深度学习的流行度和实用性。但是，这些框架只为一些服务器端GPU提供优化，并将工作负载部署到其他平台，如移动手机，嵌入式设备和其他特定的加速器(e.g., PFGAs, ASICs)，对于这些平台的优化适配是一个巨大工作量的任务。所以我们提出了TVM，一个端到端的优化堆栈，具备图形级和算子级的优化，以为多种硬件后端提供深度学习工作负载的性能可移植性。我们讨论了 TVM 所解决的深度学习优化挑战：高级算子融合（operator fusion）、多线程低级内存重用、任意硬件基元的映射，以及内存延迟隐藏。实验结果证明 TVM 在多个硬件后端中的性能可与适应低功耗 CPU 和服务器级 GPU 的当前最优库相比。我们还通过针对基于 FPGA 的通用深度学习加速器的实验，展示了 TVM 对新型硬件加速器的适应能力。该编译器基础架构已开源。

1 介绍

深度学习模型可以识别图像、处理自然语言，以及在部分具有挑战性的策略游戏中击败人类。在其技术发展的过程中，现代硬件稳步推进的计算能力扮演了不可或缺的作用。很多目前最为流行的深度学习框架，如 TensorFlow、MXNet、Caffe 和 PyTorch，支持在有限类型的服务器级 GPU 设备上获得加速，这种支持依赖于高度特化、供应商特定的 GPU 库。然而，专用深度学习加速器的种类越来越多，这意味着现代编译器与框架越来越难以覆盖所有的硬件。

显而易见，以现有的点到点方式实现不同深度学习框架对所有种类的硬件进行后端支持是不现实的。我们的最终目标是让深度学习负载可以轻松部署到所有硬件种类中，其中不仅包括 GPU、FPGA 和 ASIC（如谷歌 TPU），也包括嵌入式设备，这些硬件的内存组织与计算能力存在着显著的差异，如下图所示：

考虑到这种需求的复杂性，开发一种能够将深度学习高级程序降低为适应任何硬件后端的低级优化代码的优化框架是最好的方法。

目前的深度学习框架依赖于计算图的中间表示来实现优化，如自动微分和动态内存管理。然而，图级别的优化通常过于高级，无法有效处理硬件后端算子级别的转换。另一方面，目前深度学习框架的算子级别库通常过于僵化，难以轻松移植到不同硬件设备上。为了解决这些问题，我们需要一个可实现从计算图到算子级别的优化，为各种硬件后端带来强大性能的编译器框架。

1.1 优化的基本挑战

深度学习的优化编译器需要同时展示高级别与低级别的优化，在论文中，研究人员总结了在计算图级别与张量算子级别上的四大基本挑战：

1. 高层次数据流复写：不同的硬件设备可能具有截然不同的内存层次结构，因此，融合算子与优化数据布局的策略对于优化内存访问至关重要。

2. 跨线程内存复用：现代 GPU 与专用加速器的内存可被多个计算核心共享，传统的无共享嵌套并行模式已不再是最优方法。为优化内核，在共享内存负载上的线程合作很有必要。

3. 张量计算内联函数：最新的硬件带来了超越向量运算的新指令集，如 TPU 中的 GEMM 算子和英伟达 Volta 架构中的 Tensor Core。因此在调度过程中，我们必须将计算分解为张量算术内联函数，而非标量或向量代码。

4. 延迟隐藏（Latency Hiding）：尽管在现代 CPU 与 GPU 上，同时拥有多线程和自动缓存管理的传统架构隐藏了延迟问题，但专用的加速器设计通常是采用精简控制流和将复杂性分配到编译器堆栈上面的方案。所以设计涉及到隐藏内存访问延迟的调度器时必须要非常仔细。

1.2 TVM：一个端到端优化堆栈

下图展示了TVM的基本构成：

TVM是一个端到端优化堆栈，该端到端优化编译器堆栈可降低和调整深度学习工作负载，以适应多种硬件后端。TVM 的设计目的是分离算法描述、调度和硬件接口。该原则受到 Halide [22] 的计算／调度分离思想的启发，而且通过将调度与目标硬件内部函数分开而进行了扩展。这一额外分离使支持新型专用加速器及其对应新型内部函数成为可能。TVM 具备两个优化层：计算图优化层，用于解决第一个调度挑战；具备新型调度基元的张量优化层，以解决剩余的三个挑战。通过结合这两种优化层，TVM 从大部分深度学习框架中获取模型描述，执行高级和低级优化，生成特定硬件的后端优化代码，如树莓派、GPU 和基于 FPGA 的专用加速器。该论文做出了以下贡献：

• 我们构建了一个端到端的编译优化堆栈，允许将高级框架（如 Caffe、MXNet、PyTorch、Caffe2、CNTK）专用的深度学习工作负载部署到多种硬件后端上（包括 CPU、GPU 和基于 FPGA 的加速器）。

• 我们发现了提供深度学习工作负载在不同硬件后端中的性能可移植性的主要优化挑战，并引入新型调度基元（schedule primitive）以利用跨线程内存重用、新型硬件内联函数和延迟隐藏。

• 我们在基于 FPGA 的通用加速器上对 TVM 进行评估，以提供关于如何最优适应专用加速器的具体案例。

我们的编译器可生成可部署代码，其性能可与当前最优的特定供应商库相比，且可适应新型专用加速器后端。

2 优化计算图

2.1 计算图

在深度学习框架中计算图是一个表示程序的常用方式。下图展示了一个使用计算图表示一个2层的卷积神经网络：

高层次表示和底层的编译器中间表示(IR，像LLVM)的主要不同在于中间数据是大数据量的多维Tensor。TVM采用了在一个计算图表示上使用高层次优化的方案：node代表tensor操作，edges代表tensor之间的数据流向和数据依赖关系。

计算图为计算任务提供了一个全局视图，同时也避免了描述每个计算任务具体是如何实现的。在一个图中，静态内存计划Pass能够通过预分配所有中间结果Tensor来处理。这个分配阶段和传统编译器中的寄存器分配Pass类似。和LLVM IR类似，一个计算图应该能够被转换成函数公式图。例如，一个常数折叠Pass能够在计算图预计算阶段被静态地执行，以节省运行时代价。下图展示了在一些TVM中实现的新颖的图层次优化：操作符融合和数据布局转换。

2.2 操作符融合

对于GPU和特定加速器而言，将多次操作融合在一起的优化方法能较为明显地降低执行时间。操作符融合的想法是来源于单个Kernel函数会节省将中间结果写回全局内存的时间消耗。从具体分析来看，我们总结了四种类别的图操作符：

• injective(one-to-one map)：单射，如add/sqrt/sub
• reduction：约简，如sum/max/min
• complex-out-fusable(can fuse element-wise map to output)，如conv2d
• opaque(cannot be fused)

下图展示一个使用操作符融合进行图优化的示例：

下图展示了3种工作负载经过操作符融合之后的性能对比：

2.3 数据布局转换

Tensor操作是计算图的基本操作符，Tensor中涉及到的运算会根据不同的操作符拥有不同的数据布局需求。例如，一个深度学习加速器可能会使用4x4张量操作，所以需要数据切割成4x4的块来存储以优化局部访存效率。下图展示了一个矩阵如何布局，这种布局能够适应计算2x2张量操作：

对于优化数据布局而言，需要为每个操作符提供定制的数据布局。如果在producer(生产者)和consumer(消费者)之间出现了数据布局不匹配的情况时，此时需要我们进行数据布局转换。

2.4 计算图级别优化的限制

虽然高层次的数据流图优化能够大大提升深度学习的计算效率，但是它们和现有的Operator库相关。当前，只有少量的深度学习框架支持操作符融合。随着越来越多支持正交基的操作符被引进，能够被融合的Kernel数量正在急剧增大。但是当出现越来越多的硬件后端时，为不同的数据布局方式、数据类型和硬件内联函数手动适配的方案已经变得不可行的。为此，我们提出了一种代码生成的方案，将在下一章介绍这种方案。

3 优化张量操作

此章描述了如何为广泛的硬件后端生成优良的操作符版本代码

3.1 Tensor表达式语言

这里我们介绍一种能够进行自动代码生成的数据流Tensor表达式语言。和高层次计算图语言不同，Tensor操作的实现是不透明的(这里可以理解为Tensor描述和实现是分离的)。每个操作被描述为一个数学表达式，例如：

注解：

1. Recurrence：循环
2. cumsum：矩阵元素累计和

我们的Tensor表达式语言借鉴了一些已有的编程语言，像Halide，Darkroom和TACO。我们的Tensor表达式语言支持一般的算术和数学操作符。我们也引进了满足交换律的Reduction操作符，这个操作符能够轻松地进行跨线程调度。之后我们还会引进一个高层次的扫描操作符，这个操作符能够把基础计算操作符联合成窗体循环计算。TVM计算操作符也支持Tensor元组(可以理解多个Tensor输入或者输出)，这样有利于支持像argmax这种函数。总结陈述一下，所有Tensor操作被表示为高层次数据流图并且覆盖深度学习常用的计算模式。

3.2 调度空间

上一节已经陈述了如何描述Tensor表达式，但是如何为硬件后端创建高性能的实现依然是具有挑战性的。这里展示一个为CPU/GPU/深度学习加速器进行典型优化的示例：

上面的示例为矩阵乘法的CPU/GPU/深度学习加速器的实现

每个经过底层优化的程序针对不同硬件后端采用不同的调度策略的联合，这也为Kernel设计者带来了很大的负担。所以通过借鉴Halide，我们采用了解耦计算描述和调度器优化两个过程。调度器会根据硬件后端采用特定的规则将计算描述向下转换成已优化的硬件实现。TVM的调度空间可以展示在下图中：

这个示例展示了一个Tiled(使用分块)的矩阵乘法

为了更好地快速展示调度空间，我们需要提供有效的调度器基元。下图将展示一系列TVM常用的调度器基元，这些基元经过实践调整之后具有很高的计算性能：

我们采用一些Halide中有效的调度器基元，另外引进了一些新的调度器基元Tensorization/Virtual Thread。接下来将会详细介绍这些调度器基元。

3.3 协作式嵌套并行化

针对深度学习工作负载而言，并行编程的关键是改善计算密集型Kernel的计算效率。现代GPU提供了大规模的并行化，需要我们将并行编译模型加入到调度器里面。大部分现有的解决方案都是采用了嵌套并行编译的方法，是一种fork-join的并行模式。具体而言，我们能使用一个并行调度器基元来处理一个数据并行任务。每个并行任务都能够被递归地细分到一个子任务以实现多级线程。

我们称这种模型为 shared-nothing nested parallelism，即一个工作线程不能观察到相同计算阶段内的相邻线程的数据。相邻线程之前的交互只能发生在join阶段，join发生在子任务结束并且下个阶段所需的数据已经准备好时。这种编程模型禁止线程在相同计算阶段内进行协作。下图展示了一个遵循这种限制的矩阵乘法示例：

在一个GPU上矩阵乘法一般会被切割成小分块，每个分块指定一个线程组。如果使用shared-nothing nested parallelism的话，每个线程在Reduction阶段就需要保证获取的数据是独立无依赖的。

一个更好的替代方案是线程在获取数据时采用shared-nothing的手段。这种模式被使用在一些流行度非常高的GPU编程语言中，例如CUDA，OpenCL和Metal，但是没有一个调度器基元采用这种手段。我们将memory scopes的概念引进到调度器空间，可以标记一个阶段是可共享的。如果不指定memory scopes，自动Scope推理机制将会标记相关的阶段是thread-local的，如下图所示：

共享的任务需要计算所有组内工作线程的依赖关系，我们可以通过分发加载任务到一组线程从而有效地的调度数据加载任务。我们强制使用相同的线程工作在共享任务的一个专区上是非常不值得的，需要保证线程能持续在加载阶段和计算阶段之间切换。这种改进的实现需要额外的编译器支持。具体来说，需要设计一个处理范围约束的推理算法，这个算法能够推断出共享任务的范围，能够把具有协作关系的线程合并在一起。另外，内存同步屏障也需要正确地被插入以确保这些共享的数据对数据的消费者可见。

下图展示了shared-nothing nested parallelism和with cooperation在GPU上面的性能对比，另外也加入了TVM和Halide的对比数据：

从上面的性能对比数据来看，我们发现采用新调度器基元的方案在GPU上取得了最佳的性能。最后，上面介绍的memory scopes的概念也可以引进到针对特定深度学习加速器的代码生成过程中。

3.4 张量化：生成硬件接口

深度学习工作负载拥有比较高的计算密度，这些计算可以被典型地分解成Tensor操作，像矩阵-矩阵乘法或一维卷积。这些自然的分解引领了当前添加Tensor操作基元的趋势，这些新兴的计算原型是相当多种多样的，包括矩阵-矩阵乘法、矩阵-向量点乘、一维卷积等等。这些新的基元对于Tensor操作调度提出了新的挑战：调度器必须使用这些基元从而获得它们的加速特性。我们称之为Tensorization问题，类似于SIMD架构下的向量化问题。

Tensorization显著地不同于Vectorization。Tensor计算基元的输入是多维数据，可能是固定长度或者变长，并且存在不同的数据布局。更重要的是我们不能依靠于一组固定的基元，因为新的深度学习加速器采用特殊的Tensor指令。因此我们需要一个解决方案用来支持未来新的特定的加速器。

为了解决上面的问题，我们从调度器中分离了硬件接口。具体来说，我们引进了一套Tensor内联声明机制。我们可以使用Tensor表达式语言来声明每个新的硬件内联的行为，和给他分配底层原语是一样的。另外我们引进了一个调度器基元Tensorization来作为基元的计算单元。下图展示了一个Tensorization的示例：

Tensor表达式语言能够同时描述用户准备的计算描述信息和硬件暴露接口的抽象信息。Tensorization把调度器从硬件基元中解耦出来，这样能够使TVM更容易扩展新的硬件架构。Tensorization调度器生成的代码经过实践具有很高的计算性能：将复杂的操作分解成一系列重复的微型Kernel调用。因此我们能够使用Tensorization基元来发挥手工制作的汇编微型Kernel的性能优势，这种方法在一些硬件平台上是非常有益的。例如，在AMD Vega GPU上，通过Tensorizing半精度GEMM到一个手动制作的4x4微型Kernel获得1.5倍的性能加速效果。

3.5 编译器的延迟隐藏支持

延迟隐藏指的是针对一个在计算过程中重叠内存操作的过程最大化访存和计算利用率。针对不同的硬件它需要采用不同的策略。对于CPU而言访存延迟隐藏通过同步多线程或者硬件预取技术达到。对于GPU而言，则通过成千上万的线程组快速切换达到访存延迟隐藏。对于特定深度学习加速器，经常会倾向于使用精简控制流和编译器堆栈。

下图展示了编译器如何在一个流水线式深度学习加速器中显式地处理数据依赖：

上面的示例遵循了运算/访存分离的哲学。

4 代码生成和运行时支持

经过调度器优化之后，接下来的任务就是生成能运行在特定硬件平台下的代码，并且方便已生成Kernel的部署和性能分析。

4.1 代码生成

对于一个特定的多元数据流声明，坐标系相关的超图和调度树，我们可以通过迭代遍历调度树的方案生成Lowered代码，并且推断出输入Tensor的依赖范围(使用坐标系相关的超图)，接下来生成循环嵌套的low-level代码。low-level代码是通过类C的循环程序具象化而来，在这个过程中我们使用了一个Halide循环程序数据结构的变体，我们也使用了Halide常用的lowering primitives，像是storage flattening，循环展开，对于GPU和特定深度学习加速器而言，则使用了同步点检测、虚拟线程注入、模块化生成机制。最终，循环程序被翻译成LLVM/CUDA/Metal/OpenCL源代码。

4.2 运行时支持

对于GPU程序而言，我们会独立的构建host端和device端的模块并且提供一个运行时模块系统来启动Kernel。我们也为基于FPGA的深度学习加速器构建了一套驱动程序，这套驱动程序使用C语言API，可以实现构造指并推送到目标加速器上执行。我们的代码生成算法然后将加速器程序翻译成一系列Runtime API。

4.3 自动调优

本篇论文着重于提供一套新颖的优化框架，这套框架能够为深度学习系统编译高性能底层实现。我们已经完成的优化效果给我们展示了对于调度器生成高性能代码而言还有很大的值得探索的空间。我们正在探索一个非常早期的自动调度器技术和一个自动内联注入操作的Pass。像高维卷积、矩阵乘法和深度卷积这样的复杂操作可以被一套调度器模板进行自动调优。我们相信结合我们更精细化的优化技术，在未来还会有性能改善。

观点：虽然从现阶段来看自动调优的性能还不是很理想，但是能预测到硬件性能的自动调优可能是未来的一个大方向。所以fackbook也开源了TensorComprehensions，可以预想到不久之后自动调优很有可能会超过人类专家手动调优的水平。

4.4 远程部署性能分析

为了嵌入式设备，TVM设计了一套方便性能分析和自动调优的基础设施。传统情况下，嵌入式开发一般是在主机上进行交叉编译然后复制可执行文件到目标设备上执行，上述编译、运行和性能分析的工作都需要手动进行。在编译器堆栈中我们提供了一个远程程序调用：通过RPC接口，我们能够完全在host端完成上面的所有步骤，这样的方式可以极大地加快在嵌入式设备和基于FPGA的加速器上的优化工作。

5 评估

我们主要在如下硬件平台对TVM进行了评估：

1. 嵌入式ARM CPU
2. 服务端NVIDIA GPU
3. FPGA深度学习加速器

性能评估程序是基于真实的深度学习工作负载，包括ResNet和MobileNet。并且选取了MxNet和Tensorflow作为对比框架。

5.1 树莓派3B评估

我们在树莓派3B(4核Cortex-A53 1.2Ghz)上面评估了TVM的性能，我们使用了MxNet作为基准系统。

5.2 NVIDIA Tesla K80和GTX1080评估

待补充…

5.3 PYNQ FPGA开发板评估

待补充…

6 相关工作

Tensorflow/CNTK/Theano/MxNet这些深度学习框架为用户提供了便捷的深度学习工作负载体验接口，可以容易地将深度学习模型部署到不同的硬件平台。然而现有的深度学习框架都依赖于厂商独立定制的Tensor操作库，现在这些框架能利用TVM软件堆栈为大量种类的硬件设备生成优化良好的底层实现代码。

高层次计算图DSLs也是一个表示和执行高层次优化的典型方法。Tensorflow的XLA和之前介绍的DLVM都是属于这个类别。在这些工作中，计算图的表示方法是类似的，另外在本篇论文中我们也介绍我们的高层次计算图DSL。当图层次的表达能够很好地和高层次优化匹配的话，也意味着这些DSL对于大量不同种类的硬件后端Tensor操作层面来说已经层次太高了。之前的工作都是依赖于定制lowering规则来直接生成底层的LLVM中间码或者硬件厂商手工指定的库。这些方法为了做到多后端设备兼容需要花费大量的工程化投入。

Halide比较明智地采用了计算和调度分离的策略，我们顺理成章地在我们的编译器中采用了Halide的内在思想和现有的有效调度器基元。Tensor操作调度也参考了其他一些为GPU设计的DSL，这些DSL采用了polyhedral-based loop transformation(基于多面体的循环转换)的方案。另外TACO这款软件提出了一种在CPU上实现生成稀疏Tensor操作的通用方法。Weld是一个用来描述数据处理任务的DSL。我们主要集中处理一些在GPU和特定加速器上的深度学习工作负载的调度挑战。我们的调度器基元有可能被这些工作负载的任务所采用。更重要的是，我们提供了一个端到端的软件堆栈，能够直接从深度学习框架的描述文件生成优化良好的目标实现代码。

当前的比较流行的趋势是特定领域的定制深度学习加速器的大量出现，现在还没有编译器堆栈去适配这些新的硬件设备。VITA被设计成可以提供了对这些硬件加速器的通用支持。本篇论文为这些特定加速器提供了一个通用有效的解决方案。

7 致谢

原文略，不过值得赞赏的是TVM得到了陈天奇的谷歌博士奖学金的支持。

8 结论

我们的工作是提供了一个端到端的软件堆栈，可以解决多种硬件后端下的性能优化挑战。我们希望我们的工作能够促进更多关于编程语言、编译的研究。另外TVM也带来了深度学习系统的软硬件协同设计的机会，我们已经开源了我们的TVM软件堆栈，还将开源VITA加速器以鼓励在这个方案更有意义的研究出现。