深度学习AI编译器-TVM简介

1.为什么需要深度学习编译器

深度学习编译器主要为解决不同框架下训练的模型 部署到指定的某些设备上时所遇到的一系列复杂的问题，即将各种深度学习训练框架的模型 部署到各种硬件所面临的问题；
首先深度学习领域，从训练框架看，当前可选的框架有pytorch、TensorFlow、Mxnet、paddle，oneflow、caffe/caffe2、mindspore等，具体选择哪个，不尽相同，但如果项目要部署落地，则面临很多问题，即从推理框架角度来看，无论我们选择何种训练框架训练模型，我们最终都是要将训练好的模型部署到实际场景的，在模型部署的时候我们会发现我们要部署的设备可能是五花八门的，例如Intel CPU/Nvidia GPU/Intel GPU/Arm CPU/Arm GPU/FPGA/NPU(华为海思)/BPU(地平线)/MLU(寒武纪)，如果我们要手写一个用于推理的框架在所有可能部署的设备上都达到良好的性能并且易于使用是一件非常困难的事。

为了解决上面的问题，科学家为编译器抽象出了编译器前端，编译器中端，编译器后端等概念，并引入IR (Intermediate Representation)的概念。解释如下：

编译器前端：接收C/C++/Java等不同语言，进行代码生成，吐出IR
编译器中端：接收IR，进行不同编译器后端可以共享的优化，如常量替换，死代码消除，循环优化等，吐出优化后的IR
编译器后端：接收优化后的IR，进行不同硬件的平台相关优化与硬件指令生成，吐出目标文件

因此我们可以将各个深度学习框架训练出来的模型看做各种编程语言，传入深度学习编译器，之后吐出IR，由于深度学习的IR其实就是计算图，所以可以叫做Graph IR，针对这些Graph IR可以做一些计算图优化在吐出IR分发给各种硬件使用，这样就解决了上述很多繁琐的问题，如下图所示：

深度学习编译器VS传统编译器

图中与传统编译系统进行了对比，发现存在相似之处：如都有中间表示、前端和后端分别硬件无关优化和硬件相关优化以及最后的代码生成模块；
同时也有一定的区别：深度学习编译系统是针对神经网络这一特定领域的编译系统，采用了以Python为主的动态解释器语言的前端、多层IR（如图层/算子层/codegen）设计、以及面向神经网络的特定优化（如自动微分、量化/混合精度、大规模并行、张量运算、循环优化等）。

深度学习编译器的架构

• 前端：将现有深度学习框架中的深度学习模型作为输入，然后将模型转换为计算图表示。前端需要实现各种框架不同格式的转换，并进行了结合了通用编译系统的优化技术和深度学习特定的优化技术的硬件无关的计算图优化，优化图逻辑，减少了冗余，提高了图IR的效率。计算图优化包括代数简化、算子融合、算子下沉、CSE公共子表达式消除、DCE死代码消除、静态内存规划和布局转换等。在前端之后，将生成优化的计算图并将其传递到后端。
• 中间表示IR：分布在前端和后端，是介于源语言和目标语言之间的程序表示，能够极大程度地提高编译流程的可拓展性，同时也能降低优化流程对前端和后端的破坏
• 后端：将高级IR转换为低级IR，并执行特定于硬件的计算图优化、算子选择以及内存分配这三个任务。计算图优化是在不影响模型的数值特性的基础上，通过图变换达到减少资源开销、适配硬件的执行能力、提升执行性能的目的。例如与硬件无关的算子内存 IO 优化和为了适配特定硬件指令限制而做的子图变换。数据存在多种存储格式和计算精度，不同的存储格式和计算精度在不同场景下对算子计算性能有较大的影响，算子选择是为 IR 图中的每个计算节点选择一个最适合在设备上执行的算子。经过计算图优化和算子选择以及常用的特定于硬件的优化包括硬件内部映射、内存分配和获取、内存延迟隐藏、并行化以及面向循环的优化后。使用JIT或AOT编译优化的低级IR，以生成不同硬件目标的代码，运行时要实现算子选择和内存分配等技术。

