TVM | 一种用于深度学习的端到端自动优化编译器-六虎

携手发明，一起生长！这是我参与「日新计划 10 月更文应战」的第 13 天点击查看活动详情

欢迎关注我的大众号 [极智视界]，获取我的更多笔记共享

大家好，我是极智视界，本文解读一下 一种用于深度学习的端到端主动优化编译器 TVM。

现在越来越需求将机器学习布置到各种硬件设备，当前的结构依靠于供货商特定的算子，并仅仅针对规模较窄的服务器级的 GPU 进行了优化。将作业负载布置到新平台：例如手机、嵌入式设备和加快器 (例如 FPGA、ASIC 等)，需求很多的手艺作业。作者提出了 TVM，这是一种公开了 graph等级和算子等级的优化，在跨不同硬件后端的深度学习作业负载的状况下也能表现优秀的功用。TVM 处理了深度学习特有的优化应战，例如高档算子交融、映射到任意硬件原语和内存推迟躲藏。它还经过选用一种新颖的、依据学习的本钱建模办法来快速查找代码优化，然后主动优化初级程序以习惯硬件特性。试验成果表明，TVM 供给的跨硬件后端功用可与用于低功耗 CPU、移动 GPU 和服务器级GPU 的最先进的手动调整库相媲美。另外作者还展现了 TVM 针对新加快器后端的习惯能力，例如依据 FPGA 的通用深度学习加快器。值得一提的是，TVM 是开源的，并现已在几家大厂内用于出产。

论文地址：www.usenix.org/system/file…

github地址：github.com/apache/tvm

源码不限速下载地址(v0.11版别&包含三方库)：download.csdn.net/download/we…

源码不限速下载地址(v0.10版别&包含三方库)：download.csdn.net/download/we…

1. 简介

深度学习模型 (DL) 现在能够用来辨认图画、处理天然言语，并在具有应战性的战略游戏中打败人类。从云服务器到主动驾驶轿车和嵌入式设备，将智能运用程序布置到各种设备的需求在不断增长。由于硬件特性的多样性，包含嵌入式 CPU、GPU、FPGA 和 ASIC (例如 TPU)，将 DL 作业负载映射到这些设备变得杂乱。这些硬件方针在内存组织、核算功用单元等方面存在差异，如图1所示。

当前的 DL 结构，如 TensorFlow、MXNet、Caffe 和 PyTorch，依托核算图中心表明来完成优化，例如主动微分和动态内存办理。可是，图级优化一般过于高档，无法处理特定于硬件后端的算子等级的转化。这些结构中的大多数都专心于一小类服务器级 GPU 设备，针对特定方针的优化依靠于高度工程化和特定于供货商的算子库。这些算子等级的库需求很多的手动调整，因而过于专业且不通明，无法轻松跨硬件设备移植。现在，在各种 DL 结构中为各种硬件后端供给支撑需求很多的工程尽力。即便关于现已支撑的后端，结构也有必要在以下之间做出艰难的挑选：(1) 防止发生不在预界说算子库中的新算子的图优化；(2) 运用这些新算子的未优化的视线。为了对不同的硬件后端启用图级和算子级的优化，作者选用了一种完全不同的端到端的办法。作者构建了 TVM，这是一个编译器，它从现有结构中获取深度学习程序的高档标准，并为各种硬件后端生成初级优化代码。为了招引更多用户，TVM 需求供给与跨不同硬件后端的很多手动优化的算子库比较具有竞争力的功用。这一方针需求处理下面描绘的要害应战。

Levaraging specific Hardware Features and Abstractions DL 加快器引进了优化的张量加快原语，而 GPU 和 CPU 则不断改善其处理元素。这对为给定的算子描绘生成优化代码提出了重大应战。硬件指令的输入是多维的，具有固定或可变长度；他们决议了不同的数据布局；他们对内存层次结构有特别要求。体系有必要有效地运用这些杂乱的原语才能从加快中获益。此外，加快器规划一般也支撑愈加精简的操控，并将大多数调度杂乱性下放到编码器仓库。关于专用加快器，体系需求生成显式操控管道依靠联系的代码，以躲藏内存拜访推迟 – 这是硬件为 CPU 和 GPU 履行的一项作业。

