TVM- End-to-End Optimization Stack for Deep Learning

TVM- End-to-End Optimization Stack for Deep Learning

引言

TensorFlow、MXNet、Caffe 和 PyTorch 等可扩展框架推动了深度学习当前的普及和实用性。然而，这些框架针对范围较窄的服务器级 GPU 进行了优化，将工作负载部署到其他平台（例如手机、嵌入式设备和专用加速器（例如 FPGA、ASIC））需要费力的手动工作。我们提出 TVM，这是一个端到端的优化堆栈，它公开图形级别和运算符级别的优化，为跨不同硬件后端的深度学习工作负载提供性能可移植性。我们讨论了 TVM 解决的深度学习特有的优化挑战：高级运算符融合、跨线程的低级内存重用、映射到任意硬件基元以及内存延迟隐藏。实验结果表明，TVM 可提供跨硬件后端的性能，可与用于低功耗 CPU 和服务器级 GPU 的最先进库相媲美。我们还通过以基于 FPGA 的通用深度学习加速器为目标，展示了 TVM 以新硬件加速器后端为目标的能力。编译器基础设施和 FPGA 加速器设计将开源。

1.简介

深度学习模型现在可以识别图像、处理自然语言并在具有挑战性的策略游戏中击败人类。现代硬件稳步提升的计算能力在深度学习目前在许多问题领域的普遍性和相关性中发挥了突出作用。许多最流行的深度学习框架，如 TensorFlow、MXNet、Caffe 和 PyTorch，通过将支持集中在一小类服务器级 GPU 设备上来利用现代硬件的力量——这种支持取决于高度工程化的使用和供应商特定的 GPU 库。然而，专业深度学习加速器的数量和多样性正在迅速增加。这些加速器带来了采用挑战，因为它们引入了现代编译器和框架无法处理的新硬件抽象。以目前的临时方式为各种硬件后端提供各种深度学习框架的支持是不可持续的。最终，目标是将深度学习工作负载轻松部署到各种硬件目标，包括嵌入式设备、GPU、FPGA 和 ASIC（例如 TPU），这些硬件在内存组织、计算等方面存在显着差异，如下所示图 1. 鉴于这些要求，开发优化框架以将深度学习程序的高级规范降低为适用于任何硬件后端的低级优化代码至关重要。

当前的深度学习框架依赖于计算图中间表示来实现自动微分和动态内存管理等优化 [3, 7, 4]。然而，图级优化通常太高级而无法处理特定于硬件后端的运算符级转换。另一方面，深度学习框架所依赖的当前算子级库过于僵化和专业化，无法轻松跨硬件设备移植。为了解决这些弱点，我们需要一个编译器框架，它可以在图形和运算符级别上提供优化机会，以在硬件后端提供有竞争力的性能

基本优化挑战 用于深度学习的优化编译器需要公开高级和低级优化。我们总结了计算图级别和张量算子级别的四个基本挑战：

1. 高级数据流重写：不同的硬件设备可能具有截然不同的内存层次结构，因此启用融合运算符和优化数据布局的策略对于优化内存访问至关重要。
2. 跨线程内存重用：现代 GPU 和专用加速器具有可以跨计算核心共享的内存。传统的无共享嵌套并行模型不再是最优的。优化内核需要线程之间在加载的共享内存上进行协作。
3. 张量计算内在函数：最新的硬件提供了超越向量运算的新指令，例如 TPU 中的 GEMM 运算符或 NVIDIA Volta 中的张量核心。因此，调度程序必须将计算分解为张量算术内在函数而不是标量或向量代码。
4. 延迟隐藏：虽然具有同步多线程和自动管理缓存的传统架构隐式地隐藏了现代 CPU/GPU 中的延迟，但专门的加速器设计通常倾向于更精简的控制，并将大部分调度复杂性卸载到编译器堆栈。尽管如此，必须仔细执行调度以隐藏内存访问延迟。

TVM：端到端优化堆栈 我们展示了 TVM（如图 2 所示），这是一种端到端优化编译器堆栈，可降低和微调深度学习工作负载以适应不同的硬件后端。 TVM 旨在分离算法描述、调度和硬件接口。该原则受到 Halide [21] 的计算/调度分离的启发，但通过将调度与目标硬件内在函数分离来扩展该概念。这种额外的分离可以支持新型专用加速器及其相应的新内在函数。

