深度课堂：全角度解读神经网络编译器-技术圈

一、神经网络编译器概览

近年来，以机器学习、深度学习为核心的AI技术得到迅猛发展，深度神经网络在各行各业得到广泛应用：

1. CV（计算机视觉）：目标检测，场景识别，图像分割等。

2. 智慧语音：语音识别，声纹识别等。

3. NLP（自然语言处理）：自动搜索引擎，对话服务机器人，文本分类，智能翻译等。

4. 科学研究：应用于物理、生物、医学等多研究领域。高能粒子对撞分类，宇宙天体图数据分析，星系形状建模，生物结构的蛋白质折叠预测，精准医疗与疾病预测。

这些应用催生更多的新模型出现：CNN, RNN, LSTM, GAN, GNN，也催生着如Tensorflow, Pytorch, Mxnet, Caffe等深度学习框架出现。目前训练框架开始收敛，逐步形成了PyTorch引领学术界，TensorFlow主导工业界的一个双雄局面。

但是，深度学习算法要实现落地应用，必须被部署到硬件上，例如Google的TPU、华为麒麟NPU，各种AI专用芯片，FPGA等。

这些各训练框架训练出来的模型要如何部署到不同的终端硬件呢，这就需要深度学习神经网络编译器来解决。

在介绍神经网络编译器之前，我们先简单回顾下传统编译器。

传统编译器：

熟悉C语言编译流程的应该了解，编译器其实就是一堆工具链的组合，输入是高级编程语言，输出是目标硬件的机器码。一般编译器的流程可以分为三个部分：解析高级语言，编译优化，生成目标代码。

LLVM（low level virtual machine）是一个开源的编译器基础设施项目，通过模块分为前端，中端（优化）和后端三部分。LLVM提出了LLVM的中间表示LLVM IR。每当出现新的编程语言，只需要开发相应的前端，将编程语言转换成LLVM的中间表示；类似地，出现新的硬件架构，只需要开发相应的后端，对接上LLVM的中间表示。模块化的划分，避免了因编程语言和CPU架构的翻新而引发的编译器适配性问题，大大简化了编译器的开发工作。

神经网络编译器：

前面我们提到，神经网络编译器的出现是为了把深度学习神经网络模型部署到终端设备上，简单来说，就是让神经网络模型在设备上能跑起来，而且跑得快，能满足应用需求。如果是针对这个目标，传统编译器也可以满足模型在硬件上的部署需求，为什么要提出神经网络编译器呢？神经网络编译器和传统编译器的区别和联系是什么呢？

编译器的输入/输出：

传统编译器的输入是高级编程语言，输出是目标硬件的机器码。神经网络编译器的输入是深度学习训练框架训练出来的模型定义文件，输出是能够在不同硬件高效执行的代码。

编译器的IR：

传统编译器的IR是前端对接不同的编程语言，后端对接不同硬件，属于low level IR。神经网络编译器需要专门对神经网络进行优化，前端适配各种训练框架训练出来的神经网络模型，后端适配各种硬件，因此需要High level IR(计算图层级)以及算子层级/张量计算层级(operator/tensor)的IR，算子层级需要实现算子在目标硬件上高效执行。下图展示是神经网络编译器的层级结构：

1. 前端对接各个深度学习训练框架训练出来的算法模型（Tensorflow, Caffe, Pytorch, Mxnet等）。

2. 图层级（High-level IR）：神经网络的结构可以表示成计算图，图层级的操作则是对计算图进行一些和具体硬件和框架无关的操作，比如算子融合，内存分配优化，数据类型和数据维度的推导等。我们可通过算子融合的方式，避免中间数据频繁的在寄存器和内存直接来回读写，从而提升整体推理性能。

3. 算子层级（operator level/kernel level）算子层级主要是张量计算。为了实现这些计算在硬件上高效实现，发挥芯片的性能，通常硬件芯片配有专门优化的算子计算库，如Intel的MKL, Nvidia的CuDNN, TensorRT。这个层级需要支持每个硬件后端的每个算子实现。

4. 各硬件后端：GPU, ARM CPU, X86 CPU, NPU等。

总的来说，传统编译器是面向编程语言写的任意程序在目标硬件的优化，神经网络编译器是专门面向神经网络在目标硬件的高速部署。目前神经网络编译器处在一个发展中的阶段，早期图层级和算子层级都是手工去实现，手工实现算子融合，专家手工写神经网络算子在目标硬件的优化计算库等，后期慢慢朝着自动优化方向发展。下文将介绍几个相关的开源项目：TVM，Halide，AutoKernel。