深度学习编译系统整体架构图

此外，对应深度学习工作流，完整的深度学习编译框架还包括：编程接口、硬件加速器、数据处理、模型部署、分布式训练等模块。深度学习编译系统以加速器为核心，硬件加速器模块提供丰富的编程接口；数据处理模块负责数据的读取、存储以及预处理；分布式训练模块为应对单个机器上内存及算力资源不足，并行加速模型训练的过程；模型部署模块负责进行模型压缩、针对推理硬件平台模型算子优化、进行模型混淆设计保证模型安全等。

中间表示

网络模型传递给编译系统后会被转成IR，深度学习编译系统中IR的设计需要考虑：正确处理张量数据类型；自动微分实现的简洁性、性能以及微分的扩展能力；支持静态图和动态图，可以针对任务需求灵活选择合适的模型；支持高阶函数和闭包，以抽象通用问题、减少重复代码；编译优化；JIT及时编译能力等因素。

常采用多级IR的设计方法将前后解耦，常见的一种设计是分为High-Level IR，Low-Level IR以及代码生成时用到的LLVM IR，如图所示：

• High-Level IR ：也称为图形IR，表示计算和控制流，与硬件无关。高级IR的目标是在操作符和数据之间建立控制流和依赖关系，并为图形级优化提供接口，进行常数折叠、代数化简、公共子表达式、Layout转换、算子融合、优化图逻辑等优化。
• Low-Level IR：用于针对不同硬件目标的特定于硬件的优化和代码生成。因此，低级别IR应该足够细粒度以反映硬件特性并表示特定于硬件的优化，如算子选择、循环变换、循环切分，与硬件intrinsic映射、内存分配等后端pass优化。 深度学习编译系统中的IR设计趋势 深度学习的计算任务的特点是比较结构化，即控制依赖相对较少；且数据以张量为主，访问比较有规律。从图中可以看到深度学习编译系统的IR设计的其中一种变化趋势是变得越来越“厚”，即层数越来越多。这些不同层次间的IR之间一般通过progressive lowering可以从高层到低层进行转化。就像接力一样，每一层抽象上处理本层适合干的事，然后将变换后的IR往下层丢。另一方面，如MLIR也变得越来越灵活和开放

前端

深度学习编译系统的起点是将网络模型转换为图IR，首先介绍前端中的计算图。

计算图： 用来表示深度学习网络模型在训练与推理过程中计算逻辑与状态的工具，由张量Tensor、算子Operator以及有向线段表示的张量状态、依赖关系构成。用于对输入数据、算子和算子执行顺序进行统一表达；定义中间状态和模型状态，从而帮助框架更好地管理内存；自动化分析模型定义、在自动微分过程中计算梯度；分析算子执行关系，发现将算子进行异步执行的机会，从而优化程序执行

目前主流的深度学习框架的执行模式有两种，分别为静态图模式和动态图模式。

静态图模式下，程序在编译执行时先生成神经网络的图结构，然后再执行图中涉及的计算操作。因此，在静态图模式下，编译器利用图优化等技术对执行图进行更大程度的优化，从而获得更好的执行性能，有助于规模部署和跨平台运行，包括以下两种实现方式：

Tracing模式：框架把Python假执行一遍，记录算子执行序列作为正向图，并以此进行自动微分生成反向图，然后进行正反向图的编译优化
AST转换：框架获取Python代码的AST，通过编译技术转换成正向图，基于正向图生成反向图，同时进行编译优化

动态图模式下，程序采用解析式的执行方式，编译与执行同时发生，采用前端语言自身的解释器对代码进行解析，利用计算框架本身的算子分发功能，按照计算依赖关系进行调度执行，动态生成临时的图拓扑结构。计算框架依照算子调度顺序记录参与反向计算的算子与张量。当前向计算执行完毕，计算框架根据动态生成的前向计算图结构拓扑关系，利用记录的反向计算算子和张量动态生成反向计算图，最终完成神经网络模型的梯度计算与参数更新，更加节省内存占用，同时方便用户编写和调试神经网络模型。

