端到端TVM编译器（上）

端到端TVM编译器（上）

摘要

将机器学习引入到各种各样的硬件设备中。AI框架依赖于特定于供应商的算子库，针对窄范围的服务器级gpu进行优化。将工作负载部署到新平台，例如手机、嵌入式设备和加速器（例如，FPGA、ASIC）–需要大量手动操作。TVM，一个开源图形级的编译器和算子级优化，提供可移植到不同领域的深度学习工作负载性能硬件后端。TVM解决了特定于深度学习的优化挑战，例如高级算子融合、映射到任意硬件原语，存储潜伏期隐藏。通过采用一种新颖的基于学习的成本建模方法，用于快速探索代码优化。实验表明，TVM在多个领域都有很好的性能，富于竞争力的硬件后端，用于低功耗CPU、移动GPU和服务器级GPU的手动调优库。展示了TVM瞄准新加速器后端的能力，例如，基于FPGA的通用深度学习加速器。系统是开源的，帮助大公司产品量产落地。

1. Introduction

深度学习（DL）模型可以识别图像，处理自然语言，在具有挑战性的策略游戏中击败人类。

各种设备，从云服务器到自动驾驶汽车和嵌入式设备上，部署智能应用的需求越来越大。将由于硬件特性的多样性，包括嵌入式CPU、GPU、FPGA和ASIC（例如，TPU），DL工作负载映射到这些设备是非常困难的。

图1:CPU、GPU和类似TPU的加速器需要不同的片上内存架构和计算原语。在生成优化代码时，必须解决这种分歧。

如图1所示，这些硬件目标在内存组织，计算功能单元。当前的DL框架，如TensorFlow，MXNet，Caffe和PyTorch，依靠计算图中间表示来实现优化，例如，自动微分和动态内存管理。然而，图形级优化，通常太高，无法处理特定于硬件后端的算子级转换。其中大部分框架侧重于GPU设备服务器的一个狭窄类和代理特定于目标的优化到高度工程化和供应商特定的算子库。这些算子级库需要手动tune，过于专业化和不透明，不易于跨硬件设备移植。提供目前需要大量的工程工作后端，支持各种不同硬件的DL框架。即使对于支持的后端，框架也必须做出艰难的选择：

（1）避免在预定义的算子库中生成新算子的图优化

（2）使用这些新算子的未优化实现。

为了对不同的硬件后端实现图形级和算子级优化，采用了一种根本不同的端到端方法。构造了TVM，一种采用高级规范的编译器，从现有的框架中获取深度学习程序的高级规范，并为一组不同的硬件后端生成低级优化代码。

为了吸引用户，TVM需要在不同的硬件后端提供与大量手动优化的算子库，具有竞争力的性能。

这一目标要求解决下述关键挑战。

利用特定的硬件特性和抽象。

DL加速器引入了优化的张量计算原语，而GPU和CPU不断改进其处理元素。这对为给定的算子描述生成优化代码，提出了一个重大挑战。

硬件指令的输入是多维的，可以是固定的，也可以是可变的长度；

规定了不同的数据布局；

有内存层次结构的特殊要求。

系统必须有效地利用这些复杂的原语，才能从加速中获益。此外，加速器设计通常也倾向于精简控制，并将大部分调度复杂性转移到编译器堆栈中。对于专用的加速器，系统需要生成显式控制到的管道依赖关系的代码，隐藏CPU和GPU执行的任务的内存访问延迟–硬件。

大范围的搜索空间优化

另一个挑战，在没有手动调整算子的情况下，生成高效的代码。如果执行black盒子自动调谐内存访问、线程模式，新硬件原语的组合选择为生成的代码创建巨大的配置空间（例如，循环平铺和排序、缓存、展开）。

可以采用预先确定的cost model来指导搜索，但构建cost model准确模型是困难的，因为越来越复杂的modern硬件问题。此外，这种方法将要求我们为每种硬件类型建立单独的cost model。

TVM通过三个关键模块解决了这些挑战

