TVM：用于在硬件平台上部署深度学习工作负载的端到端 IR 堆栈

2017 年 8 月 18 日

陈天奇（项目负责人），Thierry Moreau（硬件堆栈），蒋紫珩†（图编译），沈海辰（GPU 优化）

顾问：Luis Ceze，Carlos Guestrin，Arvind Krishnamurthy

华盛顿大学保罗·G·艾伦计算机科学与工程学院

DMLC 开源社区

†亚马逊网络服务

深度学习已变得无处不在且不可或缺。这场革命的部分原因是可扩展的深度学习系统，例如 TensorFlow、MXNet、Caffe 和 PyTorch。大多数现有系统都针对狭窄范围的服务器级 GPU 进行了优化，并且需要大量努力才能部署在其他平台（如手机、物联网设备和专用加速器（FPGA、ASIC））上。随着深度学习框架和硬件后端的数量增加，我们提出了一个统一的中间表示（IR）堆栈，该堆栈将弥合以生产力为中心的深度学习框架与以性能或效率为导向的硬件后端之间的差距。

我们很高兴宣布推出 TVM，作为解决此问题的方案。TVM 是一个新颖的框架，可以

• 表示和优化 CPU、GPU 和其他专用硬件的通用深度学习计算工作负载
• 自动转换计算图，以最大限度地减少内存利用率，优化数据布局并融合计算模式
• 提供从现有前端框架到裸机硬件，一直到浏览器可执行 javascript 的端到端编译。

借助 TVM，我们可以轻松地在手机、嵌入式设备甚至浏览器上运行深度学习工作负载，而无需付出额外的努力。TVM 还为多种硬件平台（包括依赖于新型计算原语的专用加速器）上的深度学习工作负载提供了统一的优化框架。

我们采用了编译器社区的通用理念，并提供两个中间表示层，以有效地将高级深度学习算法降低到多种硬件后端。

在今天的发布中，我们开源了 TVM 包，其中包含用于 x86、ARM、OpenCL、Metal、CUDA 和 Javascript 的优化原语。我们正在积极努力添加对专用硬件加速和 Nvidia 的 GEMM 优化 Volta 架构的支持。

技术细节

TVM 堆栈的目标是提供一个可重用的工具链，用于将来自深度学习框架前端的高级神经网络描述编译到用于多个硬件后端的低级机器代码。以 Apache MXNet 作为前端示例，以下代码片段演示了如何使用 TVM 将深度学习模型的高级描述编译为针对目标硬件量身定制的优化可执行模块。

挑战在于在支持多个硬件后端的同时，将计算、内存和能量占用量保持在最低水平。我们借鉴了编译器社区的智慧，以便弥合众多深度学习框架和硬件后端之间的差距：我们构建了一个由 NNVM 组成的双层中间层，NNVM 是用于任务调度和内存管理的高级中间表示 (IR)，TVM 是用于优化计算内核的富有表现力的低级 IR。

堆栈的第一层是基于计算图的表示。计算图是一个有向无环图，它将计算表示为节点，将数据流依赖性表示为边。这种表示非常强大：它允许我们将操作属性烘焙到计算图中，并指定转换规则以迭代地优化计算图。这是大多数现有深度学习框架（包括 TVM 堆栈中的 NNVM 图表示、TensorFlow XLA 和 Intel 的 ngraph）采用的常用方法。

图优化框架可以支持许多强大的优化。例如，我们提供了一个亚线性内存优化功能，允许用户在单个 GPU 上训练 1000 层 ImageNet ResNet。

但是，我们发现仅基于计算图的 IR 不足以解决支持不同硬件后端的挑战。原因是一个像卷积或矩阵乘法这样的单个图运算符可以针对每个硬件后端以非常不同的方式进行映射和优化。这些特定于硬件的优化在内存布局、并行化线程模式、缓存访问模式和硬件原语的选择方面可能会有很大差异。我们希望能够在通用表示中显式地表达这些优化旋钮，以有效地导航优化空间。

我们构建了一个低级表示来解决这个问题。此表示基于索引公式，并额外支持递归计算。

低级 IR 采用了现有图像处理语言（如 Halide 或 darkroom）的原理，以制定富有表现力的深度学习 DSL。TVM 构建了受循环转换工具（如 loopy 和基于多面体的分析）启发的低级优化。我们还从深度学习框架（如 MXNet、TensorFlow、Theano）中使用的数据流描述语言中汲取了灵感。然后在调度阶段处理 TVM 中描述的算法，以应用针对目标硬件后端量身定制的转换。

TVM 包括 CPU 优化框架中常见的标准转换原语。更重要的是，TVM 结合了针对 GPU 的新型优化原语，通过利用线程协作模式、数据布局转换和强大的新计算原语。将 TVM 与 NNVM 结合使用，为跨软件堆栈优化深度学习工作负载提供了丰富的机会，从而实现了联合的计算图级和运算符级优化。