目前，深度学习编译系统中均动静态图两种模式，并逐步呈现由动静分离到动静融合再到动静统一的趋势，动静融合中的动静态图转换的技术常见于模型部署阶段，当动态图模式训练完成模型参数后，可以将整体网络结构转换为静态图格式，将神经网络模型和参数文件进行序列化保存，与前端代码完全解耦，扩大模型部署的硬件支持范围；动静统一具体是指静态图可以支持尽量多的动态图语法，同时使得静态图提供接近动态图的语法使用体验。

生成计算图IR之后，为了有效减少程序在运行时可能出现的错误，编译系统的前端引入了类型系统（Type System）和静态分析（Static Analysis）系统。

类型系统是类型的集合以及使用类型来规定程序行为的规则，用于定义不同的类型、指定类型的操作和类型之间的相互作用，可以防止程序在运行时发生类型错误，

静态分析是指在不实际运行程序的情况下，通过词法分析、语法分析、控制流、数据流分析等技术对代码进行分析验证的技术，能够为编译优化提供线索和信息，有效减少代码中存在的结构性错误、安全漏洞等问题。在编译系统前端的编译过程中，静态分析可能会被执行多次，有些框架还会通过静态分析的结果判断是否终止编译优化。

此外，还进行与传统编译系统通用的常见硬件无关优化：代数简化、CSE公共子表达式消除、DCE死代码消除、静态内存规划和布局转换等。

微分求解通常有手动求解、数值微分、符号微分以及自动微分几种方法。

手动求解法：求解出梯度公式，然后编写代码，代入实际数值得出真实的梯度。但每次修改算法模型都需要修改对应的梯度求解算法。
数值微分法：根据导数的原始定义 ，当h取很小的值时，用户就能根据给出的目标函数和要求解的变量就可以方便地自动给出相应的梯度。但计算量大。
符号微分法：使用代数软件实现一些公式，然后对用户提供的具有闭包形式的数学表达式进行自动微分求解。但会存在表达式膨胀问题，即可能会使得问题符号微分求解的表达式急剧膨胀，导致最终求解速度变慢。
自动微分法是一种介于符号微分和数值微分的方法，将符号微分法应用于最基本的算子（常数、幂函数、指数函数、对数函数、三角函数等），然后代入数值保存中间结果，最后再应用于整个函数。
深度学习编译系统中为何需要自动微分？

深度学习框架训练模型的基本流程是：依据模型搭建计算图；依据输入计算损失函数；计算损失函数对模型参数的导数；根据得出的导数，利用梯度下降法等方法来反向更新模型参数，让损失函数最小化。搭建计算图（依赖关系）和计算损失函数的过程，称为"正向传播”，这个是依据用户模型完成，本质上是用户自己处理。而依据损失函数求导的过程，称为"反向传播”，这个对于用户来说太繁重，所以各种深度学习框架都提供了自动求导功能，进行反向传播时的自动梯度计算和更新模型的参数。因此， 自动微分在深度学习中处于非常重要的地位，是整个训练算法的核心组件之一。同时，自动微分通常在编译系统前端优化中实现，通过对中间表示的符号解析来生成带有梯度函数的中间表示，可以灵活结合编程语言的循环结构，其过程实质上是一种图计算，可以对其做很多优化。

自动微分根据链式法则的不同组合顺序可分为前向模式和后向模式。

前向模式：从输入方向开始计算梯度值，
后向模式：从输出方向开始计算梯度值 ，
当函数的输出个数远远大于输入个数时，前向模式效率更高，反之后向模式效率更高。由于计算一个标量值的输出关于大量参数输入的梯度是深度学习中最为常见的，使得反向模式的自动微分成为反向传播算法中使用的核心技术之一。

后端

计算图优化