Large Search Space for Optimization 另一个应战是在不手动调整算子的状况下生成高效的代码。内存拜访、线程方式和新颖的硬件原语的组合挑选为生成的代码 (例如循环块、排序、缓存和打开) 创建了巨大的装备空间，如果要完成黑盒主动调整，将需求很多的查找本钱。能够选用预界说的本钱模型来辅导查找，但由于现代硬件的杂乱性日益添加，构建精确的本钱模型很困难。此外，这种办法需求我们为每种硬件类型建立独自的本钱模型。

TVM 经过三个要害模块处理了这些应战。(1) 引进张量表达式言语来构建算子并供给程序转化原语，这些原语能够生成具有各种优化的不同版别的程序。这个 layer 扩展了 Halide 的核算 / 调度别离概念，还将方针硬件内涵函数与转化原语别离，然后支撑新的加快器及其相应的新内涵函数。此外，引进了新的转化原语来处理与 GPU 相关的应战，并支撑布置到专门的加快器。然后运用不同的程序转化序列，为给定的算子声明构成一个丰厚的有效程序空间。(2) 引进了一个主动程序优化结构来寻找优化的张量算子。优化器由依据 ML 的本钱模型辅导，随着我们从硬件后端搜集更多的数据，该模型会习惯和改善。(3) 在主动代码生成器之上，引进了一个 graph重写器，它充分运用了高档和算子等级的优化。

经过结合这三个模块，TVM 能够从现有的深度学习结构中获取模型描绘，履行高档和初级联合优化，并为后端生成特定于硬件的优化代码，例如 CPU、GPU 和依据FPGA的专用加快器。本文的奉献如下：

确认了为跨不同硬件后端的深度学习作业负载供给功用可移植性方面的首要优化应战；
介绍了运用跨线程内存重用、新颖的硬件内涵函数和推迟躲藏的新颖调度原语；
提出并完成了一个依据机器学习的优化体系，以主动查找和查找优化的张量算子；
构建了一个端到端的编译和优化仓库，答应将高档结构 (TensorFlow、MXNet、PyTorch、Kears、CNTK) 中指定的深度学习作业负载布置到不同的硬件后端 (包含 CPU、服务器 GPU、移动 GPU 和依据FPGA的加快器)；
开源 TVM 在几家大公司内用于出产；

运用服务器级GPU、嵌入式GPU、嵌入式CPU 和依据FPGA 的定制通用加快器上的实在作业负载评价了 TVM。试验成果表明，TVM 供给了跨后端的可移植功用，并且比由手动优化库支撑的现有结构完成了 1.2 倍到 3.8 倍的加快。

2. 总览

本节经过一个示例来介绍 TVM 的组件。图2总结了 TVM 中的履行步骤及其在论文中的相应部分。

该体系首要将来自现有结构的模型作为输入，并将其转化为核算图表明。然后它履行高档数据流重写以生成优化图。算子级优化模块有必要为该图中的每个交融算子生成高效代码。算子以声明性张量表达式言语指定；履行细节未指定。TVM 为给定硬件方针的算子辨认一组或许的代码优化。或许的优化构成了很大的空间，因而运用依据 ML 的本钱模型来寻找优化的算子。最终，体系将生成的代码打包到一个可布置的模块中。

End-User Example 只需求几行代码，用户就能够从现有的深度学习结构中获取模型并调用 TVM API 以获取可布置的模块：

import tvm as t
# Use keras framework as example, import model
graph, params = t.frontend.from_keras(keras_model)
target = t.target.cuda()
graph, lib, params = t.compiler.build(graph, target, params)

这个编译后的运转时模块包含三个组件：终究优化的核算图 (graph)、生成的算子 (lib) 和模块参数 (params)。然后能够运用这些组件将模型布置到方针后端：

import tvm.runtime as t

module = runtime.create(graph, lib, t.cuda(0))
module.set_input(**params)
module.run(data=data_array)
output = tvm.nd.empty(out_shape, ctx=t.cuda(0))
module.get_output(0, output)

TVM 支撑多种言语的布置后端，例如 C++、Java 和 Python。后面的其余部分描绘了 TVM 的架构以及体系程序员怎么扩展它以支撑新的后端。

3. 优化核算图

核算图是在 DL 结构中表明程序的常用办法。图3展现了一个两层卷积网络的示例核算图表明。