TVM 提供了两个优化层：

一个计算图优化层来解决第一个调度挑战，

一个带有新调度原语的张量优化层来解决其余三个挑战。

通过组合这些优化层，TVM 可以从大多数深度学习框架中获取模型描述，执行联合高级和低级优化，并为后端（例如 Raspberry Pi、GPU 和 FPGA）生成特定于硬件的优化代码。基于专门的加速器。

我们的论文做出了以下贡献：

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

• 我们确定了为跨不同硬件后端的深度学习工作负载提供性能可移植性的主要优化挑战。我们引入了新颖的调度原语以利用跨线程内存重用、新颖的硬件内在函数和延迟隐藏。

• 我们在基于 FPGA 的通用加速器上评估 TVM，以提供有关如何最佳定位专用加速器的具体案例研究。我们的编译器生成的可部署代码与最先进的供应商特定库在性能上具有竞争力，并且可以针对新的专用加速器后端。

在本文的其余部分安排如下。我们在第 2 节中描述了图形优化，在第 3 节中描述了张量运算符调度。我们在这两节中包含了实验结果，以定量评估我们提出的优化。我们在第 4 节讨论运行时支持。我们的端到端评估在第 5 节。相关工作在第 6 节讨论。

2. 优化计算图

计算图： 计算图是深度学习框架中表示程序的常用方法[3、6、7、4]。图 3 显示了一个两层卷积神经网络的示例计算图表示。

在这种高级表示和低级编译器 IR（例如 LLVM）之间，中间数据项是大型多维张量。 TVM 利用计算图表示来应用高级优化：节点表示对静态已知维度的张量的操作，边表示张量操作之间的数据流数据依赖性。

计算图提供了计算任务的全局视图，但避免指定每个计算任务需要如何实现。可以在图上执行静态内存规划过程，以预分配内存来保存每个中间张量结果。这个分配阶段类似于传统编译器中的寄存器分配阶段。与 LLVM IR 类似，可以将计算图转换为功能等效的图以应用优化。例如，可以应用常量折叠传递来预先计算可以静态确定的图形部分，从而节省执行成本。图 4 概述了在 TVM 中实现的新颖图级优化：运算符融合和数据布局转换。

**运算符融合：**将多个运算符融合在一起是一种优化，可以大大减少执行时间，特别是在 GPU 和专用加速器中。这个想法是将多个运算符组合成一个内核，而不将中间结果保存回全局内存。具体来说，我们识别四类图运算符：单射（一对一映射）、缩减、复杂可融合（可以将逐元素映射融合到输出）和不透明（不能融合）。我们应用图 5 中列出的规则将计算图转换为融合版本。图 6 通过比较融合版本和非融合版本在三种工作负载中的性能，展示了这种优化的影响。

Data Layout Transformation 张量运算是计算图的基本运算符。计算中涉及的张量在不同的操作中可能有不同的布局要求。例如，深度学习加速器可能利用 4 × 4张量化操作，要求将数据平铺成 4×4 块以优化访问局部性。图 7 显示了如何转换矩阵的数据布局以适应对其数据进行 2 × 2 张量化操作。优化数据布局从指定每个运算符的首选数据布局开始，因为约束决定了它们在硬件中的实现。如果数据布局不匹配，我们将在生产者和消费者之间执行适当的布局转换

图级优化的局限性 虽然高级数据流图优化可以极大地提高深度学习工作负载的效率，但它们的有效性取决于算子库提供的功能。目前，少数支持算子融合的深度学习框架需要算子库提供融合模式的实现。随着定期引入更多网络运营商，可能融合内核的数量会急剧增加。当以越来越多的硬件后端为目标时，这种方法不再可持续，因为所需的融合模式实现数量与需要支持的数据布局、数据类型和硬件加速器内在函数的数量一起增长。为后端特定运算符的巨大空间手工制作运算符内核是不可行的。为此，我们提出了一种代码生成方法，可以在下一节中生成张量运算符。

3. 优化张量操作

本节介绍 TVM 如何为各种硬件后端生成同一运算符的微调版本。

3. 1 张量表达语言