多语言和平台支持

TVM 的众多优势之一在于其对多种平台和语言的丰富支持。我们介绍了框架的两个组成部分：编译器堆栈，其中包含用于生成优化机器代码的完整优化库；以及运行时，它轻量级且提供在不同平台上部署编译模块所需的可移植性。

TVM 目前支持 Python 和 C++ 接口到嵌入式编译器堆栈。我们在设计框架时考虑了最大的重用性，以便编译器堆栈的改进可以在 Python 和 C++ 组件之间互换应用。

我们还提供了一个轻量级运行时，可以在 javascript、java、python 和 c++ 等语言中直接运行 TVM 编译的代码，平台包括 android、iOS、raspberry pi 和 Web 浏览器。

远程部署和执行

TVM 支持使用 TVM RPC 进行交叉编译和测试嵌入式设备，TVM RPC 是一个轻量级接口，用于在远程嵌入式设备上部署和执行 TVM 交叉编译的模块。这为 TVM 用户提供了一个熟悉的高级 Python 接口，以便在各种低级嵌入式设备上远程编译、优化和测试深度学习算法。

性能

TVM 仍处于开发的早期阶段，我们可以期待更多的改进，但我们已经开始看到非常有希望的结果，本节将对此进行讨论。

TVM 使我们能够灵活地探索各种深度学习内核的丰富优化空间，以适应多种硬件平台。例如，TVM 允许我们为我们最关心的内核和平台定制数据布局和融合模式要求。请注意，基准库是为更通用的问题创建的，而 TVM 的优化内核是针对我们通过自动调优过程评估的工作负载进行了大量调整的。TVM 充当快速生成此类专用内核的桥梁。

本节列出的结果仍是正在进行的工作，并且还有改进的空间。

树莓派

在结果的第一部分，我们将 TVM CPU 调度与在 Raspberry Pi 3B 上执行 resnet 工作负载的 nnpack 进行了比较。由于时间有限，我们利用 TVM 实现了直接卷积，而 nnpack 用于执行 3x3 内核的 Winograd 卷积。

我们可以发现，借助 TVM 的自动调优内核，我们可以在树莓派实验中获得与 nnpack 中手动优化内核相似的性能。

GPU 结果

作者致谢 这些基准测试和相应的调度优化由我们的贡献者创建：Leyuan Wang (AWS / UCDavis), Yuwei Hu(TuSimple) 和 Weitang Liu (AWS/ UCDavis)。他们值得所有的赞誉。

作为概念验证，我们创建了一个端到端编译管道，可以将 MxNet 模型编译为 TVM 执行图。我们通过自动将运算符融合在一起并让 TVM 生成融合内核，在图节点内部和节点之间应用优化。我们对 mobilenet ImageNet 工作负载进行了基准测试，并在下面讨论了结果

我们可以发现，TVM 在速度方面可以优于我们的基线方法。更有趣的是，内核融合带来了额外的加速。值得一提的是，TVM 完全依靠自身生成所有优化的 GPU 内核，而无需依赖像 CuDNN 这样的外部库。

我们正在进行更多实验，并将发布获得的新结果。

开源工作

TVM 最初是华盛顿大学保罗·G·艾伦计算机科学与工程学院的研究项目。TVM 堆栈旨在支持 DLPack，这是多个主要深度学习框架对张量数据结构的共识。我们收到了来自 UW、AWS、奇虎 360、Facebook、香港科技大学、TuSimple、UCDavis、上海交通大学以及 DMLC 开源社区和 DLPack 倡议成员的早期贡献。展望未来，该项目将遵循 Apache 开源模式，创建一个社区维护的项目。非常欢迎您加入我们并领导这项工作。

致谢

如果没有我们早期的贡献者，这个项目就不可能实现。我们要感谢刘轶之（奇虎 360）、胡宇威（TuSimple）、史兴佳（香港科技大学）、王乐元（AWS/UCDavis）、Nicolas Vasilache（Facebook）、翁剑（UCLA）、刘伟堂（AWS/UCDavis）、Edward Z. Yang（Facebook）、郑立民（上海交通大学）、张乔（UW）、William Moses（Facebook/MIT）和胡士文。作者还要感谢张先轶（PerfXLab）的帮助性讨论。

在构建 TVM 时，我们也从以下项目中学习了很多。

• Halide：TVM 使用 HalideIR 作为数据结构进行算术简化和低级降低。HalideIR 派生自 Halide。我们在 TVM 中实现降低管道时也从 Halide 中学习。
• Loopy：整数集分析及其循环变换原语的使用。
• Theano：用于递归的符号扫描运算符的设计灵感。

源代码

• Github 页面可以在这里找到：https://github.com/dmlc/tvm
• TVM 与 DLPack 兼容，这使得它很容易支持采用该标准的框架，例如 MXNet、PyTorch、Caffe2 和 tiny-dnn。