在AI框架设计方面，目前业界主流采用图层和算子层分层的实现方法。在前端，图层负责对计算图进行融合或重组，图层通常采用基于Tensor的High-Level IR的处理和优化；后端中仍需要对计算图针对硬件作进一步优化，之后算子层负责将融合或重组后的算子编译为高性能的可执行算子，算子层则采用基于计算指令的Low-Level IR进行分析和优化。

在计算图中尝试匹配特定的子图结构，找到目标子图结构后通过等价替换方式将其替换成对硬件友好的子图结构，进行通用硬件优化，如主流AI计算框架对用户提供的算子通常是从用户可理解、易使用角度进行定义。每个算子承载的计算量不等，计算复杂度也各不相同。但从硬件执行角度看，基于用户角度的算子划分，并不高效，也无法充分发挥硬件资源计算能力，因此可通过算子融合（将计算密集型算子与访存密集型算子融合）以及自动算子生成技术进行优化。

此外还需要进行特定硬件优化，包括由于硬件指令限制而做的优化（IR中计算节点没有直接对应的硬件算子只能通过子图的变换进行与硬件已存在的相对应），特定于硬件存储格式的优化（算子执行完成后的格式转换操作）。

华为编译器在后端实现的图算融合通过对计算图结构的拆解和聚合可以实现跨算子边界的联合优化，在算子重建过程中通过通用的计算规则以必要的访存作为代价，生成对硬件友好、执行更为高效的新算子。图算融合在图层主要思路是把复合算子打开，然后进行跨边界聚合和优化（如代数化简、公共子表达式提取（CSE）等），最后进行Kernel算子拆分。优化后的计算图会以一个个子图的方式传给MindSpore AKG继续进行后端优化、生成目标代码。

算子选择

得到优化的IR图后，由于对于每个节点可以选择不同的输入输出格式和数据类型等，需要为其选出最为合适的算子，生成完整的算子序列。算子信息：主要包括支持设备类型、数据排布格式、计算精度等。基于数据排布格式和数据精度，不同硬件会有不同的算子支持，此时硬件上支持的所有算子的集合组成算子信息库，算子选择过程就是从算子信息库中选择合适的算子。首先，选择算子执行的硬件设备，然后后端会根据IR图中推导出的数据类型和内存排布格式选择对应的算子。一个好的算子选择算法应该尽可能的保持数据类型与用户设置的数据类型相一致，且尽可能少的出现数据格式转换。

TVM编译系统中Ansor就高效地实现了算子选择功能，Ansor 有三个主要部分：

程序采样器（program sampler）：构建一个大的搜索空间，并从中采样不同 的程序；
性能调整器（performance tuner）：对采样程序的性能进行微调；
任务调度器（task scheduler）：为优化 DNN 的多个子图分配时间资源。

Ansor实现算子选择流程
针对算子选择，Ansor主要解决了以下两个问题：

如何自动化的构造一个更大的搜索空间？Ansor 使用了一个层次化的搜索空间。
如何更有效的进行搜索？Ansor 在搜索过程中增加了采样，先对完整的程序进行采样然 后再调整，提高了搜索效率。

内存分配

遍历算子序列，根据IR图中算子的输入输出的形状、数据类型、存储格式等信息为每个算子计算分配相应的输入输出内存，然后将算子加载至设备上执行计算。内存复用是一个重要的内存分配优化手段，可以让设备上容纳更大的网络模型。常见的内存分配优化手段包括内存融合和In-Place算子，分别能够将通信算子的内存进行融合以分配连续的内存，可以提高通信的效率以及合理分配In-Place算子的内存，尽量减少为每个算子的输入输出分配不同的内存，以节省内存使用并且提高计算效率。

计算调度与执行

模型训练就是计算图调度图中算子的执行过程。宏观来看训练任务是由设定好的训练迭代次数来循环执行计算图，此时我们需要优化迭代训练计算图过程中数据流载入和模型训练（推理）等多个任务之间的调度执行。微观上单次迭代需要考虑计算图内部的调度执行问题，根据计算图、计算依赖关系、计算控制分析算子的任务调度队列。