我们引入了一种数据流张量表达式语言来支持自动代码生成。与高级计算图形语言不同，张量运算的实现是不透明的，每个运算都用索引公式表达式语言描述，如图 8 所示。我们的张量表达式语言借鉴了 Halide [21]、Darkroom 等语言的线索[13] 和炸玉米饼 [18]。我们的张量表达式语言支持 C 等常见语言中常见的算术和数学运算。我们明确引入了交换归约运算符，以轻松安排跨多个线程的交换归约。我们进一步介绍了一种高阶扫描运算符，它可以结合基本的计算运算符来随着时间的推移形成循环计算。 TVM 计算操作还支持张量元组之间的缩减，从而可以轻松支持 argmax 等函数。这种表示可以描述高级数据流图中使用的所有张量操作，并涵盖深度学习中展示的常见模式。

3.2 排程空间

给定张量表达式，为每个硬件后端创建高性能实现仍然具有挑战性。

例如，图 9 突出显示了应用于 CPU、GPU 和深度学习加速器设计的典型优化。每个优化的低级程序都是调度策略不同组合的结果，给内核编写者带来了很大的负担。我们采用从 **Halide [21] 的调度优化中解耦计算描述的原则。**调度是将计算描述降低到后端优化实现的特定规则。这个想法是对调度空间和用于遍历该空间的转换进行正式建模，从而提供生成低级代码的不同方法。 TVM 的调度空间如图 10 所示。为了快速探索调度空间，我们需要提供有效的调度原语来转换调度。图 11 显示了 TVM 中使用的一组通用调度原语。许多这些原语与高性能计算中的实践相呼应。我们采用了 Halide 中有用的调度原语并引入了新的调度原语来应对 GPU 和专用硬件加速器带来的挑战。我们将在接下来的三个小节中详细描述这些原语。

3.3 嵌套并行与协作

并行编程是提高深度学习工作负载中计算密集型内核效率的关键。现代 GPU 提供大规模并行性，要求我们将并行编程模型烘焙到调度转换中。大多数现有解决方案都采用称为嵌套并行程序的并行编程模型，这是一种 fork-join 并行形式。具体来说，我们可以使用并行调度原语来并行化数据并行任务。每个并行任务可以进一步递归细分为子任务，以利用目标架构上的多级线程层次结构（例如，GPU 中的线程组）

我们称此模型为无共享嵌套并行，因为一个工作线程无法在同一并行计算阶段查看其同级线程的数据。兄弟线程之间的交互发生在连接阶段，当子任务完成并且下一阶段可以使用前一阶段产生的数据时。这种编程模型不允许线程相互协作以在同一并行阶段执行集体任务。图 12 提供了一个矩阵乘法示例来说明此限制。 GPU 上的矩阵乘法需要将工作划分为分布在线程组中的块，然后分配给每个单独的线程。在无共享嵌套并行性下，每个线程都必须在缩减阶段独立获取其计算块所需的数据。 1 无共享方法的更好替代方法是跨线程协作获取数据。这种模式在使用 CUDA、OpenCL 和 Metal 等语言的 GPU 编程中广为人知，但尚未在调度原语中实现。我们将内存范围的概念引入到调度空间中，以便可以将一个阶段标记为共享。如果没有内存范围，自动范围推断会将相关阶段标记为线程本地，如图 12 所示。共享任务需要计算组中所有工作线程的依赖关系。我们可以通过使用线程绑定原语在同一组线程中分配加载任务来高效地安排输入数据加载。值得注意的是，我们被迫使用相同的线程来处理除法共享工作负载，因此线程可以在加载和计算阶段持续存在。这种改进的实现需要额外的编译器支持。具体来说，边界推理算法需要能够通过将所有协作线程的任务合并在一起来推断共享任务的边界。此外，需要正确插入内存同步屏障，以保证共享加载的数据对消费者可见。图 14 比较了使用无共享嵌套并行生成与协作生成的 GPU 内核的性能。我们还将 TVM 与 Halide 进行了比较，后者采用无共享嵌套并行方法。我们发现采用这些新的调度原语对于在 GPU 上获得最佳性能至关重要。最后，除了对 GPU 有用之外，内存作用域还允许我们标记特殊的内存缓冲区，并在针对专门的深度学习加速器时创建特殊的降低规则.

3.4 张量化：泛化硬件接口