这种高档表明与初级编译器中心表明 (IR) (如 LLVM) 之间的首要区别在于：中心数据项是大型的多维张量。核算图供给了算子的全局视图，但它们防止指定每个算子有必要怎么完成。与 LLVM IR 一样，核算图能够转化为功用等效的图以运用优化。作者还运用常见 DL 作业负载中的形状特异性来优化一组固定的输入形状。

TVM 运用核算图表明来运用高档优化：节点表明对张量或程序输入的操作，边表明操作之间的数据依靠联系。它完成了许多图级优化，包含：算子交融，将多个小操作交融在一起；常量折叠，预先核算能够静态确认 graph 部分，节约履行本钱；一个静态内存规划 pass，它预先分配内存来保存每个中心张量；数据排布转化，将内部数据排布转化为后端友爱的方式。接下来评论算子交融和数据排布转化。

Operator Fusion 算子交融将多个算子组合到一个内核中，而不会将中心成果保存在内存中。这种优化能够大大削减履行时刻，特别是在 GPU 和专用加快器中的履行时刻。详细来说，有四类 graph 算子：(1) injective (一对一映射，例如 add)；(2) reduction (例如 sum)；(3) complex-out-fusable (能够将元素映射交融到输出，例如 conv2d)；(4) opaque (不能交融的，例如 sort)。这里供给了交融这些算子的通用规矩，如下所示。多个 injective 算子能够交融成另一个 injective 算子。reduction 算子能够与输入 injective 算子交融 (例如交融 scale 和 sum)。比如 conv2d 之类的算子是 complex-out-fusable 的，我们能够将元素算子交融到它的输出中。能够运用这些规矩将核算图转化为交融版别。图4展现了这种优化对不同作业负载的影响。作者发现交融算子经过削减内存拜访会发生高达 1.2 到 2 倍的加快。

Data Layout Transformation 有多种办法能够在核算图中存储给定的张量。最常见的数据排布挑选是列主序和行主序。在实践中，我们或许会愈加喜欢运用愈加杂乱的数据排布。例如，DL 加快器或许会选用 4X4 矩阵运算，需求将数据排布平铺成 4X4 块以优化拜访局部性。

数据排布优化将核算图转化为能够运用愈加好的内部数据排布以在方针硬件上履行核算图。它首要在给定内存层次结构规矩的约束条件下为每个算子指定首选数据排布。如果首选数据布局不匹配，就在出产者和顾客之间履行适当的布局转化。

尽管高档图优化能够极大地进步 DL 作业负载的功率，但他们也仅仅与算子库差不多有效。现在，支撑算子交融的少量 DL 结构需求算子库来供给交融方式的完成。随着定期引进更多的网络算子，或许的交融内核数量会急剧添加。当需求面对越来越多的硬件后端时，这种办法不再可持续，由于所需求的交融方式完成数量与有必要支撑的数据排布、数据类型和加快器内涵函数的数量相结合。为程序所需的各种操作和每个后端手艺规划算子内核是不可行的。为此，接下来提出了一种代码生成办法，能够为给定模型的算子生成各种或许的完成。

4. 生成张量算子

TVM 经过在每个硬件后端生成许多有效的完成并挑选优化的完成，为每个算子生成高效的代码。这个进程建立在 Halide 将描绘和核算规矩 (或调度优化) 解耦的想法的根底上，并将其扩展为支撑新的优化 (嵌套并行、张量和推迟躲藏) 和广泛的硬件后端。

4.1 张量表达和调度空间

引进了张量表达式言语来支撑主动代码生成，与高档核算图表明不同，张量算子的完成是不通明的，每个算子都用索引公式表达言语来描绘。以下代码展现了核算转置矩阵乘法的示例张量表达式。

每个核算操作都指定输出张量的形状和描绘怎么核算它的每个元素的表达式。张量表达式言语支撑常见的算术和数学运算，并包含常见的 DL 算子方式。该言语没有指定循环结构和许多其他履行细节，它供给了为各种后端添加硬件感知优化的灵活性。选用来自 Halide 的解耦核算/调度准则，运用调度来表明从张量表达式到初级代码的特定映射。

经过增量运用坚持程序逻辑等价的根本转化 (调度原语) 来构建调度。图5展现了在专用加快器上调度矩阵乘法的示例。在内部，TVM 运用调度转化，运用数据结构来盯梢循环结构和其他信息。然后此信息能够协助按给定的终究调度生成初级代码。

