大家好，我是极智视界，本文解读一下一种用于深度学习的端到端自动优化编译器TVM。
现在越来越需要将机器学习部署到各种硬件设备，当前的框架依赖于供应商特定的算子，并只是针对范围较窄的服务器级的GPU进行了优化。将工作负载部署到新平台：例如手机、嵌入式设备和加速器(例如FPGA、ASIC等)，需要大量的手工工作。作者提出了TVM，这是一种公开了graph级别和算子级别的优化，在跨不同硬件后端的深度学习工作负载的情况下也能表现优秀的性能。TVM解决了深度学习特有的优化挑战，例如高级算子融合、映射到任意硬件原语和内存延迟隐藏。它还通过采用一种新颖的、基于学习的成本建模方法来快速搜索代码优化，从而自动优化低级程序以适应硬件特性。实验结果表明，TVM提供的跨硬件后端性能可与用于低功耗CPU、移动GPU和服务器级GPU的最先进的手动调整库相媲美。另外作者还展示了TVM针对新加速器后端的适应能力，例如基于FPGA的通用深度学习加速器。值得一提的是，TVM是开源的，并已经在几家大厂内用于生产。
深度学习模型(DL)现在可以用来识别图像、处理自然语言，并在具有挑战性的策略游戏中击败人类。从云服务器到自动驾驶汽车和嵌入式设备，将智能应用程序部署到各种设备的需求在不断增长。由于硬件特性的多样性，包括嵌入式CPU、GPU、FPGA和ASIC(例如TPU)，将DL工作负载映射到这些设备变得复杂。这些硬件目标在内存组织、计算功能单元等方面存在差异，如图1所示。
当前的DL框架，如TensorFlow、MXNet、Caffe和PyTorch，依靠计算图中间表示来实现优化，例如自动微分和动态内存管理。然而，图级优化通常过于高级，无法处理特定于硬件后端的算子级别的转换。这些框架中的大多数都专注于一小类服务器级GPU设备，针对特定目标的优化依赖于高度工程化和特定于供应商的算子库。这些算子级别的库需要大量的手动调整，因此过于专业且不透明，无法轻松跨硬件设备移植。目前，在各种DL框架中为各种硬件后端提供支持需要大量的工程努力。即使对于已经支持的后端，框架也必须在以下之间做出艰难的选择：(1)避免产生不在预定义算子库中的新算子的图优化；(2)使用这些新算子的未优化的视线。为了对不同的硬件后端启用图级和算子级的优化，作者采用了一种完全不同的端到端的方法。作者构建了TVM，这是一个编译器，它从现有框架中获取深度学习程序的高级规范，并为各种硬件后端生成低级优化代码。为了吸引更多用户，TVM需要提供与跨不同硬件后端的大量手动优化的算子库相比具有竞争力的性能。这一目标需要解决下面描述的关键挑战。
LevaragingspecificHardwareFeaturesandAbstractionsDL加速器引入了优化的张量加速原语，而GPU和CPU则不断改进其处理元素。这对为给定的算子描述生成优化代码提出了重大挑战。硬件指令的输入是多维的，具有固定或可变长度；他们决定了不同的数据布局；他们对内存层次结构有特殊要求。系统必须有效地利用这些复杂的原语才能从加速中受益。此外，加速器设计通常也支持更加精简的控制，并将大多数调度复杂性下放到编码器堆栈。对于专用加速器，系统需要生成显式控制管道依赖关系的代码，以隐藏内存访问延迟-这是硬件为CPU和GPU执行的一项工作。
LargeSearchSpaceforOptimization另一个挑战是在不手动调整算子的情况下生成高效的代码。内存访问、线程模式和新颖的硬件原语的组合选择为生成的代码(例如循环块、排序、缓存和展开)创建了巨大的配置空间，如果要实现黑盒自动调整，将需要大量的搜索成本。可以采用预定义的成本模型来指导搜索，但由于现代硬件的复杂性日益增加，构建准确的成本模型很困难。此外，这种方法需要我们为每种硬件类型建立单独的成本模型。
通过结合这三个模块，TVM可以从现有的深度学习框架中获取模型描述，执行高级和低级联合优化，并为后端生成特定于硬件的优化代码，例如CPU、GPU和基于FPGA的专用加速器。本文的贡献如下：
使用服务器级GPU、嵌入式GPU、嵌入式CPU和基于FPGA的定制通用加速器上的真实工作负载评估了TVM。实验结果表明，TVM提供了跨后端的可移植性能，并且比由手动优化库支持的现有框架实现了1.2倍到3.8倍的加速。
本节通过一个示例来介绍TVM的组件。图2总结了TVM中的执行步骤及其在论文中的相应部分。
End-UserExample只需要几行代码，用户就可以从现有的深度学习框架中获取模型并调用TVMAPI以获取可部署的模块：
importtvmast#Usekerasframeworkasexample,importmodelgraph,params=t.frontend.from_keras(keras_model)target=t.target.cuda()graph,lib,params=t.compiler.build(graph,target,params)这个编译后的运行时模块包含三个组件：最终优化的计算图(graph)、生成的算子(lib)和模块参数(params)。然后可以使用这些组件将模型部署到目标后端：
importtvm.runtimeastmodule=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的架构以及系统程序员如何扩展它以支持新的后端。
计算图是在DL框架中表示程序的常用方法。图3展示了一个两层卷积网络的示例计算图表示。
这种高级表示与低级编译器中间表示(IR)(如LLVM)之间的主要区别在于：中间数据项是大型的多维张量。计算图提供了算子的全局视图，但它们避免指定每个算子必须如何实现。与LLVMIR一样，计算图可以转换为功能等效的图以应用优化。作者还利用常见DL工作负载中的形状特异性来优化一组固定的输入形状。
TVM利用计算图表示来应用高级优化：节点表示对张量或程序输入的操作，边表示操作之间的数据依赖关系。它实现了许多图级优化，包括：算子融合，将多个小操作融合在一起；常量折叠，预先计算可以静态确定graph部分，节省执行成本；一个静态内存规划pass，它预先分配内存来保存每个中间张量；数据排布转换，将内部数据排布转换为后端友好的形式。接下来讨论算子融合和数据排布转换。
DataLayoutTransformation有多种方法可以在计算图中存储给定的张量。最常见的数据排布选择是列主序和行主序。在实践中，咱们可能会更加喜欢使用更加复杂的数据排布。例如，DL加速器可能会采用4X4矩阵运算，需要将数据排布平铺成4X4块以优化访问局部性。
数据排布优化将计算图转换为可以使用更加好的内部数据排布以在目标硬件上执行计算图。它首先在给定内存层次结构规定的约束条件下为每个算子指定首选数据排布。如果首选数据布局不匹配，就在生产者和消费者之间执行适当的布局转换。
虽然高级图优化可以极大地提高DL工作负载的效率，但他们也只是与算子库差不多有效。目前，支持算子融合的少数DL框架需要算子库来提供融合模式的实现。随着定期引入更多的网络算子，可能的融合内核数量会急剧增加。当需要面对越来越多的硬件后端时，这种方法不再可持续，因为所需要的融合模式实现数量与必须支持的数据排布、数据类型和加速器内在函数的数量相结合。为程序所需的各种操作和每个后端手工设计算子内核是不可行的。为此，接下来提出了一种代码生成方法，可以为给定模型的算子生成各种可能的实现。
TVM通过在每个硬件后端生成许多有效的实现并选择优化的实现，为每个算子生成高效的代码。这个过程建立在Halide将描述和计算规则(或调度优化)解耦的想法的基础上，并将其扩展为支持新的优化(嵌套并行、张量和延迟隐藏)和广泛的硬件后端。
引入了张量表达式语言来支持自动代码生成，与高级计算图表示不同，张量算子的实现是不透明的，每个算子都用索引公式表达语言来描述。以下代码展示了计算转置矩阵乘法的示例张量表达式。
每个计算操作都指定输出张量的形状和描述如何计算它的每个元素的表达式。张量表达式语言支持常见的算术和数学运算，并涵盖常见的DL算子形式。该语言没有指定循环结构和许多其他执行细节，它提供了为各种后端添加硬件感知优化的灵活性。采用来自Halide的解耦计算/调度原则，使用调度来表示从张量表达式到低级代码的特定映射。
通过增量应用保持程序逻辑等价的基本转换(调度原语)来构建调度。图5展示了在专用加速器上调度矩阵乘法的示例。在内部，TVM应用调度转换，使用数据结构来跟踪循环结构和其他信息。然后此信息可以帮助按给定的最终调度生成低级代码。
作者的张量表达式借鉴了Halide、Darkroom和TACO。为了在许多后端实现高性能，必须要支持足够多的调度原语来涵盖不同硬件后端的各种优化。图6总结了TVM支持的算子代码生成过程和调度原语。作者重用了有用的原语和来自Halide的低级循环程序AST，并且引入了新的原语来优化GPU和加速器的性能。新的原语是实现最佳GPU性能所必需的，也是加速器所必需的。CPU、GPU和类TPU加速器是三种重要的深度学习硬件。
并行性是提高DL工作负载中计算密集型内核效率的关键。现代GPU提供了大规模的并行性，要求咱们将并行模式融合到调度转换中。大多数现有的解决方案都采用一种称为嵌套并行的模型，这是一种fork-join的形式。该模型需要并行调度原语来并行化数据并行任务；每个任务可以进一步递归地细分为子任务，以利用目标架构的多级线程层次结构(例如GPU中的线程组)。将此模型称为无共享嵌套并行，因为一个工作线程无法在同一并行计算阶段查看其兄弟的数据。
无共享方法的替代方法是协作获取数据。具体来说，线程组可以协作获取它们都需要的数据并将其放入共享内存空间中。这种优化可以利用GPU内存层次结构，并通过共享内存区域实现跨线程的数据重写。TVM使用调度原语来支持这种众所周知的GPU优化，以实现最佳性能。以下GPU代码示例对矩阵乘法进行了优化。
图7展示了这种优化的影响，将内存范围的概念引入调度空间，以便可以将计算阶段(代码中的AS和BS)标记为共享。
如果没有显式内存范围，自动范围推理会将计算阶段标记为thread-local。共享任务必须计算组中所有工作线程的依赖关系，此外还必须正确插入内存同步屏障，以保证共享加载的数据对消费者可见。最后，除了对GPU有用之外，内存范围还让咱们可以标记特殊的内存缓冲区，并在针对专门的DL加速器时创建特殊的降低规则。
DL工作负载具有很高的算术强度，通常可以分解为张量算子，如矩阵乘或一维卷积。这些自然分解导致了最近添加的张量计算原语的趋势。这些新的原语为基于调度的编译创造了机遇和挑战，使用它们可以提高性能，所以编译框架必须无缝的集成它们。这种张量化类似于SIMD架构的矢量化，但又有着显著差异。指令输入是多维的，具有固定或可变的长度，并且具有不同的数据排布。更为重要的是，这种原语不能是固定的，因为新的加速器有可能会出现张量指令的变体，因此需要一个可扩展的解决方法。
延迟隐藏是指将内存操作与计算重叠以最大限度地利用内存和计算资源的过程。它需要不同的策略，具体取决于目标硬件后端。在CPU上，内存延迟隐藏是通过同步多线程或硬件预取来隐式实现的。GPU依赖于许多线程的快速上下文切换。相比之下，诸如TPU之类的专用DL加速器通常倾向于使用解耦访问执行(DAE)架构记性更加精简的控制，并将细粒度同步的问题转移给软件。
图9显示了减少运行时延迟的DAE硬件管道。与单片硬件设计相比，流水线可以隐藏大部分内存访问开销，几乎可以充分利用计算资源。为了实现更加高的利用率，指令流必须增加细粒度的同步操作。没有它们，就无法强制执行依赖关系，从而导致错误执行。因此，DAE硬件流水线需要在流水线阶段依赖于细粒度的入队/出队操作，以保证正确执行，如图9的指令流所示。
对需要显式低级同步的DAE加速器进行编程很困难。为了减少编程负担，引入了虚拟线程调度原语，让程序员可以指定高级数据并行程序，就像他们指定支持多线程的硬件后端一样。然后TVM会自动将程序降低为具有低级显式同步的单个指令流，如图8所示。该算法从高级多线程程序调度开始，然后插入必要的低级同步操作以保证每个线程的正确执行。接下来，它将所有虚拟线程的操作交付给单个指令流中。最后，硬件恢复指令流中低级同步所规定的可用流水线的并行度。
HardwareEvaluationofLatencyHiding展示了延迟隐藏在基于FPGA的定制加速器设计上的有效性。在加速器上运行ResNet的每一层，并使用TVM生成两个调度：一个有延迟隐藏，一个没有。具有延迟隐藏的调度将程序与虚拟线程并行化以进行管道并行，从而隐藏了内存访问延迟。结果在图10中显示为roofline图表；roofline性能图可以深入了解给定系统在不同基准测试中使用计算和内存资源的情况。总体而言，延迟隐藏提高了所有ResNet层的性能。峰值计算利用率从没有延迟隐藏的70%提高到了隐藏延迟的88%。
鉴于丰富的调度原语集合，剩下的问题是为DL模型的每一层找到最佳的算子实现。在这里，TVM为与每一层关联的特定输入形状和布局创建了一个专门的算子。这种专门的优化提供了显著的性能优势(与针对较小的形状和布局多样性的手工代码相比)，但它也提出了自动化的挑战。系统需要选择调度优化：例如修改循环顺序或优化内存层次结构，以及特定于调度的参数，例如切片大小和循环展开因子。这样的组合选择为每个硬件后端创建了算子实现的大的搜索空间。为了应对这一挑战，构建了一个具有两个主要组件的自动调度优化器：一个新配置的调度搜索器以及一个预测给定配置性能的机器学习成本模型。如图11展示了这些组件和TVM的自动优化流程。
从大型配置空间中找到最佳调度的一种方法是通过黑盒优化，即自动调整，该方法用于调优高性能计算库。然而，自动调整需要许多实验来确定一个好的配置。
另一种方法是构建一个预定义的成本模型来指导对特定硬件后端的搜索，而不是运行所有可能性并测量它们的性能。理想情况下，一个完美的成本模型会考虑影响性能的所有因素：内存访问模式、数据重用、管道依赖性和线程模式等。不幸的是，由于现代硬件日益复杂，这种方法很繁琐。此外，每个新的硬件目标都需要一个新的(预定义的)成本模型。
作者在ML优化器中实现了几种类型的模型。采用梯度树提升模型(基于XGBoost)，它根据从循环程序中提取的特征进行预测；这些功能包括每个循环级别的每个内存缓冲区的内存访问计数和重用率，以及循环注释的one-hot编码，例如vectorize、unroll和parallel。作者还评估了一个神经网络模型，该模型使用TreeRNN在没有特征工程的情况下总结循环程序的AST。图13总结了成本模型的工作流程。
一旦咱们选择了一个成本模型，我们就可以使用它来选择建议的配置，在这些配置上迭代地运行实际测量。在每次迭代中，搜索器使用ML模型的预测来选择一批候选者来运行测量。然后将收集的数据用作训练数据以更新模型。如果不存在初始训练数据，则搜索器会选择随机候选者进行测量。
最简单的搜索算法通过成本模型枚举并运行每个配置，选择前k个预测的执行者。但是这种策略再搜索空间很大时变得难以处理。相反，作者采用并行模拟退火算法。搜索器从随机配置开始，并且在每一步中，随机走到附近的配置。如果成本如成本模型所预测的那样降低，则这种转变是成功的。如果目标配置成本较高，则可能会失败(拒绝)。这样随机游走在倾向于收敛于成本模型预测的成本较低的配置。搜索状态在成本模型更新中持续存在；咱们从这些更新后的最后一个配置继续搜索。
分布式设备池可以扩展硬件试验的运行，并在多个优化作业之间实现细粒度的资源共享。TVM实现了一个定制的、基于RPC的分布式设备池，使客户端能够在特定类型的设备上运行程序。咱们可以使用此接口在主机编译器上编译程序，请求远程设备、远程运行函数，并在主机上访问相同脚本中的结果。TVM的RPC支持动态上传并运行使用其运行时约定的交叉编译模块和函数。因此，相同的基础架构可以执行单个工作负载优化和端到端graph推理。作者的方法可以跨多个设备自动执行编译、运行和配置步骤。这种基础设施对于嵌入式设备尤其重要，因为传统上这些设备需要繁琐的手动操作来进行交叉编译、代码部署和测量。
图14展示了TVM、MXNet、Tensorflow和TensorflowXLA在NVIDIATitanX上的性能评估。
图15展示了cuDNN、TensorComprehensions、MXKernel、TVM、TVMPT在TITANX上对ResNet18中所有conv2d算子，和MobileNet中所有深度conv2d算子的相对加速情况。
图16展示了TVM和TFLite在ARMA53上端到端性能评估。
图17展示了cuDNN、TensorComprehensions、MXKernel、TVM、TVMPT在ARMA53上对ResNet18中所有conv2d算子，和MobileNet中所有深度conv2d算子的相对加速情况。
图18展示了手工优化、TVM单线程和TVM多线程对ResNet中conv2d算子的相对加速。
图19展示了在Mali-T860MP4的端到端实验结果。
图20展示了VDLA硬件设计总览。
图21展示了将ResNet工作负载中的卷积卸载到基于FPGA的加速器上的加速效果。
作者提出了一个端到端的编译堆栈，以解决跨各种硬件后端的深度学习的基本优化挑战。提出的系统包括自动化的端到端的优化，这在以前是一项劳动密集型和高度专业化的任务。作者希望这项工作将鼓励对端到端编译方法的更多研究，并为DL系统软硬件协同设计技术开辟新的机会。
[1]TVM:AnAutomatedEnd-to-EndOptimizingCompilerforDeepLearning.