其他部件

• 模型训练 分布式训练是为应对单个机器上内存及算力资源不足，设计分布式训练系统将一个机器学习模型任务拆分成多个子任务，并将子任务分发给多个计算节点（切分—并行—合并模式），解决资源瓶颈方法包括：数据并行、模型并行、混合并行以及流水线并行。

数据并行

数据并行： 单程序多数据模式，解决算力不足。其中不同设备上对应的局部梯度会进行聚合，从而计算平均梯度。聚合过程往往由集合通讯库（Collective Communication）的 Allreduce 操作来完成。

模型并行

模型并行： 多程序单数据模式，解决内存不足问题，可分为算子内并行和算子间并行两种方式。

混合并行

混合并行： 多程序多数据模式

流水线并行：

流水线并行：解决计算图下游设备长时间处于空闲状态的模型并行空洞问题，将一个数据小批量划分为多个微批量，每个微批量对应的梯度将会缓存，等到全部微批量完成，缓存的梯度会被加和，算出平均梯度，更新模型参数。关键因素是流水线泡沫（设备完成前向传播后等到全部反向传播开始期间设备处于空闲状态）。

• 数据处理 数据处理包括数据加载、数据混洗、数组组装、数据发送等过程。

数据加载（Load，负责从存储设备中加载读取数据集，提供统一数据存储格式进行读取）
数据混洗（Shuffle，将输入数据按用户指定顺序随机打乱以提升模型鲁棒性）
数据变换（Map，内置数据预处理算子进行数据变换处理，如图像缩放翻转、音频随机加噪变调、文本停词去除随机遮盖）
数据组装（组装一个batch的数据进行训练）
数据发送（将处理后数据发送至加速器设备中进行后续模型计算和更新）
数据处理模块的设计需要满足易用性、高效性以及保序性等特性。

易用性是指提供良好的编程抽象和接口使得用户方便构建一个数据处理流水、支持用户灵活地在数据流水中注册自定义算子；
高效性是指支持随机读取且具备高读取吞吐率的文件格式、合理的并行架构执行数据流水线：数据读取、数据预处理计算、芯片上模型计算三个步骤异步并行执行；
保序性：使数据最终送入模型的顺序由数据混洗组件的数据输出顺序唯一确定。

• 模型部署 模型部署是指将训练好的模型部署到运行环境中进行推理的过程，需要提供训练模型到推理模型的转换、部署阶段的性能优化、模型压缩、模型推理优化以及安全保护等功能。实现训练模型到推理模型的转换需要提供了不同训练框架的转换器，转换为统一的数据结构；完成模型转换后，会将一些不依赖于输入的工作提前完成，且有些优化工作只有在部署阶段才能进行或者彻底进行，包括算子融合、算子替换、算子重排等；模型压缩算法包括减小模型精度的量化、剪枝神经网络的稀疏、有选择地从教神经网络中任意的一层来传递有用的信息给另一更小的神经网络的知识蒸馏；之后在计算硬件上进行前处理、执行推理与后处理；此外还需要对计算图网络结构加扰、节点匿名化、参数权重加扰、算子变换等进行模型混淆以实现模型安全保护。

2. TVM

基于上面深度学习编译器的思想，陈天奇领衔的TVM横空出世。TVM就是一个基于编译优化的深度学习推理框架（暂且说是推理吧，训练功能似乎也开始探索和接入了），我们来看一下TVM的架构图，图片来自：https://tvm.apache.org/2017/10/06/nnvm-compiler-announcement
从这个图中我们可以看到，TVM架构的核心部分就是NNVM编译器（注意一下最新的TVM已经将NNVM升级为了Realy，所以后面提到的Relay也可以看作是NNVM）。NNVM编译器支持直接接收深度学习框架的模型，如TensorFlow/Pytorch/Caffe/MxNet等，同时也支持一些模型的中间格式如ONNX、CoreML。这些模型被NNVM直接编译成Graph IR，然后这些Graph IR被再次优化，吐出优化后的Graph IR，最后对于不同的后端这些Graph IR都会被编译为特定后端可以识别的机器码完成模型推理。比如对于CPU，NNVM就吐出LLVM可以识别的IR，再通过LLVM编译器编译为机器码到CPU上执行。