作者的张量表达式学习了 Halide、Darkroom 和 TACO。为了在许多后端完成高功用，有必要要支撑足够多的调度原语来包含不同硬件后端的各种优化。图6总结了 TVM 支撑的算子代码生成进程和调度原语。作者重用了有用的原语和来自 Halide 的初级循环程序 AST，并且引进了新的原语来优化 GPU 和加快器的功用。新的原语是完成最佳 GPU 功用所必需的，也是加快器所必需的。CPU、GPU 和类TPU加快器是三种重要的深度学习硬件。

4.2 嵌套并行与协作

并行性是进步 DL 作业负载中核算密集型内核功率的要害。现代 GPU 供给了大规模的并行性，要求我们将并行方式交融到调度转化中。大多数现有的处理方案都选用一种称为嵌套并行的模型，这是一种 fork-join 的方式。该模型需求并行调度原语来并行化数据并行使命；每个使命能够进一步递归地细分为子使命，以运用方针架构的多级线程层次结构 (例如 GPU 中的线程组)。将此模型称为无同享嵌套并行，由于一个作业线程无法在同一并行核算阶段查看其兄弟的数据。

无同享办法的代替办法是协作获取数据。详细来说，线程组能够协作获取它们都需求的数据并将其放入同享内存空间中。这种优化能够运用 GPU 内存层次结构，并经过同享内存区域完成跨线程的数据重写。TVM 运用调度原语来支撑这种众所周知的 GPU 优化，以完成最佳功用。以下 GPU 代码示例对矩阵乘法进行了优化。

图7展现了这种优化的影响，将内存规模的概念引进调度空间，以便能够将核算阶段 (代码中的 AS 和 BS) 标记为同享。

如果没有显式内存规模，主动规模推理会将核算阶段标记为 thread-local。同享使命有必要核算组中一切作业线程的依靠联系，此外还有必要正确刺进内存同步屏障，以保证同享加载的数据对顾客可见。最终，除了对 GPU 有用之外，内存规模还让我们能够标记特别的内存缓冲区，并在针对专门的 DL 加快器时创建特别的下降规矩。

4.3 张量化

DL 作业负载具有很高的算术强度，一般能够分化为张量算子，如矩阵乘或一维卷积。这些天然分化导致了最近添加的张量核算原语的趋势。这些新的原语为依据调度的编译发明了机遇和应战，运用它们能够进步功用，所以编译结构有必要无缝的集成它们。这种张量化类似于 SIMD 架构的矢量化，但又有着明显差异。指令输入是多维的，具有固定或可变的长度，并且具有不同的数据排布。更为重要的是，这种原语不能是固定的，由于新的加快器有或许会出现张量指令的变体，因而需求一个可扩展的处理办法。

经过运用张量内涵声明机制将方针硬件内涵特性与调度别离，然后使张量化变得能够扩展。作者运用相同的张量表达式言语来声明每个新硬件内涵行为和与之相关的下降规矩。以下代码展现了怎么声明一个 8×8 张量的硬件内涵函数。

此外，引进了张量调度原语，以用相应的内涵函数替换核算单元。编译器将核算方式与硬件声明相匹配，并将其下降到相应的硬件内涵特性。张量化将调度与特定的硬件原语别离，然后能够轻松扩展 TVM 以支撑新的硬件架构。生成的张量调度代码符合高功用核算特性：将杂乱的算子分化为一系列 micro-kernel 调用。我们还能够运用 tensorize 原语来运用手艺规划的 micro-kernel，这在某些平台上或许是有利的。例如，经过运用 bit-serial 矩阵向量乘法 micro-kernel 为移动 CPU 完成超低精度算子，这些算子对一比特或两比特宽的数据类型进行操作。这个 micro-kernel 将成果累积成越来越大的数据类型，以最大极限地削减内存占用。将 micro-kernel 表达为 TVM 固有的张量与非张量化版别比较可发生高达1.5倍的加快。

4.4 显式内存推迟躲藏

