并行编译技术_「独家」漫谈深度学习的编译/优化器（下）

终于到了一窥究竟的时候了！最后一篇文章，我们重点解释3个最核心的问题：

1. NNVM如何支持不同DL框架的模型？

2. TVM做了哪些优化？

3. TVM如何支持不同的硬件平台？

Let's go!

第一个问题： NNVM如何支持不同DL框架的模型？

还记得LLVM中前端的定义吗？！所有的DL框架对NNVM而言，都是前端。So，我们可以很轻易地找到NNVM对所有DL框架支持的代码，look，如下图！

可以看到NNVM对caffe2、Darkne、Tensorflow、Mxnet等DL框架的支持，我们随便挑Mxnet作为样例，主要涉及的源码有From-mxnet.py(使用样例)、mxnet.py等，大致步骤如下：

首先通过关键函数from_mxnet(symbol, arg_params=None, aux_params=None)的作用是将MxNet模型转化为NNVM可以编译的格式，第一个输入参数是MxNet网络模型定义的Symbol，后2个输入参数是MxNet模型的参数【这3个输入参数可以由mxnet框架API获取，样例代码：mx_sym, args, auxs = mx.model.load_checkpoint('mobilenet', 0)】。

然后使用NNVM进行编译，函数如下：

graph, lib, params = nnvm.compiler.build(net, target=target, target_host=target_host, shape={'data': input_shape}, params=params, dtype=dtype)

第1个输入参数是mxnet模型转后后的symbol(上面提及的mx_sym)；

第2个输入参数是TVM支持的硬件平台类型，如下图所示：

对应源码文件Target.py，可支持的硬件平台有ARM CPU、NVIDIA GPU、AMD GPU等；

第3、4、5输入参数是模型参数，不必多说。

最后将模型转换为可在TVM执行的库文件，对应源码Module.py中的export_libray函数，使用样例：lib.export_library("mobilenet_deploy.so")。

*完整示例：

第二个问题：TVM做了哪些优化？

TVM是一个端到端优化堆栈，主要有如下优化：

内容源于大神的论文《TVM: End-to-End Optimization Stack for Deep Learning》https://arxiv.org/abs/1802.04799v1

2.1 操作符融合：

对于GPU和特定加速器而言，将多次操作融合在一起的优化方法能较为明显地降低执行时间。操作符融合的想法是来源于单个Kernel函数会节省将中间结果写回全局内存的时间消耗。从具体分析来看，我们总结了四种类别的图操作符：

• injective(one-to-one map)：单射，如add/sqrt/sub
• reduction：约简，如sum/max/min
• complex-out-fusable(can fuse element-wise map to output)，如conv2d
• opaque(cannot be fused)

2.2 计算预优化

一些常数折叠Pass能够在计算图预计算阶段被静态地执行，节约运行时的代价。

2.3 数据共享及数据布局转换

Tensor操作是计算图的基本操作符，Tensor中涉及到的运算会根据不同的操作符拥有不同的数据布局需求。例如，一个深度学习加速器可能会使用4x4张量操作，所以需要数据切割成4x4的块来存储以优化局部访存效率。下图展示了一个矩阵如何布局，这种布局能够适应计算2x2张量操作：

2.4 Tensor操作优化

TVM的Tensor表达式语言借鉴了一些已有的编程语言，像Halide，Darkroom和TACO。每个操作被描述为一个数学表达式，例如：

优化1：TVM通过借鉴Halide，采用了解耦计算描述和调度器优化两个过程。调度器会根据硬件后端采用特定的规则将计算描述向下转换成已优化的硬件实现。

优化2：协作式嵌套并行化

针对深度学习工作负载而言，并行编程的关键是改善计算密集型Kernel的计算效率。现代GPU提供了大规模的并行化，需要将并行编译模型加入到调度器里面。大部分现有的解决方案都是采用了嵌套并行编译的方法，是一种fork-join的并行模式。TVM能使用一个并行调度器基元来处理一个数据并行任务。每个并行任务都能够被递归地细分到一个子任务以实现多级线程。这种模型为 shared-nothing nested parallelism，即一个工作线程不能观察到相同计算阶段内的相邻线程的数据。相邻线程之前的交互只能发生在join阶段，join发生在子任务结束并且下个阶段所需的数据已经准备好时。这种编程模型禁止线程在相同计算阶段内进行协作。

2.5 编译器延迟隐藏

