TVM：一种自动端到端优化的深度学习编译器
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
提出背景
现有的 DL 框架依赖于计算图 IR 来实现优化，比如自动微分（Auto Differentiation）和动态内存管理。然而计算图层级（Graph-level）的优化对硬件后端特定算子层级（Operator-level）的变换来说，往往视角太高了。这些框架大部分关注的都是服务器级 GPU 设备中的一小撮，而把目标机相关的优化委派给高度工程化的、厂家特化的算子库。这类算子库需要大量的人工调优，也因此过于特殊化（普适性差）和不透明，导致不易于在硬件平台之间移植。即使是框架支持的后端，优化计算图时也需要在以下两个选择中权衡：
- 避免引入算子库中未定义的新算子（扩展性差）
- 新算子使用未优化的实现代替（性能降级）
为了使得各种硬件后端的计算图层级和算子层级优化成为可能，TVM 作为一种全新意义上的端到端（End-to-end）方法被提出。TVM 编译器从现有框架中取得 DL 程序的高层级表示，并产生多硬件平台后端上低层级的优化代码。TVM 的目标是展示与人工调优的竞争力，关键的挑战是：
- 平衡特定硬件的特性和抽象。DL 加速器引入了张量计算原语（Tensor Compute Primitives），但 CPU 和 GPU 有各自的数据处理形式（标量、向量），如何将多维度的数据布局通过变换，使之适合处理器和内存层级（Memory Hierarchy），是一个巨大的挑战。此外，加速器设计普遍偏好控制精简性的设计，而将任务调度的复杂性下放到编译栈上。对于某些特定的加速器，编译器甚至需要产生能够显式解决流水线数据依赖的代码，来“隐藏”内存访问的延迟（Memory Access Latency Hiding）。
- 优化存在海量的搜索空间。安排内存访问、线程模式和新硬件原语（Hardware Primitives）等元素是排列组合级别的复杂度，如果实现一个黑箱来自动调优会带来巨大的搜索开销。也许我们可以预定义一个开销模型来指导搜索，但是现代硬件的复杂程度是不断增长的（而且速度很快），精确建模的难度非常大，更不用说还得给每一种硬件类型都独立建模了。
主要工作
TVM 提出了三个重要模块：
- 张量表达语言（Tensor Expression Language）和变换原语（Transformation Primitives）。这是对 Halide 计算调度解耦理念的扩展，把硬件本质（Hardware Intrinsic）和变换原语也分离了，使得 TVM 可以支持新的加速器和对应的硬件本质。
- 自动程序优化框架（Automated Program Optimization Framework）。使用一个基于 ML 的开销模型来指导如何寻找最优张量算子，当收集到更多硬件后端的数据时，模型的适应性和表现会不断提升。
- 计算图重写（Graph Rewriter）。作用于自动代码生成器上层，统筹计算图级和算子级的联合优化。
工作流程
TVM 的工作流程：
- 读取现有框架中的模型，产生计算图表示
- 进行高层级数据流重写（High-level Dataflow Rewriting），产生优化后的计算图
- 进行算子层级优化，为计算图中的融合算子（Fused Operator）产生高效代码
- 算子由张量表达语言声明，执行细节不需要指定
- 利用 ML 开销模型从目标机潜在的优化集合中，搜索算子的最优代码
- 将生成的代码包装为可部署的模块
优化计算图（Optimizing Computational Graphs）
计算图是一种高级表示，提供了对于算子的全局视野，而不需要指明实现的细节。就像 LLVM IR，计算图也可以转换成功能等价的子图，以适合各种优化手段。TVM 对计算图的优化包括：算子融合（Operator Fusion）、常量折叠（Constant Folding）、静态内存规划（Static Memory Planning）、数据布局变换（Data Layout Transformation）等。
算子融合
算子融合应用了多个算子、一次计算的理念，避免保存中间结果时对内存的访问，从而减少执行时间。TVM 将算子分为以下几类：
- 单射（Injective），一对一的运算，例如元素和
- 规约（Reduction），多对一的运算，例如求和
- 复杂运算、输出可融合（Complex-out-fusable），可以将元素映射融合到输出，例如二维卷积
- 不透明（Opaque），不可融合，例如排序
针对这几类算子，TVM 提出了泛用的融合规则：
- 多个单射算子可以融合为单个单射算子
- 一个规约算子可以和多个单射算子融合，例如缩放后求和
- 逐元素（Element-wise）类型的算子可以融合到复杂运算、输出可融合算子的输出
算子融合通过减少访存可以实现 1.2× 到 2× 的加速。
数据布局变换
最常见的数据布局是列优先（Column Major）和行优先（Row Major）的方式。事实上，DL 加速器往往会采用更复杂的数据布局，比如 4×4 的矩阵，为了能够充分利用空间局部性，要求数据能够平铺成 4×4 的“小砖块”。数据布局变换将计算图转换成可以更好地利用内部数据布局的形式，首先需要为每个算子规定来自内存层级的约束，如果数据布局不符合要求，就进行变换，这里采用的是生产者消费者模式。
生成张量运算（Generating Tensor Operations）
虽然计算图优化能极大地提高 DL 工作负载，但是它的效果与算子库提供的算子水平有很大关系。现在支持算子融合的 DL 框架很少有要求算子库也提供算子融合模式的实现，因为随着神经网络算子的不断提出，融合算子的数量也经历了排列组合级别的增长，再考虑到各种不同硬件后端的出现，这种方式明显是不可持续的。出于同样的原因，理想的、多样的算子也不可能经由手工调制，于是张量算子的自动生成就成了迫切需要。
张量表达式（Tensor Expression）和调度空间（Schedule Space）
张量表达式由结果形状和运算规则两部分组成，支持常见的算数运算和 DL 算子。它无需指明循环结构和其他执行细节，提供给硬件后端优化更大的灵活性。
在维持程序逻辑等价性的前提下，TVM 对张量表达式逐次使用基本变换（调度原语），并记录下过程中的循环结构等其他所需信息，这些信息用来帮助生成最终调度（Final Schedule）的低层级代码。
协作嵌套并行（Nested Parallelism with Cooperation）
嵌套并行是 Fork-join 模型的一种形式，指的是每一个子任务都可以递归地被更进一步划分成子任务并行处理，从而深度利用目标架构的多级线程层级（Multi-level Thread Hierarchy），比如 GPU 的线程组（Thread Group）。如果在并行计算阶段中，一个线程无需访问相邻线程的数据，这种模型又被叫做无共享嵌套并行（Shared-nothing Nested Parallelism）。
代替无共享方式的一种选择是协作获取数据：一组线程共同获取到一块数据，然后各取所需。这样做的好处是能够充分利用 GPU 显存层级，通过共享内存也使得线程之间的数据重用成为可能。
TVM 在调度空间种引入了内存作用域（Memory Scope）的概念，计算阶段（Compute Stage）可以被标记为共享（Shared）。如果没有显式的内存作用域，自动作用域推导会把计算阶段标记为线程局部（Thread-local）。共享任务必须计算彼此之间的依赖关系，内存同步屏障（Memory Synchronization Barrier）技术也需要用来保证数据对数据的消费者是可见的。另外，内存作用域还可以标记特殊内存缓存，这对 GPU 来说很有用；当以 DL 加速器为目标机时，内存作用域还可以创建额外的代码低层级化规则。
张量化（Tensorization）
类比向量化（Vectorization）之于 SIMD。
显式内存延迟隐藏（Explicit Memory Latency Hiding）
CPU 隐藏内存延迟的方式是多线程，GPU 隐藏内存延迟的方式是线程组的快速上下文切换，但是特定 DL 加速器（比如 TPU）偏好控制精简型的解耦访问执行（Decoupled Access Execute）架构，会将细粒度的同步控制下放到软件处理。
DAE 架构流水线需要保证正确的依赖关系，可以通过使用细粒度的依赖队列实现。直接在低层级上实现 DAE 加速器的同步控制是相对困难的，TVM 引入了虚拟线程调度原语（Virtual Threading Scheduling Primitive），开发者可以假装指定的硬件后端拥有多线程支持，TVM 来负责插入确保执行顺序所必须的低层级同步操作，并自动生成单指令流。
自动优化（Automating Optimization）
TVM 为 DL 模型的每一层产生针对输入形状和布局优化过的算子，从而带来巨大的性能增益，但如何选择调度优化（比如改变循环顺序、平铺大小、展开因子等）却是排列组合级别的复杂度。为此，TVM 提出了自动调度优化器（Automated Schedule Optimizer），它包含两个主要组件：
- 调度探索器（Schedule Explorer），用来提出潜在的、有前途的优化配置
- 基于机器学习的开销模型（ML-based Cost Model），用来预测和评估给定配置的表现
TVM 提出了调度模板规范（Schedule Template Specification）API，使开发者可以在调度空间中定义锚点，包括一些额外的特定领域的背景知识。TVM 为各种硬件后端创建了通用主模板（Generic Master Template），用来从张量表达语言中自动提取可能的锚点。
基于机器学习的开销模型
比较而言，黑箱自动调优（Blackbox Auto-tuning）通常用来调优高性能计算的运行库，但是为了取得较好的结果需要大量的实验。另一种方法是预定义开销模型（Predefined Cost Model），理想情况下，它应该能够综合考虑内存访问模式、数据重用、流水线依赖、线程模式等各种因素，然而不幸的是，为当今愈加复杂的硬件架构创建预定义开销模型困难重重。
TVM 采用了数据驱动的方法，ML 模型将低层级的循环程序作为输入，预测它在指定硬件后端上的运行时长。模型使用探索过程中的运行时刻测量数据来训练，不需要用户输入硬件细节信息，模型的准确率会随着试验次数的增加而改进，它是对前两种方法的折衷。
模型的选择上面，质量和速度是关键考量。调度探索器会频繁地对开销模型发起查询请求，使调度优化过程中引入了模型预测和模型更新的时间花费，这类时间花费在真实硬件上应当被严格限制。因此与传统的超参数调优不同，大模型在调度优化里不是一种好的选择。目标函数或者损失函数可以采用预测运行时长与真实运行时长的偏差，由于调度探索器只会从最优秀的候选项里选择，因此事实上不必预测运行时长的绝对大小，TVM 让目标函数支持排名作为替代。
TVM 采用了基于 XGBoost 的梯度树状提升模型（Gradient Tree Boosting Model），从循环程序提取的内存访问计数、每级循环缓存重用率、循环 One-hot 向量表示等特征预测运行时长；另一种基于神经网络的 TreeRNN 模型则是从循环程序的 AST 中提取特征，不需要自行构造特征。两者预测质量接近，但前者训练和推导都更快。
调度探索（Schedule Exploration）
调度探索器最简单的策略就是让每一个配置都跑一边开销模型，然后选择前几个预测表现好的，问题是搜索空间大起来时间花销就不能接受了。TVM 采用的是模拟退火算法，从一个随机配置开始，每次在邻近配置中随机游走，如果开销降低了就接收下一个配置，否则以一定概率拒绝该配置。
分布式设备池（Distributed Device Pool）和远过程调用（Remote Process Call）
分布式设备池使得模型在硬件上的试验次数大大增加，同时多个优化任务之间能够进行细粒度的资源共享。TVM 的分布式设备池基于 RPC 技术的、可定制化的，支持动态装载和运行交叉编译得到的模块。这样一套相同的基础设施可以进行单工作负载的优化和端到端的图推导任务。
参考资料
T. Chen, T. Moreau, Z. Jiang, et al. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning. OSDI’18. https://arxiv.org/abs/1802.04799