推迟躲藏是指将内存操作与核算重叠以最大极限地运用内存和核算资源的进程。它需求不同的战略，详细取决于方针硬件后端。在 CPU 上，内存推迟躲藏是经过同步多线程或硬件预取来隐式完成的。GPU 依靠于许多线程的快速上下文切换。比较之下，比如 TPU 之类的专用 DL 加快器一般倾向于运用解耦拜访履行 (DAE) 架构记性愈加精简的操控，并将细粒度同步的问题转移给软件。

图9显示了削减运转时推迟的 DAE 硬件管道。与单片硬件规划比较，流水线能够躲藏大部分内存拜访开支，简直能够充分运用核算资源。为了完成愈加高的运用率，指令流有必要添加细粒度的同步操作。没有它们，就无法强制履行依靠联系，然后导致错误履行。因而，DAE 硬件流水线需求在流水线阶段依靠于细粒度的入队/出队操作，以保证正确履行，如图9的指令流所示。

对需求显式初级同步的 DAE 加快器进行编程很困难。为了削减编程担负，引进了虚拟线程调度原语，让程序员能够指定高档数据并行程序，就像他们指定支撑多线程的硬件后端一样。然后 TVM 会主动将程序下降为具有初级显式同步的单个指令流，如图8所示。该算法从高档多线程程序调度开端，然后刺进必要的初级同步操作以保证每个线程的正确履行。接下来，它将一切虚拟线程的操作交付给单个指令流中。最终，硬件恢复指令流中初级同步所规矩的可用流水线的并行度。

Hardware Evaluation of Latency Hiding 展现了推迟躲藏在依据 FPGA 的定制加快器规划上的有效性。在加快器上运转 ResNet 的每一层，并运用 TVM 生成两个调度：一个有推迟躲藏，一个没有。具有推迟躲藏的调度将程序与虚拟线程并行化以进行管道并行，然后躲藏了内存拜访推迟。成果在图10 中显示为 roofline图表；roofline功用图能够深入了解给定体系在不同基准测验中运用核算和内存资源的状况。整体而言，推迟躲藏进步了一切 ResNet 层的功用。峰值核算运用率从没有推迟躲藏的 70% 进步到了躲藏推迟的 88%。

5. 主动优化

鉴于丰厚的调度原语调集，剩下的问题是为 DL 模型的每一层找到最佳的算子完成。在这里，TVM 为与每一层关联的特定输入形状和布局创建了一个专门的算子。这种专门的优化供给了明显的功用优势 (与针对较小的形状和布局多样性的手艺代码比较)，但它也提出了主动化的应战。体系需求挑选调度优化：例如修正循环次序或优化内存层次结构，以及特定于调度的参数，例如切片大小和循环打开因子。这样的组合挑选为每个硬件后端创建了算子完成的大的查找空间。为了应对这一应战，构建了一个具有两个首要组件的主动调度优化器：一个新装备的调度查找器以及一个猜测给定装备功用的机器学习本钱模型。如图11展现了这些组件和 TVM 的主动优化流程。

5.1 调度空间标准

作者构建了一个调度模版标准 API，让开发人员在调度空间中声明旋钮。在指定或许的调度时，模版标准答应在必要时结合开发人员的特定领域知识。作者还为每个硬件后端创建了一个通用主模板，该模板依据运用张量表达式言语表达的核算描绘主动提取或许的旋钮。在高层次上，我们或许希望尽或许多的装备，并让优化器办理挑选担负。因而优化器有必要为我们试验中运用的实在世界 DL 作业负载查找数十亿种或许的装备。

5.2 依据机器学习的本钱模型

从大型装备空间中找到最佳调度的一种办法是经过黑盒优化，即主动调整，该办法用于调优高功用核算库。可是，主动调整需求许多试验来确认一个好的装备。

另一种办法是构建一个预界说的本钱模型来辅导对特定硬件后端的查找，而不是运转一切或许性并丈量它们的功用。抱负状况下，一个完美的本钱模型会考虑影响功用的一切要素：内存拜访方式、数据重用、管道依靠性和线程方式等。不幸的是，由于现代硬件日益杂乱，这种办法很繁琐。此外，每个新的硬件方针都需求一个新的 (预界说的) 本钱模型。

