陈天奇等人提出TVM：深度学习自动优化代码生成器

2018 年 2 月 17 日 机器之心

选自arXiv

作者：陈天奇等

机器之心编译

参与：李泽南、路雪

TVM 是由华盛顿大学在读博士陈天奇等人提出的深度学习自动代码生成方法，去年 8 月机器之心曾对其进行过简要介绍。该技术能自动为大多数计算硬件生成可部署优化代码，其性能可与当前最优的供应商提供的优化计算库相比，且可以适应新型专用加速器后端。近日，这项研究的论文《TVM: End-to-End Optimization Stack for Deep Learning》终于完成，内容包含新方法的介绍与讨论，以及 TVM 在英伟达、AMD 的 GPU、树莓派及一些 FPGA 上的性能评估。

项目链接：https://github.com/dmlc/tvm

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

显而易见，以现有的点到点方式实现不同深度学习框架对所有种类的硬件进行后端支持是不现实的。我们的最终目标是让深度学习负载可以轻松部署到所有硬件种类中，其中不仅包括 GPU、FPGA 和 ASIC（如谷歌 TPU），也包括嵌入式设备，这些硬件的内存组织与计算能力存在着显著的差异（如图 1 所示）。考虑到这种需求的复杂性，开发一种能够将深度学习高级程序降低为适应任何硬件后端的低级优化代码的优化框架是最好的方法。

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

图 1：CPU、GPU 与类 TPU 加速器需要不同的片上存储架构和计算基元。在生成优化代码时我们必须考虑这一问题。

图 2：TVM 堆栈图。目前的堆栈支持多种深度学习框架以及主流 CPU、GPU 以及专用深度学习加速器。

优化的四大基本挑战

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

高级数据流复写：不同的硬件设备可能具有截然不同的内存层次结构，因此，融合算子与优化数据布局的策略对于优化内存访问至关重要。
跨线程内存复用：现代 GPU 与专用加速器的内存可被多个计算核心共享，传统的无共享嵌套并行模式已不再是最优方法。为优化内核，在共享内存负载上的线程合作很有必要。
张量计算内部函数：最新的硬件带来了超越向量运算的新指令集，如 TPU 中的 GEMM 算子和英伟达 Volta 架构中的 Tensor Core。因此在调度过程中，我们必须将计算分解为张量算术内部函数，而非标量或向量代码。
延迟隐藏（Latency Hiding）：尽管在现代 CPU 与 GPU 上，同时拥有多线程和自动缓存管理的传统架构隐藏了延迟问题，但专用的加速器设计通常使用精简控制与分流，这为编译器堆栈的调度带来了复杂性。所以，调度仍需仔细，以隐藏内存访问延迟。

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

我们构建了一个端到端的编译优化堆栈，允许将高级框架（如 Caffe、MXNet、PyTorch、Caffe2、CNTK）专用的深度学习工作负载部署到多种硬件后端上（包括 CPU、GPU 和基于 FPGA 的加速器）。
我们发现了提供深度学习工作负载在不同硬件后端中的性能可移植性的主要优化挑战，并引入新型调度基元（schedule primitive）以利用跨线程内存重用、新型硬件内部函数和延迟隐藏。
我们在基于 FPGA 的通用加速器上对 TVM 进行评估，以提供关于如何最优适应专用加速器的具体案例。

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

图 3：两层卷积神经网络的计算图示例。图中每个节点表示一次运算，它消耗一或多个张量，并生成一或多个张量。张量运算可以通过属性进行参数化，以配置其行为（如 padding 或 stride）。

论文：TVM: End-to-End Optimization Stack for Deep Learning

论文链接：https://arxiv.org/abs/1802.04799

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