二、TVM的前世今生

在编译器快速发展的浪潮中，较为突出的便是TVM（Tensor Virtual Machine)。 TVM最早提出是2017年，是深度学习系统的编译器堆栈。

第一代TVM的设计借鉴了借鉴传统编译器框架LLVM的设计思路，设计抽象出中间表示层，不同的模型只需要开发相应的前端接口，不同的后端只需要开发相应的后端接口。

TVM全称为Tensor Virtual Machine，属于算子层级，主要用于张量计算，提供独立于硬件的底层计算中间表示，采用各种方式（循环分块，缓存优化等）对相应的计算进行优化。第一代的图层级表示叫NNVM（Neural Network Virtual Machine）。NNVM的设计目标是：将来自深度学习框架的计算图转换为统一的计算图中间表示（IR)，对之进行优化。

第一代的静态图存在一定的缺陷：

1. 不能较好支持控制流，如分支跳转，循环等。

2. 不能支持计算图输入形状，取决于输入tensor大小的模型，比如word2vec等。

虽然Tensorflow有如tf.cond.Tf.while_loop的控制接口来在某种程度上解决第一个问题，tf.fold来解决第二个问题，但是这种方式对刚刚接触深度学习框架的小白来说不是特别友好。

后面出现的动态图摒弃了传统的计算图先定义，后执行的方式，采用了计算图在运行时定义的模式，这种计算图就称为动态图。

第二代TVM 的图计算层变为Relay VM，Relay和第一代的图计算表示NNVM的最主要区别是Relay IR除了支持dataflow（静态图），能够更好地解决control flow（动态图）。它不仅是一种计算图的中间表示，也支持自动微分。

总结一下，目前TVM的架构是：

1. 最高层级支持主流的深度学习前端框架，包括TensorFlow,MXNet,Pytorch等。

2. Relay IR支持可微分，该层级进行图融合，数据重排等图优化操作。

3. 基于tensor张量化计算图，并根据后端进行硬件原语级优化，autoTVM根据优化目标探索搜索空间，找到最优解。

4. 后端支持ARM、CUDA/Metal/OpenCL、加速器VTA（Versatile Tensor Accelerator）。

三、Halide

Halide于2012年提出，主要用于自动优化。其嵌入到C++中，是MIT研究人员专门为图像处理设计的一种程序语言。Halide语言易于编写，语法简单，数据结构清晰，能过自动对代码进行优化，使得程序获得比较好的执行效率。

它设计的核心思想是把算法和调度分离。这样做的好处是，在给定算法的情况下只需要去调整它的Schedule 调度选择，不用重写算法实现不同的Schedule。当调整Schedule、探索设计空间时也不会担心因为重写算法而导致计算的正确性会发生变化。

Algorithm部分主要是算法描述和计算的数学表达式。Schedule部分则是告诉机器什么时候分配内存，如何计算（分块计算还是顺序计算）——目前已经提供了一些调度策略。

不同调度策略考虑重复冗余计算和局部性（locality)的权衡。

四、AutoKernel

深度学习模型能否成功在终端落地应用，满足产品需求，一个关键的指标就是神经网络模型的推理性能。

目前的高性能算子计算库主要是由高性能计算优化工程师进行手工开发。然而新的算法/硬件的不断涌现，导致了算子层级的优化开发工作量巨大。同时优化代码的工作并不是一件简单的事，它要求工程师既要精通计算机体系架构，又要熟悉算子的计算流程。

人才少，需求多，技术门槛高，因此我们认为算子优化自动化是未来的大趋势。而提出AutoKernel的初衷便是希望能把这个过程自动化，从小处入手，在算子层级的优化，实现优化代码的自动生成。

AutoKernel的输入是算子的计算描述（如Conv、Poll、Fc），输出是经过优化的加速源码。

这一工具的开发旨在降低优化工作的门槛，不需要有底层汇编的知识门槛，不用手写优化汇编。可通过直接调用开发的工具包便可生成汇编代码。同时还提供了包含CPU、GPU的docker环境，无需部署开发环境，只需使用docker便可。还可通过提供的插件——plugin，可以把自动生成的算子一键集成到推理框架中——Tengine。