与此相反，作者选用计算办法来处理本钱建模问题。在这种办法中，调度查找器提出能够进步算子功用的装备。关于每个调度装备，运用一个 ML 模型，该模型将下降循环程序作为输入，并猜测其在给定硬件后端的运转时刻。该模型运用查找期间搜集的运转时丈量数据进行练习，不需求用户输入详细的硬件信息。当在优化进程中查找更多装备时，会定期更新模型，这也会进步其他相关作业负载的精确性。经过这种方式，ML 模型的质量会随着更多的试验而进步。表1总结了主动化办法之间的首要区别。依据 ML 的本钱模型在主动调整和预界说本钱建模之间获得平衡，并且能够从相关作业负载的前史功用数据中获益。

Machine Learning Model Design Choices 在调度查找器挑选将运用哪种 ML 模型时，我们有必要考虑两个要害要素：质量和速度。调度查找器频繁查询本钱模型，由于模型猜测时刻和模型调整时刻而发生开支。为了可用，这些开支有必要小于在实在硬件上丈量功用所需的时刻，依据详细的作业负载/硬件方针，这或许是几秒钟的数量级。这种速度要求将我们的问题与传统的超参数调整问题区别开来，在传统超参数调整问题中，履行丈量的本钱相关于模型开支非常高，并且能够运用愈加贵重的模型。除了模型的挑选，我们还需求挑选一个方针函数来练习模型，例如装备猜测运转时刻的差错。可是，调度查找器仅依据猜测的相对次序 (A 运转速度比 B 快) 来挑选最佳候选，而不需求直接猜测绝对履行时刻。相反，作者选用 rank 方针来猜测运转时本钱的相对次序。

作者在 ML 优化器中完成了几种类型的模型。选用梯度树提高模型 (依据 XGBoost)，它依据从循环程序中提取的特征进行猜测；这些功用包含每个循环等级的每个内存缓冲区的内存拜访计数和重用率，以及循环注释的 one-hot 编码，例如 vectorize、unroll 和 parallel。作者还评价了一个神经网络模型，该模型运用 TreeRNN 在没有特征工程的状况下总结循环程序的 AST。图13总结了本钱模型的作业流程。

作者发现树增强和 TreeRNN 具有相似的猜测质量。可是，前者的猜测速度要快两倍，并且练习时刻要少得多。因而，在试验中挑选梯度树提高作为默认本钱模型。尽管如此，作者还是信任这两种办法都是有价值的，并等待在这个问题上进行愈加多的研讨。均匀而言，树增强模型能够在 0.67 毫秒内进行猜测，这比运转实际丈量要快数千倍。图12将依据 ML 的优化器与黑盒主动调整办法进行了比较；前者比后者愈加快地找到了更好的装备。

5.3 调度查找

一旦我们挑选了一个本钱模型，我们就能够运用它来挑选主张的装备，在这些装备上迭代地运转实际丈量。在每次迭代中，查找器运用 ML 模型的猜测来挑选一批候选者来运转丈量。然后将搜集的数据用作练习数据以更新模型。如果不存在初始练习数据，则查找器会挑选随机候选者进行丈量。

最简略的查找算法经过本钱模型枚举并运转每个装备，挑选前 k 个猜测的履行者。可是这种战略再查找空间很大时变得难以处理。相反，作者选用并行模拟退火算法。查找器从随机装备开端，并且在每一步中，随机走到附近的装备。如果本钱如本钱模型所猜测的那样下降，则这种转变是成功的。如果方针装备本钱较高，则或许会失利 (回绝)。这样随机游走在倾向于收敛于本钱模型猜测的本钱较低的装备。查找状态在本钱模型更新中持续存在；我们从这些更新后的最终一个装备持续查找。

5.4 分布式设备池和 RPC

分布式设备池能够扩展硬件试验的运转，并在多个优化作业之间完成细粒度的资源同享。TVM 完成了一个定制的、依据 RPC 的分布式设备池，使客户端能够在特定类型的设备上运转程序。我们能够运用此接口在主机编译器上编译程序，恳求长途设备、长途运转函数，并在主机上拜访相同脚本中的成果。TVM 的 RPC 支撑动态上传并运转运用其运转时约定的交叉编译模块和函数。因而，相同的根底架构能够履行单个作业负载优化和端到端 graph 推理。作者的办法能够跨多个设备主动履行编译、运转和装备步骤。这种根底设施关于嵌入式设备特别重要，由于传统上这些设备需求繁琐的手动操作来进行交叉编译、代码布置和丈量。