延迟隐藏指的是针对一个在计算过程中重叠内存操作的过程最大化访存和计算利用率。针对不同的硬件它需要采用不同的策略。对于CPU而言访存延迟隐藏通过同步多线程或者硬件预取技术达到。对于GPU而言，则通过成千上万的线程组快速切换达到访存延迟隐藏。对于特定深度学习加速器，经常会倾向于使用精简控制流和编译器堆栈。下图展示了编译器如何在一个流水线式深度学习加速器中显式地处理数据依赖：

第三个问题： TVM如何支持不同的硬件平台？

在回答这个问题前，介绍另外一个东东：Apache Calcite，之所以介绍这个，是因为Calcite对不同平台的SQL支持方式，与TVM支持不同硬件平台是一个套路。

看Calcite前2个步骤是SQL语句的解析、逻辑计划的优化，要注意的是最后一步物理计划的生成，它是将标准的SQL Operation转换为不同平台的执行实现。

在深度学习中，也有非常多的标准OP，比如Relu、SoftMax、CNN、Pool之类，TVM中都统一到Tensor中，TVM中将优化后的计算图算子转换为具体平台物理执行代码的过程成为：张量化，生成硬件接口。在论文中，TVM对不同硬件平台的支持是这样表述的：

小结一下就是：从调度器中分离了硬件接口，引进了一套Tensor内联声明机制。TVM使用Tensor表达式语言来声明每个新的硬件内联的行为，引进了一个调度器基元Tensorization来作为基元的计算单元。如下图所示：

Tensor表达式语言能够同时描述用户准备的计算描述信息和硬件暴露接口的抽象信息。Tensorization把调度器从硬件基元中解耦出来，这样能够使TVM更容易扩展新的硬件架构。Tensorization调度器生成的代码经过实践具有很高的计算性能：将复杂的操作分解成一系列重复的微型Kernel调用。因此开发者能够使用自定义Tensorization基元的方式来提升性能，这种方法在一些硬件平台上是非常有益的。例如，在AMD Vega GPU上，通过Tensorizing半精度GEMM到一个手动制作的4x4微型Kernel获得1.5倍的性能加速效果。

当然在TVM中，这些代码的生成过程正常情况下是自动的，论文中表述如下：

代码生成：TVM通过迭代遍历调度树的方案生成Lowered代码，low-level代码是通过类C的循环程序具象化而来。在这个过程中TVM使用了一个Halide循环程序数据结构的变体，也使用了Halide常用的lowering primitives，像是storage flattening，循环展开，对于GPU和特定深度学习加速器而言，则使用了同步点检测、虚拟线程注入、模块化生成机制，最终，循环程序被翻译成LLVM/CUDA/Metal/OpenCL源代码。

运行时支持：对于GPU程序而言，TVM会独立构建host端和device端的模块并且提供一个运行时模块系统来启动Kernel。对于FPGA程序，TVM也构建了一套驱动程序，这套驱动程序使用C语言API，可以实现构造指并推送到目标加速器上执行。

到了这里，宝宝心里是不是可以松一口气了呢？！

嘿嘿，大神们其实很早就发现：深度学习硬件真正的目标，是希望大家都可以把自己的深度学习的模型跑在加速器上。而加速器带来的特殊定制，导致对于每个加速器，我们都需要重新设计一套driver，以及上面的软件栈---包含编译，优化和部署模块。现在的行业现状，不论是高校还是公司，大家都面临的同样的一个问题 -- 深度学习加速器从来不只是硬件的问题，而是涉及到一个从硬件，到驱动，再到编译优化部署到模型本身的全栈问题。

So，出现了一个专门的子项目：VTA！VTA是什么？VTA的全名叫做Versatile Tensor Accelerator，直译过来为灵活的张量加速器。首先，VTA是一个完全开源的深度学习加速器。但是VTA不光包含了加速器设计本身，完整的驱动，tvm编译的整合和直接从tvm前端python编译部署深度学习模型的完整开源工具链。TVM发布的VTA包含了模拟器和FPGA部署模块。

关于VTA就不在DL 编译/优化器文章中赘述了，有空的时候再单开文章来絮叨絮叨，小伙伴们知道它是“张量化：硬件接口生成”的扩展项目就好。

好了，关于《漫谈深度学习的编译/优化器》的文章终于可以告一段落了，更多的细节可以伙伴们可以关注留言或私信交流。