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 Rew