TVM整体结构如下图

Frontend：这个就是将来自不同深度学习框架中的神经网络转化成TVM自己的IR表示。神经网络模型的输入是protoBuf文件，比如在tensorflow中就是pbtxt文件，这个文件中是protobuf形式的网络结构，包括各个节点的特性，参数，算符等。

IRModule：TVM定义的IR结构有两类，一个是relay中的IR，一个是TIR。两者的区别，relay中的IR是一个高层次的抽象表示，而TIR是更接近硬件的表示结构。Relay中IR通过relay::function来描述，function描述了整个图结构，是图结构的另外一种描述，function有输入输出变量，内部是计算序列。很像编程中定义的函数。TIR中用PrimFunc描述，它更加底层一些，包含了tensor计算算符，load/store等指令。用户如果有自定义的指令，就可以在TIR中定义这些指令。

Passes：这些pass是完成对网络图结构的一些优化和转化，比如常数折叠，算符融合，去除死代码等等，都使用pass来构造。TVM通过这些pass来遍历，然后对节点进行修改。对于内存分配，针对硬件特点的图优化策略都可以用pass来进行构造。

Transform：transform包括两方面：优化和lowering，这些都是在pass架构中构造的。Lower实现高层次IR到TIR的转化，或者完成到硬件实现的转化。比如将多维tensor转化成一维向量。当然很多lowering操作都是硬件相关的，需要开发人员自行编写，LLVM，CUDA等就有自己的lowering方法。

Auto TVM：这个是TVM一种自动化优化方法，其主要思想就是基于强化学习的方法，搜寻一个最优的优化空间，实现网络调优。它的采集硬件实现的不同表现，来不断更新新的优化变量。通常一个网络都很大，搜索空间巨大，所以这是一项耗费量大的工程。而且开发者的硬件特点不同，可能有自己的独特优化方式，所以autoTVM是一个可选择优化方式。

Runtime：编译器提供了一个端到端的解决方案，runtime是最终可以在硬件平台上运行的程序，它包含了编译出的代码或者指令，硬件驱动，软件端调用。

以上是TVM的主要结构，在来看一下TVM代码的基本构成。

图中箭头表示了相互逻辑关系。

Support：架构的一些通用组件，比如socketing，logging等。

Frontends：TVM可以实现tensorflow，pytorch，onnx等很多深度学习框架到其自身IR表示的转化，这是frontends的主要工作，frontends就包含了针对每种深度学习框架图结构的转化函数。

Relay：这是一个高层次图结构的描述，它有自己的IR表示，用这些IR表示来描述神经网络结构。一些高层次的优化也是在relay IR的基础上进行的。

Autotvm：可选的基于强化学习的网络优化方法。

Topi：这是一个tensor计算库，里边包含了很多神经网络通用的计算单元，比如矩阵乘法，点乘，卷积等。我们可以通过这些topi来进行计算。

Tir：相比于IR来说，这个层次的计算单元更加底层和原始，更接近硬件实现。

Te：te是tensor expression，用户可以通过调用te中的函数来构建tir。TVM不仅可以将tensorflow等深度学习框架模型转化成自己的IR结构，还允许用户自己通过这些te来构建一个网络。

Node：node是在IR的基础上增加了一些新特性，可以允许用户对一些函数进行访问。这样就能够实现更加复杂的表达函数。比如如下的代码：

A和b就是tir.add的节点。

Driver：针对硬件的驱动。

Arith：这个紧紧绑定了TIR，可以帮助进行TIR优化的时候进行一些分析。

Runtime：runtime封装了图结构的转化，优化，代码生成，以及程序在硬件上的执行，为客户提供一个API接口完成所有的编译过程。