TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
TVM：一个用于深度学习的自动化端到端优化编译器
作者：Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, Arvind Krishnamurthy
机构：华盛顿大学 Paul G. Allen 计算机科学与工程学院，AWS，上海交通大学，加州大学戴维斯分校，康奈尔大学
A1 主要贡献
核心问题与研究目标：随着机器学习应用需求的增长，需要将其部署到多样化的硬件设备上，从云服务器到自动驾驶汽车和嵌入式设备。然而，当前的深度学习（DL）框架（如 TensorFlow、MXNet）严重依赖于供应商特定的算子库，且主要为服务器级 GPU 优化。将工作负载部署到新平台（如手机、嵌入式设备、FPGA、ASIC）需要大量手动工作。这是因为不同硬件在内存组织、计算单元等方面存在巨大差异（如图1所示），而现有的算子库过于专用和不透明，难以跨硬件移植。此外，现有框架难以在图优化（可能产生新算子）和使用未优化的新算子实现之间做出选择。
本文的研究目标是构建一个端到端的编译器TVM，它能够接收来自现有框架的深度学习程序高级规范，并为各种硬件后端生成低级别的优化代码。TVM旨在提供与手动优化的算子库相媲美的性能，同时解决深度学习优化的特定挑战，包括高级算子融合、映射到任意硬件原语以及内存延迟隐藏。
创新点与主要贡献：
TVM通过三个关键模块来应对上述挑战：
1. 张量表达式语言和调度原语：引入一种张量表达式语言来构建算子，并提供程序转换原语来生成具有不同优化的程序版本。此层扩展了Halide【【32】，Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，PLDI】的计算/调度分离概念，通过将目标硬件内联函数与转换原语分离，从而支持新型加速器及其新内联函数。此外，还引入了新的调度原语来解决GPU相关的挑战（如跨线程内存复用）和专用加速器的需求（如延迟隐藏）。
2. 基于机器学习的自动优化框架：引入一个自动程序优化框架来寻找优化的张量算子。该优化器由一个基于机器学习的成本模型指导，该模型会随着从硬件后端收集更多数据而自适应地改进。这种方法能够快速探索巨大的代码优化配置空间。
3. 图重写器：在自动代码生成器之上，引入一个图重写器，以充分利用高级（图级别）和算子级的优化。
本文的主要贡献总结如下：
* 识别了在多样化硬件后端上为深度学习工作负载提供性能可移植性的主要优化挑战。
* 引入了新的调度原语，以利用跨线程内存复用、新型硬件内联函数和延迟隐藏。
* 提出并实现了一个基于机器学习的优化系统，用于自动探索和搜索优化的张量算子。
* 构建了一个端到端的编译和优化栈，允许将高级框架（包括TensorFlow、MXNet、PyTorch、Keras、CNTK）中指定的深度学习工作负载部署到多样化的硬件后端（包括CPU、服务器GPU、移动GPU和基于FPGA的加速器）。TVM系统是开源的，并已在多家大公司中投入生产使用。
实验结果表明，TVM在服务器级GPU、嵌入式GPU、嵌入式CPU和定制的基于FPGA的通用加速器上，提供了与现有手工优化库相媲美甚至超越的性能，速度提升范围为1.2倍至3.8倍。
A3 背景知识与设计原则
本节通过一个示例来逐步介绍TVM的组件。图2总结了TVM中的执行步骤。系统首先接收来自现有框架的模型作为输入，并将其转换为计算图表示。然后，它执行高级数据流重写以生成优化的图。算子级优化模块必须为该图中的每个融合算子生成高效的代码。算子在声明性张量表达式语言中指定，执行细节未指定。TVM为给定硬件目标的算子识别了一系列可能的代码优化。可能的优化形成了一个巨大的空间，因此我们使用基于机器学习的成本模型来寻找优化的算子。最后，系统将生成的代码打包成一个可部署的模块。
终端用户示例。用户只需几行代码，就可以从现有的深度学习框架中获取模型，并调用TVM API来获得一个可部署的模块：
import tvm as t
# 以keras框架为例，导入模型
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等语言。
3 优化计算图
计算图表示。计算图是深度学习框架中表示程序的常用方法【【3】，Tensorflow: A system for large-scale machine learning，2016，OSDI】，【【4】，An introduction to computational networks and the computational network toolkit，2014，MSR-TR】，【【7】，Theano: new features and speed improvements，2012，NIPS Workshop】，【【9】，MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems，2015，NIPS Workshop】。图3展示了一个两层卷积神经网络的计算图示例。这种高级表示与LLVM等低级编译器中间表示（IR）的主要区别在于，中间数据项是大的多维张量。计算图提供了算子的全局视图，但避免指定每个算子的实现方式。与LLVM IR一样，计算图可以被转换为功能上等价的图以应用优化。我们还利用常见深度学习工作负载中形状的特异性，为一组固定的输入形状进行优化。
图级别优化。TVM利用计算图表示来应用高级优化：节点代表对张量或程序输入的操作，边代表操作之间的数据依赖。它实现了许多图级优化，包括：算子融合，将多个小操作融合成一个；常量折叠，预计算图中可以静态确定的部分，节省执行成本；静态内存规划传递，预先分配内存以容纳每个中间张量；以及数据布局转换，将内部数据布局转换为后端友好的形式。下面我们讨论算子融合和数据布局转换。
算子融合。算子融合将多个算子组合成一个单一的内核，而无需在内存中保存中间结果。这种优化可以极大地减少执行时间，特别是在GPU和专用加速器上。具体来说，我们将图算子分为四类：（1）内射（injective，一对一映射，如加法），（2）归约（reduction，如求和），（3）复杂输出可融合（complex-out-fusable，可以将元素级映射融合到输出，如conv2d），以及（4）不透明（opaque，不能融合，如排序）。我们提供了融合这些算子的通用规则。多个内射算子可以融合成另一个内射算子。归约算子可以与输入的内射算子融合（例如，融合缩放和求和）。像conv2d这样的算子是复杂输出可融合的，我们可以将元素级算子融合到其输出。我们可以应用这些规则将计算图转换为融合版本。图4展示了这种优化在不同工作负载上的影响。我们发现融合算子通过减少内存访问，带来了高达1.2倍到2倍的加速。
数据布局转换。在计算图中，有多种方式可以存储一个给定的张量。最常见的数据布局选择是列主序和行主序。在实践中，我们可能更倾向于使用更复杂的数据布局。例如，一个深度学习加速器可能利用4×4矩阵操作，要求数据被平铺成4×4的块以优化访问局部性。数据布局优化将计算图转换为一个可以使用更优内部数据布局的图，以便在目标硬件上执行。它首先为每个算子指定首选的数据布局，给定内存层次结构所决定的约束。然后，如果生产者和消费者的首选数据布局不匹配，我们就在它们之间执行适当的布局转换。
高级优化的局限性。虽然高级图优化可以极大地提高深度学习工作负载的效率，但它们的有效性受限于算子库所能提供的内容。目前，少数支持算子融合的深度学习框架要求算子库为融合模式提供实现。随着新网络算子的不断引入，可能的融合内核数量会急剧增长。当目标硬件后端数量增加时，这种方法不再可持续，因为所需的融合模式实现数量随着数据布局、数据类型和加速器内联函数的数量呈组合式增长。为程序所需的各种操作以及每个后端手工制作算子内核是不可行的。为此，我们接下来提出一种代码生成方法，可以为给定模型的算子生成各种可能的实现。
A2 方法细节
TVM通过为每个算子在每个硬件后端上生成许多有效的实现，并选择一个优化的实现来产生高效的代码。这个过程建立在Halide【【32】，Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，PLDI】将描述与计算规则（或调度优化）解耦的思想之上，并扩展它以支持新的优化（嵌套并行、张量化和延迟隐藏）和广泛的硬件后端。我们现在重点介绍TVM的特有功能。
4.1 张量表达式与调度空间
声明式张量表达式。我们引入了一种张量表达式语言来支持自动代码生成。与高级计算图表示中张量操作的实现是不透明的不同，每个操作都在索引公式表达式语言中进行描述。以下代码显示了用于计算转置矩阵乘法的张量表达式示例：
m, n, h = t.var('m'), t.var('n'), t.var('h')
A = t.placeholder((m, h), name='A')
B = t.placeholder((n, h), name='B')
# computing rule
k = t.reduce_axis((0, h), name='k')
C = t.compute((m, n),
lambda y, x: t.sum(A[k, y] * B[k, x], axis=k))
# result shape
每个compute操作都指定了输出张量的形状和一个描述如何计算其每个元素的表达式。我们的张量表达式语言支持常见的算术和数学运算，并覆盖了常见的深度学习算子模式。该语言没有指定循环结构和许多其他执行细节，并为各种后端添加硬件感知优化提供了灵活性。
调度与调度原语。我们采用Halide【【32】，Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，PLDI】的计算/调度解耦原则，使用一个调度（schedule）来表示从张量表达式到低级代码的特定映射。许多可能的调度都可以执行此功能。我们通过增量应用保持程序逻辑等价性的基本转换（调度原语）来构建一个调度。图5展示了在一个专用加速器上调度矩阵乘法的示例。在内部，TVM使用一个数据结构来跟踪循环结构和其他信息，随着我们应用调度转换。这些信息随后可以帮助为给定的最终调度生成低级代码。
TVM对调度的扩展。我们的张量表达式借鉴了Halide【【32】，Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，PLDI】、Darkroom【【17】，Darkroom: Compiling high-level image processing code into hardware pipelines，2014，ACM Trans. Graph.】和TACO【【23】，The tensor algebra compiler，2017，Proc. ACM Program. Lang.】。其主要增强功能包括支持下面讨论的新调度优化。为了在许多后端上实现高性能，我们必须支持足够多的调度原语，以覆盖不同硬件后端上的各种优化。图6总结了TVM支持的操作代码生成过程和调度原语。我们重用了Halide中有用的原语和低级循环程序AST，并引入了新的原语来优化GPU和加速器性能。新的原语对于实现最佳GPU性能至关重要，对加速器来说更是必不可少。CPU、GPU和类TPU加速器是深度学习的三种重要硬件类型。本节描述了针对CPU、GPU和类TPU加速器的新优化原语，而第5节解释了如何自动推导出高效的调度。
4.2 合作性嵌套并行
并行性与共享内存。并行性是提高深度学习工作负载中计算密集型内核效率的关键。现代GPU提供大规模并行性，要求我们将并行模式融入调度转换中。大多数现有解决方案采用一种称为嵌套并行的模型，这是一种分叉-连接（fork-join）的形式。该模型需要一个并行调度原语来并行化数据并行任务；每个任务可以进一步递归地细分为子任务，以利用目标架构的多级线程层次结构（例如，GPU中的线程组）。我们称这种模型为无共享嵌套并行（shared-nothing nested parallelism），因为一个工作线程无法查看同一并行计算阶段内其兄弟线程的数据。
合作式数据获取。无共享方法的替代方案是合作式地获取数据。具体来说，一组线程可以合作获取它们都需要的数据，并将其放入共享内存空间。这种优化可以利用GPU内存层次结构，并通过共享内存区域实现跨线程的数据重用。TVM使用一个调度原语支持这种众所周知的GPU优化，以实现最佳性能。下面的GPU代码示例优化了矩阵乘法。
# 标记计算阶段为共享
s[AS] = t.compute_at(s[C], k)
s[BS] = t.compute_at(s[C], k)
s[AS].set_scope("shared")
s[BS].set_scope("shared")
# 绑定线程轴
ty, tx = s[C].op.thread_axis
s[C].bind(ty, t.thread_axis("blockIdx.y"))
s[C].bind(tx, t.thread_axis("blockIdx.x"))
# 合作式获取
ty, tx = s[AS].op.thread_axis
s[AS].bind(ty, t.thread_axis("threadIdx.y"))
s[AS].bind(tx, t.thread_axis("threadIdx.x"))
图7展示了这种优化的影响。我们向调度空间引入了内存作用域（memory scopes）的概念，以便计算阶段（代码中的AS和BS）可以被标记为共享。如果没有显式的内存作用域，自动作用域推断会将计算阶段标记为线程局部（thread-local）。共享任务必须计算组内所有工作线程的依赖关系。此外，必须正确插入内存同步屏障，以保证加载到共享内存的数据对消费者可见。最后，除了对GPU有用之外，内存作用域还让我们能够标记特殊的内存缓冲区，并在针对专用深度学习加速器时创建特殊的降级规则。
4.3 张量化
背景与挑战。深度学习工作负载具有高算术强度，通常可以分解为张量算子，如矩阵-矩阵乘法或一维卷积。这些自然的分解导致了最近增加张量计算原语的趋势【【1】，NVIDIA Tesla V100 GPU Architecture，2017】，【【12】，Eyeriss: A spatial architecture for energy-efficient dataflow for convolutional neural networks，2016，ISCA】，【【21】，In-datacenter performance analysis of a tensor processing unit，2017，ISCA】。这些新原语为基于调度的编译带来了机遇和挑战；虽然使用它们可以提高性能，但编译框架必须无缝地集成它们。我们称之为张量化（tensorization）：它类似于SIMD架构的向量化，但有显著差异。指令输入是多维的，具有固定或可变长度，并且每个都有不同的数据布局。更重要的是，我们不能支持一组固定的原语，因为新的加速器正在出现，它们有自己变化的张量指令。因此，我们需要一个可扩展的解决方案。
可扩展的张量化机制。我们通过将目标硬件内联函数与调度分离的机制，使张量化变得可扩展，该机制用于张量内联函数的声明。我们使用相同的张量表达式语言来声明每个新硬件内联函数的行为以及与之关联的降级规则。以下代码显示了如何声明一个8×8张量硬件内联函数。
# 声明一个8x8的gemm内联函数
def gemm_intrin_lower(y, x):
# ... 实现细节 ...
return tvm.call_intrin(..., op="gemm", ...)
# 注册内联函数
gemm8x8 = t.decl_tensor_intrin(y.op, gemm_intrin_lower)
此外，我们引入了一个tensorize调度原语，用相应的内联函数替换一个计算单元。编译器将计算模式与硬件声明进行匹配，并将其降级为相应的硬件内联函数。张量化将调度与特定的硬件原语解耦，使得扩展TVM以支持新的硬件架构变得容易。
张量化的应用。张量化调度的生成代码与高性能计算中的实践保持一致：将复杂操作分解为一系列微内核调用。我们也可以使用tensorize原语来利用手工制作的微内核，这在某些平台上可能是有益的。例如，我们通过利用位串行矩阵向量乘法微内核，为移动CPU实现了操作于一位或两位宽数据类型的超低精度算子。该微内核将结果累加到逐渐增大的数据类型中，以最小化内存占用。将该微内核作为张量内联函数呈现给TVM，比非张量化版本带来了高达1.5倍的加速。
4.4 显式内存延迟隐藏
延迟隐藏的背景。延迟隐藏是指将内存操作与计算重叠以最大化内存和计算资源利用率的过程。它根据目标硬件后端需要不同的策略。在CPU上，内存延迟隐藏是通过同时多线程（SMT）【【14】，Simultaneous multithreading: a platform for next-generation processors，1997，IEEE Micro】或硬件预取【【10】，Effective hardware-based data prefetching for high-performance processors，1995，IEEE Transactions on Computers】，【【20】，Improving direct-mapped cache performance by the addition of a small fully-associative cache and prefetch buffers，1990，ISCA】隐式实现的。GPU依赖于大量线程束（warps）的快速上下文切换【【44】，Understanding Latency Hiding on GPUs，2016，PhD thesis】。相比之下，像TPU【【21】，In-datacenter performance analysis of a tensor processing unit，2017，ISCA】这样的专用深度学习加速器通常倾向于更精简的控制，采用解耦访问-执行（decoupled access-execute, DAE）架构【【35】，Decoupled access/execute computer architectures，1982，ISCA】，并将细粒度同步的问题卸载给软件。
DAE架构与编译挑战。图9展示了一个减少运行时延迟的DAE硬件流水线。与单体硬件设计相比，该流水线可以隐藏大部分内存访问开销，并几乎完全利用计算资源。为了实现更高的利用率，指令流必须用细粒度的同步操作来增强。没有它们，依赖关系无法强制执行，导致错误执行。因此，DAE硬件流水线要求在流水线阶段之间进行细粒度的依赖入队/出队操作，以保证正确执行，如图9的指令流所示。
虚拟线程调度原语。对需要显式低级同步的DAE加速器进行编程是困难的。为了减轻编程负担，我们引入了一个虚拟线程（virtual threading）调度原语，让程序员可以像为支持多线程的硬件后端那样，指定一个高级数据并行程序。然后，TVM会自动将程序降级为带有低级显式同步的单指令流，如图8所示。该算法从一个高级多线程程序调度开始，然后插入必要的低级同步操作以保证每个线程内的正确执行。接下来，它将所有虚拟线程的操作交错成一个单指令流。最后，硬件根据指令流中的低级同步恢复可用的流水线并行性。
延迟隐藏的硬件评估。我们现在在一个定制的基于FPGA的加速器设计上展示延迟隐藏的有效性，我们将在6.4小节中深入描述。我们在该加速器上运行ResNet的每一层，并使用TVM生成两种调度：一种带延迟隐藏，一种不带。带延迟隐藏的调度使用虚拟线程将程序并行化，以暴露流水线并行性，从而隐藏内存访问延迟。结果如图10所示，以roofline图【【47】，Roofline: An insightful visual performance model for multicore architectures，2009，Commun. ACM】的形式呈现；roofline性能图提供了关于一个给定系统在不同基准测试中利用计算和内存资源情况的洞察。总的来说，延迟隐藏改善了所有ResNet层的性能。峰值计算利用率从没有延迟隐藏的70%增加到有延迟隐藏的88%。
5 自动化优化
自动化挑战。给定丰富的调度原语集，我们剩下的问题是为深度学习模型的每一层找到最优的算子实现。在这里，TVM为每一层特定的输入形状和布局创建一个专门的算子。这种特化提供了显著的性能优势（与针对较少形状和布局多样性的手工代码相比），但它也带来了自动化挑战。系统需要选择调度优化——例如修改循环顺序或为内存层次结构进行优化——以及特定于调度的参数，如分块大小和循环展开因子。这些组合选择为每个硬件后端创建了一个巨大的算子实现搜索空间。
自动化优化框架。为了应对这一挑战，我们构建了一个自动化的调度优化器，它有两个主要组成部分：一个提出有前途的新配置的调度探索器，以及一个预测给定配置性能的机器学习成本模型。本节描述了这些组件和TVM的自动化优化流程（图11）。
5.1 调度空间规范
模板化API。我们构建了一个调度模板规范API，让开发者可以在调度空间中声明可调参数（knobs）。模板规范允许在指定可能的调度时，根据需要融入开发者的领域特定知识。我们还为每个硬件后端创建了一个通用的主模板，该模板根据使用张量表达式语言表达的计算描述自动提取可能的旋钮。在高层次上，我们希望考虑尽可能多的配置，并让优化器来管理选择的负担。因此，在我们的实验中，优化器必须为现实世界的深度学习工作负载搜索数十亿种可能的配置。
5.2 基于机器学习的成本模型
自动化方法的比较。从大型配置空间中找到最佳调度的一种方法是通过黑盒优化，即自动调优。这种方法被用于调优高性能计算库【【15】，Fftw: an adaptive software architecture for the fft，1998，IEEE International Conference on Acoustics, Speech and Signal Processing】，【【46】，Automatically tuned linear algebra software，1998，Supercomputing Conference】。然而，自动调优需要大量实验来确定一个好的配置。另一种方法是为特定硬件后端构建一个预定义的成本模型来指导搜索，而不是运行所有可能性并测量其性能。理想情况下，一个完美的成本模型会考虑所有影响性能的因素：内存访问模式、数据重用、流水线依赖和线程模式等。不幸的是，由于现代硬件日益复杂，这种方法很繁琐。此外，每个新的硬件目标都需要一个新的（预定义的）成本模型。
表1：自动化方法比较。模型偏差指建模导致的不准确性。
统计方法。我们采用一种统计方法来解决成本建模问题。在这种方法中，调度探索器提出可能改善算子性能的配置。对于每个调度配置，我们使用一个机器学习模型，该模型将降级后的循环程序作为输入，并预测其在给定硬件后端上的运行时间。该模型使用在探索过程中收集的运行时测量数据进行训练，不需要用户输入详细的硬件信息。我们随着在优化过程中探索更多配置而定期更新模型，这也提高了对其他相关工作负载的准确性。通过这种方式，机器学习模型的质量随着实验次数的增加而提高。表1总结了各种自动化方法之间的关键差异。基于机器学习的成本模型在自动调优和预定义成本建模之间取得了平衡，并可以从相关工作负载的历史性能数据中受益。
机器学习模型设计选择。在选择调度探索器将使用的机器学习模型时，我们必须考虑两个关键因素：质量和速度。调度探索器频繁查询成本模型，这会因模型预测时间和模型重新拟合时间而产生开销。为了有用，这些开销必须小于在真实硬件上测量性能所需的时间，后者根据具体的工作负载/硬件目标可能在秒的量级。这一速度要求将我们的问题与传统的超参数调优问题区分开来，在后者中，进行测量的成本相对于模型开销非常高，可以使用更昂贵的模型。除了模型的选择，我们还需要选择一个目标函数来训练模型，例如配置预测运行时间的误差。然而，由于探索器仅根据预测的相对顺序（A比B快）选择顶尖候选者，我们不需要直接预测绝对执行时间。相反，我们使用一个排序目标（rank objective）来预测运行时成本的相对顺序。
模型实现。我们在我们的ML优化器中实现了几种类型的模型。我们采用了一个梯度树提升模型（基于XGBoost【【8】，Xgboost: A scalable tree boosting system，2016，KDD】），它基于从循环程序中提取的特征进行预测；这些特征包括每个循环级别上每个内存缓冲区的内存访问次数和重用率，以及循环注解（如“vectorize”、“unroll”和“parallel”）的独热编码。我们还评估了一个使用TreeRNN【【38】，Improved semantic representations from tree-structured long short-term memory networks，2015，arXiv】来总结循环程序AST的神经网络模型，无需特征工程。图13总结了成本模型的工作流程。我们发现树提升和TreeRNN具有相似的预测质量。然而，前者的预测速度快两倍，训练时间成本也低得多。因此，在我们的实验中，我们选择梯度树提升作为默认的成本模型。尽管如此，我们相信两种方法都很有价值，并期待未来对这个问题进行更多的研究。平均而言，树提升模型在0.67毫秒内完成预测，比进行实际测量快数千倍。图12将基于ML的优化器与黑盒自动调优方法进行了比较；前者比后者更快地找到更好的配置。
5.3 调度探索
探索策略。一旦我们选择了一个成本模型，我们就可以用它来选择有希望的配置，并迭代地进行实际测量。在每次迭代中，探索器使用ML模型的预测来选择一批候选配置进行测量。收集到的数据随后用作训练数据来更新模型。如果不存在初始训练数据，探索器会选择随机候选进行测量。最简单的探索算法是枚举并通过成本模型运行每个配置，选择预测性能最好的前k个。然而，当搜索空间很大时，这种策略变得不可行。
并行模拟退火。相反，我们运行一个并行的模拟退火算法【【22】，Optimization by simulated annealing，1983，Science】。探索器从随机配置开始，在每一步，随机走到一个附近的配置。如果成本模型预测成本降低，则此转换成功。如果目标配置的成本更高，则很可能失败（拒绝）。随机行走倾向于收敛到成本模型预测成本较低的配置。探索状态在成本模型更新之间保持不变；我们在这些更新后从最后一个配置继续。
5.4 分布式设备池和RPC
扩展测量。一个分布式设备池可以扩展硬件试验的运行规模，并能在多个优化任务之间实现细粒度的资源共享。TVM实现了一个定制的、基于RPC的分布式设备池，使客户端能够在特定类型的设备上运行程序。我们可以使用这个接口在主机编译器上编译程序，请求一个远程设备，远程运行函数，并在主机的同一脚本中访问结果。
RPC功能。TVM的RPC支持动态上传和运行使用其运行时约定的交叉编译模块和函数。因此，相同的基础设施可以执行单个工作负载优化和端到端图推理。我们的方法自动化了跨多个设备的编译、运行和性能分析步骤。这种基础设施对于嵌入式设备尤其关键，这些设备传统上需要繁琐的手动工作来进行交叉编译、代码部署和测量。
A4 实验环境与结果
实验环境
- 模型与工作负载：实验基于真实的深度学习推理工作负载，包括ResNet【【16】，Identity mappings in deep residual networks，2016，arXiv】、MobileNet【【19】，Mobilenets: Efficient convolutional neural networks for mobile vision applications，2017，arXiv】、LSTM语言模型【【48】，Recurrent neural network regularization，2014，arXiv】、深度Q网络（DQN）【【28】，Human-level control through deep reinforcement learning，2015，Nature】和深度卷积生成对抗网络（DCGAN）【【31】，Unsupervised representation learning with deep convolutional generative adversarial networks，2015，arXiv】。具体算子配置见表2。
- 硬件配置：
- 服务器级GPU：NVIDIA Titan X。
- 嵌入式GPU：ARM Mali-T860MP4（在Firefly-RK3399板上）。
- 嵌入式CPU：ARM Cortex A53（四核1.2GHz）。
- FPGA加速器：一个在低功耗PYNQ板上实现的VDLA（Vanilla Deep Learning Accelerator）设计，该板包含一个ARM Cortex A9双核CPU（667MHz）和一个基于Artix-7的FPGA fabric。VDLA设计了一个16x16的矩阵向量单元，时钟频率为200MHz。
- 软件配置：
- TVM实现：核心用C++实现（约5万行代码），提供Python和Java语言绑定。
- 对比基线：
- MXNet (v1.1)
- TensorFlow (v1.7)
- TensorFlow XLA
- TensorFlow Lite (TFLite, commit: 7558b085)
- ARM Compute Library (v18.03)
- Caffe2 (commit: 39e07f7) 用于超低精度算子
- TensorComprehension (TC, commit: ef644ba)
- 依赖库：cuDNN v7，cuBLAS v8。
表2：在单核实验中使用的ResNet-18中所有conv2d算子和MobileNet中所有depthwise conv2d算子的配置。H/W表示高度/宽度，IC输入通道数，OC输出通道数，K核大小，S步幅大小。所有操作均使用“SAME”填充。所有depthwise conv2d操作的通道乘数为1。
实验结果
6.1 服务器级GPU评估
- 实验内容：在NVIDIA Titan X上，对TVM、MXNet、Tensorflow和Tensorflow XLA进行了端到端性能比较。基线框架依赖cuDNN和cuBLAS，而TVM是自动生成和优化代码。
- 实验结果：如图14所示，TVM在所有工作负载上均优于基线，速度提升范围为1.6倍到3.8倍。DQN的3.8倍加速尤为显著，因为它使用了非传统的算子（4x4 conv2d, strides=2），而cuDNN对此类算子优化不佳。
- 分析结论：TVM的优势来自于联合图优化和自动优化器生成的高性能融合算子。即使对于cuDNN已高度优化的标准2D卷积，TVM也能为大多数层生成更好的GPU核（图15）。对于新兴的depthwise卷积，TVM和TC都能找到比MXNet手工核更快的实现，这得益于TVM对大调度空间的探索和有效的基于ML的搜索算法。
6.2 嵌入式CPU评估
- 实验内容：在ARM Cortex A53上，将TVM与Tensorflow Lite（TFLite）进行比较。
- 实验结果：图17显示，TVM生成的算子在ResNet和MobileNet上均优于TFLite手工优化的版本。端到端比较（图16）也显示TVM在三个工作负载上都优于TFLite。
- 超低精度算子：实验还展示了TVM支持低于8位的超低精度推理的能力。通过实现ARM特定的张量化内联函数，TVM生成的代码在2位激活、1位权重的ResNet推理上，其单线程版本优于Caffe2的手工优化库，多线程版本则进一步提升了性能（图18）。
- 分析结论：TVM能够为嵌入式CPU快速优化新兴的张量算子，并且通过张量化机制，能够灵活支持和优化各种精度的算子，甚至超越手工优化的库。
6.3 嵌入式GPU评估
- 实验内容：在ARM Mali-T860MP4 GPU上，将TVM与厂商提供的ARM Compute Library进行端到端比较。
- 实验结果：如图19所示，对于float16和float32两种数据类型，TVM在所有可用模型上均优于基线，速度提升范围为1.2倍到1.6倍。
- 分析结论：TVM在移动GPU平台上也表现出强大的性能可移植性。
6.4 FPGA加速器评估
- 实验内容：在PYNQ板上实现了一个名为VDLA的通用推理加速器（图20），并使用TVM为该新硬件生成ResNet推理代码。VDLA是一个张量处理器，具有专门的片上内存和显式同步控制。
- 实验结果：我们将ResNet中的卷积层卸载到VDLA上执行。对于可卸载的卷积层，FPGA实现了相对于Cortex A9 CPU 40倍的加速。然而，由于Amdahl定律，整个系统的性能受限于必须在CPU上执行的部分（如第一个卷积层、残差连接和激活函数），如图21所示。
- 分析结论：这个实验展示了TVM能够适应新架构及其硬件内联函数，并为专门加速器生成高效代码。增加后端（VDLA）的支持仅需约2000行Python代码。
A7 补充细节
与深度学习框架的关系。深度学习框架【【3】，Tensorflow: A system for large-scale machine learning，2016，OSDI】，【【4】，An introduction to computational networks and the computational network toolkit，2014，MSR-TR】，【【7】，Theano: new features and speed improvements，2012，NIPS Workshop】，【【9】，MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems，2015，NIPS Workshop】为用户提供了便捷的接口来表达和部署深度学习工作负载。虽然现有框架目前依赖于特定供应商的张量算子库来执行其工作负载，但它们可以利用TVM的堆栈为更多的硬件设备生成优化代码。
与高级计算图DSL的关系。高级计算图DSL是表示和执行高级优化的典型方式。Tensorflow的XLA【【3】，Tensorflow: A system for large-scale machine learning，2016，OSDI】和最近引入的DLVM【【45】，Dlvm: A modern compiler infrastructure for deep learning systems，2017，arXiv】属于此类。这些工作中计算图的表示是相似的，本文也使用高级计算图DSL。虽然图级表示非常适合高级优化，但对于在各种硬件后端下优化张量算子来说，它们的级别太高。先前的工作依赖于特定的降级规则直接生成低级LLVM，或求助于供应商精心制作的库。这些方法对于每个硬件后端和算子变体的组合都需要大量的工程努力。
与领域特定语言（DSL）和编译技术的关系。Halide【【32】，Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，PLDI】引入了将计算与调度分离的思想。我们采纳了Halide的见解，并在我们的编译器中重用了其现有的有用调度原语。我们的张量算子调度也与其他针对GPU的DSL工作【【18】，Futhark: Purely functional gpuprogramming with nested parallelism and in-place array updates，2017，PLDI】，【【24】，Loo.py: transformation-based code generation for GPUs and CPUs，2014，ARRAY Workshop】，【【36】，Lift: A functional data-parallel ir for high-performance gpu code generation，2017，CGO】，【【37】，Optiml: An implicitly parallel domain-specific language for machine learning，2011，ICML】以及基于多面体的循环变换【【6】，Pencil: A platform-neutral compute intermediate language for accelerator programming，2015，PACT】，【【43】，Polyhedral parallel code generation for cuda，2013，ACM Trans. Archit. Code Optim.】相关。TACO【【23】，The tensor algebra compiler，2017，Proc. ACM Program. Lang.】引入了一种在CPU上生成稀疏张量算子的通用方法。Weld【【30】，Weld: Rethinking the interface between data-intensive applications，2017，arXiv】是用于数据处理任务的DSL。我们专注于解决针对GPU和专用加速器的深度学习工作负载的新调度挑战。我们的新原语有可能被这些工作中的优化流程所采用。
与自动化调优和成本模型的关系。像ATLAS【【46】，Automatically tuned linear algebra software，1998，Supercomputing Conference】和FFTW【【15】，Fftw: an adaptive software architecture for the fft，1998，IEEE International Conference on Acoustics, Speech and Signal Processing】这样的高性能库使用自动调优来获得最佳性能。Tensor comprehension【【42】，Tensor comprehensions: Frameworkagnostic high-performance machine learning abstractions，2018，arXiv】将黑盒自动调优与多面体优化相结合来优化CUDA内核。OpenTuner【【5】，Opentuner: An extensible framework for program autotuning，2014，PACT】和现有的超参数调优算法【【26】，Efficient hyperparameter optimization and infinitely many armed bandits，2016，arXiv】应用了领域无关的搜索。在Halide中，使用预定义的成本模型来自动调度图像处理流水线【【29】，Automatically scheduling halide image processing pipelines，2016，ACM Trans. Graph.】。TVM的ML模型使用有效的、考虑程序结构的领域感知成本建模。基于此的分布式调度优化器可扩展到更大的搜索空间，并能在广泛支持的后端上找到最先进的内核。更重要的是，我们提供了一个端到端的堆栈，可以直接从深度学习框架中获取描述，并与图级堆栈一起进行联合优化。
与加速器编译的关系。尽管用于深度学习的加速器【【11】，Dadiannao: A machine-learning supercomputer，2014，MICRO】，【【21】，In-datacenter performance analysis of a tensor processing unit，2017，ISCA】日益普及，但如何构建一个编译堆栈以有效地针对这些设备仍不清楚。我们评估中使用的VDLA设计提供了一种通用方法来总结类TPU加速器的特性，并为如何为加速器编译代码提供了一个具体的案例研究。我们的方法也可能使现有将深度学习编译到FPGA的系统受益【【34】，From high-level deep neural models to fpgas，2016，MICRO】，【【40】，FINN: A framework for fast, scalable binarized neural network inference，2016，arXiv】。本文提供了一种通过张量化和编译器驱动的延迟隐藏来有效针对加速器的通用解决方案。
A5 结论
本文提出了一个端到端的编译栈，旨在解决深度学习在多样化硬件后端上部署时面临的基本优化挑战。我们的系统包含自动化的端到端优化能力，而这在历史上是一项劳动密集且高度专业化的任务。通过引入新的调度原语、基于机器学习的成本模型以及图级与算子级联合优化，TVM能够为CPU、GPU、移动设备和新型加速器生成高性能代码，其性能与甚至超越了手工优化的库。
我们希望这项工作能鼓励更多关于端到端编译方法的研究，并为深度学习系统的软硬件协同设计技术开辟新的机会。
💬 评论讨论
欢迎在这里分享您的想法和见解！