(1）引入了张量表达式语言，构建算子，提供程序转换。通过各种优化生成不同版本程序的原语。该层扩展了Halid的计算/调度分离概念，将目标硬件内部函数与转换原语分离，从而支持新的加速器及其相应的新内部函数。此外，引入了新的转换原语，解决GPU相关的挑战，并支持部署到专用加速器。然后，可以应用不同的程序转换序列，为给定的算子声明形成一个丰富的有效程序空间

(2) 引入了一个自动程序优化框架，搜索优化的张量算子。优化器是在基于ML的cost model的指导下，从硬件后端收集更多数据，该模型会进行调整和改进

(3）在自动代码生成器上，引入一个充分高级和算子级优化。

结合这三个模块，TVM可以从现有的深度学习框架中进行模型描述，执行高级和低级联合优化，为后端（如CPU、GPU和基于FPGA的专用处理器）生成特定于硬件的优化代码加速器。

本文的贡献如下：

• 确定了在为跨不同硬件后端的深度学习工作负载，提供性能可移植性方面的主要优化挑战
• 引入了新颖的调度原语，利用了跨线程内存重用、新颖的硬件内在的，潜在的隐藏。
• 提出并实现了基于机器学习的自动探索和优化系统，搜索优化的张量算子。
• 构建了一个端到端的编译和优化堆栈，允许多种硬件后端（包括CPU、服务器GPU、移动GPU和基于FPGA的加速器），部署深度学习在高级框架（包括TensorFlow、MXNet、PyTorch、Keras、CNTK）中指定的工作负载。

开源的TVM在几家大公司内部使用。使用真实世界的服务器类GPU，嵌入式GPU，嵌入式

CPU和基于FPGA的定制通用加速器工作负载，对TVM进行了评估。

实验表明，TVM具有跨后端的便携性性能，实现范围从1.2× 至3.8× 超越现有框架，支持手工优化库加速。

1. Overview

本节通过使用组件walk示例来描述TVM。

图2: TVM的系统概述。当前的堆栈支持从许多深度学习框架和交换格式（如CoreML和ONNX）到主要CPU、GPU和专用加速器的描述。

图2总结了TVM中的执行步骤及其相应部分

论文。系统首先从一个现有的框架，将其转换为计算图表示。然后执行高级数据流重写，生成优化图。这个算子级优化模块必须为图中的每个融合算子生成有效的代码。算子是用声明性张量表达式语言指定的；未指定执行详细信息。TVM标识一组给定硬件目标的算子代码优化。可能的优化形成一个很大的空间，使用一个基于ML的cost model来搜索优化的算子。最后，系统将生成的代码打包到可部署模块中。

终端用户示例。只需几行代码，用户就可以借鉴现有的深度学习框架，调用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的体系结构和系统的组成程序，可以扩展来支持新的后端。

1. Optimizing Computational Graphs

计算图是一种常用的DL框架中的程序表示方法。

图3：两层结构的计算图示例

卷积神经网络。图中的每个节点表示使用一个或多个张量，生成一个或多个张量的操作。张量运算可以通过属性参数化来配置行为（例如，填充或跨步）。

图3显示了一个两层卷积神经网络的示例计算图表示。这种高级表示和低级编译器中间表示（IR）之间的主要区别，例如LLVM，中间数据项很大，多维张量。计算图提供了算子的全局视图，避免指定每个算子必须如何实现。像在IRs中LLVM，一个计算图可以转化为功能等价图来应用优化。可以利用普通DL的形状特异性，实现一组固定的输入形状进行优化的工作负载。

应用高级优化，TVM利用计算图表示：节点表示对张量或程序输入的算子，边表示操作之间的数据依赖关系。

实现了许多图级优化，包括：

算子融合，将多个小算子融合在一起；

常量折叠，预先计算可静态确定，节省执行成本；

静态内存规划过程，预先分配内存来容纳每个中间张量；

数据布局转换，用于转换内部数据布局到后端友好的形式。

现在讨论算子融合和数据布局转换。

算子融合。

算子融合结合了多个算子，不将中间结果保存在内存中。这种优化可以大大减少执行时间，特别是在GPU和专用加速器中。具体来说，认识四个图算子的范畴：

（1）内射（一对一）map，例如add。

（2）reduction（例如sum）

（3）complexout
fusable（可以将元素映射到输出，例如，conv2d）

（4）不透明（不能融合，例如排序）。

提供通用规则来融合这些算子，如下所示。

图4：融合和非熔合算子。TVM生成这两个算子。在NVIDIA Titan X上测试。

多个内射算子可以融合成另一个内射算子。还原算子，可以与输入内射算子（例如，fuse
scale和sum）。像conv2d这样的算子是复合输出融合的，可以将元素级算子融合到输出中。可以应用这些规则，变换计算图变成融合版。图4展示了影响在不同的工作负载上进行优化。发现了融合算子生成高达1.2× 到2× 加速，减少内存访问。

数据布局转换。

有多个在计算图中存储给定张量的方法。最常见的数据布局选择是列主布局和行布局。实际上，可能更喜欢使用更复杂的数据布局。例如，DL加速器可能利用4×4矩阵运算，需要数据平铺为4×4个区块以优化访问地址。

数据布局优化。将计算图形转换为一个可以使用更好的内部数据布局的图形，在目标硬件上执行。首先为给定的每个算子指定首选的数据布局，由内存层次结构决定的约束。如果生产者和消费者的首选数据布局确实如此，则在生产者和消费者之间执行适当不匹配的布局转换。

图5：在专用加速器上优化矩阵乘法的调度转换示例。

虽然高级图优化可以极大地提高DL工作负载的效率，但只是与算子库提供的一样有效。目前，为数不多的支持算子融合的DL框架，要求算子库提供融合模式的实现。随着更多的网络算子定期引入，可能的融合的结果可以迅速生长。这种方法是，当针对越来越多的硬件后端时，不再可持续融合模式实现的数量组合增长，必须支持的数据布局、数据类型和加速器内部函数的数量。对于一个程序和每个后端所需的各种操作，手工制作算子内核是不可行的。接下来提出一种代码生成方法，可以为一个给定模型的算子。

1. Generating Tensor Operations

TVM通过在每个硬件上生成许多有效的实现，为每个算子生成有效的代码后端和选择优化的实现。

这个过程建立在Halide将描述与计算规则（或调度优化）解耦的思想基础上，扩展以支持新的优化（嵌套并行、张量化和延迟隐藏），以及各种各样的硬件后端。现在重点介绍TVM的特定功能。

4.1张量表达式与表空间

图6: TVM调度降低和代码生成过程。下表列出了现有的Halid和新型TVM用于优化调度的调度原语，用于CPU、GPU和加速器后端。张量化对加速器来说是必要的，但也可以使用

用于CPU和GPU。特殊内存范围启用GPU中的内存重用和加速器中ONCHIP内存的显式管理。延迟隐藏是特定的，就像加速器一样。

引入张量表达式语言来支持自动代码生成。与高级计算不同，在图表示中，张量运算的实现是不透明的，每个运算都用索引公式表达式语言描述。以下代码显示了一个要计算的张量表达式示例。

转置矩阵乘法：

每个计算算子都指定输出张量，描述如何计算每个元素。

张量表达式

语言支持常见的算术和数学运算，并涵盖了常见的DL运算符模式。这个语言没有指定循环结构和许多其它执行细节，并为各种后端添加硬件感知优化。采用Halide中的解耦计算/调度原则，使用调度来表示从张量表达式到低级代码的特定映射。许多可能的调度，可以执行此功能。通过逐步应用basic，保留程序的逻辑等价性。图5显示了在专用加速器上调度矩阵乘法的示例。在内部，当应用调度转换时，TVM使用数据结构跟踪循环结构和其它信息。此信息可以帮助为给定的final生成低级代码调度表。张量表达式从Halid中提取线索，Darkroom，and
TACO。主要增强功能包括，支持下面讨论的新的调度表优化。实现许多后端的高性能，必须支持足够的调度表原语，涵盖不同硬件后端上的一组不同的优化。图6总结了TVM支持的算子代码生成过程和调度原语。重用有用的原语Halid的低级循环程序AST，引入新的原语来优化GPU和加速器的性能。新的原语是必要的，实现最佳的GPU性能和必不可少的加速器。CPU、GPU、TPU类加速器是深度学习的三种重要硬件类型。

本节介绍新的优化原语CPU、GPU和TPU类加速器，而下一节说明如何自动导出有效的调度。

4.2嵌套并行与协作

并行性是提高DL工作负载中的计算密集型内核并行计算效率的关键。现代GPU提供了大量的并行性，需要将并行模式bake焙到调度表转换中。大多数现有的解决方案都采用了一种称为嵌套并行（nestedparallelism）的模型叉形连接形式。该模型需要一个并行调度原语来并行化数据并行任务；每项任务可以进一步递归细分为子任务，利用目标体系结构的多级线程层次结构（例如，GPU中的线程组）。称这个模型为sharednothing嵌套并行，因为一个工作线程不能在同一并行计算阶段中查看其同级的数据。无共享方法的另一种替代方法是协同获取数据。特别是线程组可以协作获取所有需要的数据，并放置一个共享的内存空间。此优化可以利用GPU内存层次结构，通过共享内存实现跨线程的数据重用区域。TVM支持这个众所周知的GPU优化，使用调度原语来实现最佳性能。下面的GPU代码示例优化矩阵乘法。

图7:TVM与在矩阵乘法工作负载上没有协作共享内存获取。在NVIDIA泰坦X上测试。

图7展示了这种优化的影响。将内存作用域的概念引入到调度空间，以便计算阶段（代码）可以标记为共享。没有明确的存储范围，自动范围推断将标记计算阶段作为线程本地。共享任务必须计算组中所有工作线程的依赖关系。

此外，必须设置内存同步屏障正确插入，确保共享加载的数据对消费者可见。最后，除了对GPU有用之外，内存作用域还允许标记特殊内存缓冲区和创建特殊的降低规则时，目标专门的DL加速器。