TVM（端到端的优化栈）概述

清疚 2021-11-17 12:44 3313阅读 0赞

陈天奇团队宣布推出 TVM，在微博上表示，「我们今天发布了 TVM，和 NNVM 一起组成深度学习到各种硬件的完整优化工具链，支持手机，cuda, opencl, metal, javascript 以及其它各种后端。欢迎对于深度学习，编译原理，高性能计算，硬件加速有兴趣的同学一起加入 dmlc 推动领导开源项目社区 。」

大多数现有系统针对窄范围的服务器级 GPU 进行优化，且需要在包括手机、IOT 设备及专用加速器上部署大量工作。而 TVM 是一种将深度学习工作负载部署到硬件的**端到端** IR（中间表示）堆栈。也就是说，这类解决方案能够把深度学习模型分发到各种硬件设备上、实现端到端的调优。

--------------------

## 引子 ##

众所周知，当前的深度学习框架基于数据流图（如tensorflow等），这种方法过于高层。并不适合硬件界别的设备执行。如下图所示：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70]

**TVM 旨在分离算法描述、调度（Schedule）和硬件接口。**

首先，为了设计出适合深度学习的优化编译器，我们需要解决四大挑战：

*  高级数据流重写：如今不同的硬件设备有不同的存储体系，所以设计出好的策略来将数据运算与优化数据布局进行融合是很重要的。
 *  跨线程极内存共享与协作：传统的无共享嵌套并行模型不再是最佳模型。优化内核需要对加载的共享内存的线程进行协作。
 *  内部函数的张量化：The latest hardware provides new instructions that go beyond vector operations, like the GEMM operator in the TPU or the tensor core in NVIDIA’s Volta. Consequently, the scheduling procedure must break computation into tensor arithmetic intrinsics instead of scalar or vector code.
 *  内存访问延迟、进程（线程）调度的问题：尽量减少延迟是一大挑战。

为了解决上述挑战，TVM设计了两种层级结构：

*  计算优化层（computation graph optimization layer）来解决第一个问题。
 *  张量优化层（tensor optimization layer）使用全新的调度基元来解决剩下的三个问题。

通过 TVM ，我们可以用大多数深度学习框架构建出相应的模型——既能顾及到高层次的语言表述，又能兼顾底层的优化（如GPU等）。  
下面分别介绍两种层级。

## 计算图优化层（Optimizing Computational Graphs） ##

首先我们先看一下Optimizing Computational Graphs的总体实现方法，如下图所示：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 1]  
可以看出传统的计算流图在此处大体被分成两部分：Operator Fusion、Data Layout Transformation。

#### Operator Fusion ####

这个较好理解，其实就是把运算进行了混合，压缩或者简化了原本的计算流程，这种融合过程可用于预先计算图形中可静态确定的部分,从而节省执行成本。TVM 提出了四种融合方式：

*  单射的：injective (one-to-one map)
 *  降维的：reduction
 *  复杂可利用的：omplex-out-fusable (can fuse element-wise map to output)
 *  不可利用的：opaque (cannot be fused)  
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 2]

#### Data Layout Transformation ####

这个功能个人理解应为数据的重构，从而改变存储方式，进而优化访问区域。如下图所示。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 3]  
优化数据布局（Data Layout Transformation）首先指定每个运算符的首选数据布局（数据的shape）,因为这些限制决定了它们在硬件中的实现。然后,如果生产者和使用者的数据布局不匹配,我们会在它们之间执行正确的布局转换。  
**但是，在图层级的优化方式是有局限性的**（因为算法运行的效率很大程度上取决于机器提供的编码或指令集）。事实上，这种数据布局方式会随着数据的类型、硬件加速器种类等而变得不稳定，所以 TVM 又提出来了一种代码自动生成策略。

## 张量优化层（tensor optimization layer） ##

这个层级结构实际上是针对各种硬件后端生成了同一种运算符的微调版本，即对其中的张量运算进行了相关优化。

#### Tensor Expression Language ####

首先，TVM 提出了一种张量的描述语言，用来描述各类算法。如图所示：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 4]  
这种张量描述语言支持各类常见运算。TVM 显式引入交换缩减运算符,以便跨多个线程轻松安排交换性缩减。

#### Schedule Space ####

TVM提出 Schedules 的概念，指的是一种将计算描述（张量运算）降低到后端（底层）优化实现的特定规则。这也是TVM实现的核心。

其理念是对 Schedules 空间和用于遍历此空间的转换进行建模,从而提供生成低级代码的不同方法。TVM的 Schedules 空间如图所示：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 5]

> The dataflow representation provides an expression specifying the computation of each tensor element. The axis relation hyper-graph defines a bijective relation between the original loop axis and the transformed loop axis. The schedule tree represents the loop nest structure and producer store location of the computing structure. The schedule tree representation is derived from Halide, with a special emphasis on store location relation. Together, the dataflow representation, axis relation graph, and schedule tree describe a schedule that can be lowered to the code shown in the rightmost panel.

其实这种实现方法是从 Halide 这门语言为源头开发的，TVM 提出了一种新的语言来处理GPU等硬件加速设备所提出的挑战。  
同时，为了加速我们在Schedule Space中遍历的速度，TVM提出了Schedule基元：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 6]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70]: /images/20211116/77e25b4227ba4b2fa2f9047dc5c88377.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 1]: /images/20211116/5cfd7ef1e41d40a785a99fd79b5da6b0.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 2]: /images/20211116/30aad8d7ebf64224bf440141d6460537.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 3]: /images/20211116/6a8ead55347343adbe008448b3c6f26b.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 4]: /images/20211116/859c2fbd38af439d8ab0664768f9a3dd.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 5]: /images/20211116/f0f28f0b8f534f3ea595a598f6c09b27.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhbmdib3dqMTIz_size_16_color_FFFFFF_t_70 6]: /images/20211116/6e94515cebe745b79514239a2b250029.png