深度学习工作负载具有很高的算术强度，通常可以分解为张量运算符，如矩阵-矩阵乘法或一维卷积。这些自然分解导致了最近的趋势，即添加超越向量指令的张量计算原语。新兴的计算内在函数非常多样化，包括矩阵-矩阵乘法 [17]、矩阵-向量积 [1] 和一维卷积 [9] 等示例。这些新原语为张量运算符调度带来了新挑战：调度必须使用这些原语才能从加速中获益。我们将其称为张量化问题，类似于 SIMD 架构的矢量化问题。张量化与矢量化有很大不同。张量计算原语的输入是多维的，具有固定或可变长度，并规定了不同的数据布局。更重要的是，我们不能求助于一组固定的原语，因为新的深度学习加速器正在出现，它们都有自己的张量指令。因此，我们需要一种经得起未来考验的解决方案来支持新一代的专用加速器。为了解决这一挑战，我们将硬件接口与时间表分开。具体来说，我们引入了张量内在声明机制。我们可以使用张量表达式语言来声明每个新硬件内在的行为，以及与之相关的降低规则。此外，我们引入了张量化调度原语，以用相应的张量内在函数替换计算单元。编译器将计算模式与硬件声明相匹配，并将其降低为相应的硬件内在函数。图 15 显示了一个张量化示例。张量表达式语言既描述了用户的预期计算描述，也描述了硬件公开的抽象。张量化将时间表与特定硬件原语分离，使 TVM 易于扩展以支持新硬件架构。生成的张紧调度代码与高性能计算中的常见做法一致：将复杂的操作分解为重复的微内核调用序列。因此，我们还可以使用 tensorize 原语来利用手工制作的汇编微内核，这在某些平台上可能是有益的。例如，在 AMD Vega GPU 上将半精度 GEMM 张量为 4 × 4 手工制作的微内核可以产生超过最佳非张量版本 1.5 倍以上的加速.

3.5 编译器支持延迟隐藏

延迟隐藏是指将内存操作与计算重叠以最大化内存和计算利用率的过程。它需要不同的策略，具体取决于目标硬件后端。在 CPU 上，内存延迟隐藏是通过同步多线程 [11] 或硬件预取技术 [16、8] 隐式实现的。 GPU 依靠许多线程扭曲的快速上下文切换来最大化功能单元的利用率 [25]。另一方面，专门的深度学习加速器通常支持更精简的控制，并将此问题卸载到编译器堆栈。图 16 演示了编译器如何显式处理流水线深度学习加速器的数据依赖性，该加速器遵循解耦访问/执行理念 [22、17]。我们假设硬件流水线由可以同时执行的内存和计算阶段组成，其中每个阶段都由独立的指令流控制。显式同步指令用于在流水线阶段之间发送信号以指示给定任务何时完成，以便下一个相关阶段可以开始使用或覆盖数据。这种类型的显式依赖跟踪是在具有 FIFO 队列的硬件中实现的，并为程序员和编译器提供了对硬件任务执行方式的直接控制。通过使用硬件公开的低级同步原语，我们可以有效地隐藏延迟。在如此低的层次上对具有显式同步操作的硬件进行编程是一项艰巨而艰苦的任务。为了减轻程序员的负担，我们提供了一个虚拟线程调度原语，让程序员指定一个高级数据并行程序，TVM 自动将其降低为低级显式数据依赖程序。降低过程如图 17 所示。该算法从高级并行程序开始，然后插入必要的同步指令以保证每个线程内操作的正确执行顺序。最后，所有虚拟线程的操作交织到一个线程中。

降低虚拟线程的正确性 虚拟线程的不正确交错会导致死锁。从概念上讲，在降低之前，每个线程都有自己的私有虚拟指令流。降低过程将这些虚拟指令流的指令映射到有限的物理指令流中。此过程需要正确性保证，因为物理指令流之间的物理消息队列需要在线程之间共享以传递依赖关系。只要按照图 18 中显示的规则保留顺序，我们就可以证明降低是正确的。以下定理给出了该降低过程正确性的充分必要条件

8. 结论

我们的工作提供了一个端到端的堆栈来解决跨各种硬件后端的基本优化挑战。我们希望我们的工作能够鼓励对编程语言、编译的更多研究，并为深度学习系统的硬件协同设计技术开辟新的机会。我们计划开源我们的编译器堆栈和 VITA 设计，以鼓励朝这个方向进一步研究。