Skip to content

TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

Updated: at 15:06

摘要 :有越来越多的需求将机器学习引入各种硬件设备。当前的框架依赖于特定供应商的操作库,并针对一小部分服务器级别的GPU进行优化。将工作负载部署到新平台——如手机、嵌入式设备和加速器(例如FPGA、ASIC)——需要大量手动工作。我们提出了TVM,这是一个编译器,通过暴露图层次和运算符层次的优化来实现深度学习工作负载在不同硬件后端上具有可移植性的性能。TVM解决了与深度学习相关的优化挑战,例如高级运算符融合、映射到任意硬件原语以及内存延迟隐藏等问题。它还通过采用一种基于学习的成本建模方法自动优化低级程序以适应硬件特性,从而快速探索代码优化。实验结果表明,TVM在各种硬件后端上提供了与最先进、手动调整库相竞争的性能,包括低功耗CPU、移动GPU和服务器级别GPU。我们还展示了TVM针对新型加速器后端(如基于FPGA 的通用深度学习加速器)进行目标设置的能力。该系统已经开源,并在几家大型公司内部使用中。

1.Intro

模型被部署在内存组织、计算单元各不相同的设备上,需要进行优化。传统的TF、Torch之类的方法依赖计算图进行优化,层次太高 ,无法处理特定硬件后端算子级别的优化和转换,并且不好移植。

需要实现一个,从现有框架中获取深度学习程序的高级规范,并且为各种硬件后端生成底层优化代码的编译器,有如下挑战:

TVM为解决这些挑战提出的三个方法:

Contributions :

2.Overview

Untitled.png

从ML frameworks获得模型,转换为计算图,然后进行graph-level的优化、operator-level的优化,然后codeGen输出到对应的后端

3.Optimizing Computational Graphs

TVM优化的方法:operator fusion、constant-folding、static memory planning pass、data layout transformations…

3.1. Operator Fusion

将多个运算结合成一个,从而避免保存中间结果的访存,减少运行时间。算子可以分为

Untitled.png

3.2. Data Layout Transformation

目标是将计算图转换为可以更好利用内部数据结构在目标硬件上执行的计算图。首先需要指定每个算子的“首选data layout”,然后根据一系列约束条件调整data layout。

上面这些high-level的图优化都可以提高效率,但是实际上效果取决于运算库的支持。只有少量框架支持算子融合,而随着算子数量增加,data layout、数据类型 、加速器指令数量都会快速增加,如果要求框架对于所有增加的算子、后端都编写相应的kernel是不现实的。

4.Generating Tensor Operations

4.1. Tensor Expression and Schedule Space

Tensor expression:

Untitled.png

指定了输出tensor的形状和如何计算的表达式。 tensor expression里面避免出现循环之类的结构,是为了将计算和调度解耦。

Scheduling: scheduling的过程是增量地applying basic transformation,中间要保证逻辑等价性。

Untitled.png

4.2. Nested Parallelism with Cooperation

一般的方法是递归地切分任务成子任务从而提高并行性,线程之间不共享内存,称为“shared-nothing nested parallelism”。 另一种办法是让多个thread一同取他们需要的所有数据,并且共享这组数据,提高了数据的利用率与并行性。要实现这种办法TVM引入了内存作用域的概念,便于在调度空间中标记计算阶段。

Untitled.png

4.3. Tensorization

DL负载基本可以分解成GEMM或者conv,所以现代处理器正在逐步添加对于这些操作的支持。利用这些支持为编译调度带来了新的挑战。同时,TVM不能只支持一种固定的(特定硬件的)操作原语,因为硬件正在不断更新、而TVM支持多种硬件。TVM的设计需要一种可拓展的解决方案。 TVM将算子的实际操作与算子声明进行的操作相分离。下面是一个GEMM8*8的tensor定义和实际操作。 ![[Pasted image 20230831154233.png]]

4.4. Explicit Memory Latency Hiding

通过将内存操作和计算相重叠,隐藏内存读写的延迟。在不同硬件上实现内存延迟隐藏的方法是不同的,需要根据硬件后端采取不同的策略。通过”decoupled”的方式实现延迟隐藏:

Untitled.png

Hardware Evaluation of Latency Hiding : 在FPGA上跑ResNet检查每一层的加速情况

Untitled.png

在添加了latency hiding机制之后每层都更靠近roofline了,证明了有效性。

5.Automating Optimization

TVM为每层针对特定输入形状和布局创建专用算子,系统需要选择调度优化以及特定于优化的参数。TVM构建了一个自动调度优化器,主要包括一个提出有前景的新配置方案的explorer和一个预测给定配置性能的模型。

5.1. Schedule Space Specification

给出API,让开发者可以规定哪些量是优化过程中可以变化的。同时还为硬件后端创建了通用的主模版,主模版可以从tensor expression中提取可能的优化量。

5.2. ML-Based Cost Model

讲的内容跟intro里面差不多,额外讲了关于优化过程中要考虑speed和quality的tradeoff。 参数预测用过TreeRNN和XGBoost,两者预测质量相似,但是XGBoost速度将近为两倍并且训练时间少的多。

Untitled.png

5.3. Schedule Exploration

用前面获得的Cost Model,模拟退火去搜一个比较好的config。

5.4. Distributed Device Pool and RPC

分布式计算相关,没怎么仔细讲

6.Evaluation

questions:

Untitled.png

Untitled.png

Untitled.png

Untitled.png


Previous Post
Optimizing Bit-Serial Matrix Multiplication for Reconfigurable Computing
Next Post
Roofline: An Insightful Visual Performance Model for Floating-Point Programs and Multicore Architectures