BYOC：加速器硬件接入深度学习编译器的统一框架

本文是对Bring Your Own Codegen to Deep Learning Compiler的文章解读，内容结合TVM Conf视频和相关slides。部分内容的关键节点，会参考TVM里的代码进行简要分析。代码分析部分难免出错，如有错误还望指出，非常感谢。
深度学习的应用落地，体现在在不同硬件加速器上的部署，以支持快速高效的任务推理。通常来说，每个硬件加速器GPU、NPU、XPU都有各自的繁杂的软件栈，乃至软件生态。仅编译软件栈就包括要对代码的指令翻译、优化、最终在自家设备上把DNN给它跑起来。虽然是硬件厂家，但在软件栈上的投入的开发以及维护精力是巨大的。

图1 从底层硬件、计算库、中间表达到深度学习框架

此外，硬件厂商还需要迎合市面上DNN模型架构迭代以及算子库的变化，满足客户的需求，同步地、不断地更新硬件和软件。

图2 开源深度学习框架

以上这些都带来了巨大的软件工作，为了解决这一问题，本文提出一个开源框架，可以基于现有深度学习编译器，让开发者快速接入自家硬件加速器，后续只需要关注自家代码生成工具的工作，从而复用编译器的其他模块。换句话说，用户通过本文提出的框架的接口，可以很容易地将模型切分为子图，并且对异构环境而言，不同子图可以跑在性能最佳的硬件设备上。

图3 深度学习框架 Vs. 深度学习编译器

我们的工作也在多个硬件设备上实现了最佳性能的部署，而代码量仅有千行。

图4 深度学习编译器的组成

1 Introduction & Motivation

因为云上部署在存在网络延迟和数据隐私问题，所以现在在一些特定常见下，边缘端部署成为主流，为了能应对计算密集型任务，过去几年也有不少专用加速器上市，以达到硬件峰值性能，如Nvidia Jetson系列、Xilinx Vitis AI、ARM Ethos-N系列以及Google Edge TPU。

模型往往由重复的计算密集但是参数不同的block组成，每个block通常都是特定顺序的卷积、矩阵乘法等其他算子组成。相应，硬件加速器就会为其定制化算子，以保证在矩阵计算等操作上的性能。但随着算法模型的不断演进，更加复杂的算子如数据处理，更着重受限于带宽而计算量不大的算子出现，如控制流一类的模型架构，这都在挑战硬件厂商的软硬件团队，软件团队需要基于现有硬件设计更好性能的算子库，硬件团队也要考虑新的尚未支持但常用的算子计算范式。

图5 R-CNN模型结构，位于区域候选网络（RPN）的非极大值抑制（NMS）和后处理的循环较难支持与优化

除上述模型演进带来的软硬件团队工作外，为了满足模型的表达，不同框架模型的格式。加速器厂商需要开发完整的编译软件栈，来接收不同框架的模型格式、对模型优化、算子翻译成对应硬件兼容的代码（其中要兼容不同架构），最终才在硬件设备上运行模型。对于不同硬件厂商来说，如果有一个统一的框架，可以减少很大程度的重复性工作，让各自厂商更关注于自身加速器计算内核的优化和代码生成。总之，将刚刚内容总结一下，现有的边缘加速器编译器堆栈的研究和开发由于多种原因而受到阻碍：

挑战1：算法迭代导致运行不起来。模型结构迭代，如ResNet在AlexNet基础上增加了残差链接，Inception又引入bottleneck，到SSD和Mask R-CNN再次引入控制流。控制流由于存在IF节点，很难被现在基于数据流架构的加速器优化；
挑战2：低效算子的性能问题。虽然整体模型结构简单，但是部分算子具有复杂控制逻辑，加速器并没有支持。这部分体现在目标检测网络的尾部，如R-CNN的RPN中的NMS等，非极大值抑制用来过滤候选框减少冗余。通常，加速器重点优化的是张量计算，这种类型的操作虽然占模型比重不大但耗时占比高，且并不好支持如矢量化的优化；
挑战3：通用优化的巨量工作。即使模型中所有算子都支持，但其中在性能优化会花费巨量工作，因为缺少一个统一的且完整的编译软件栈，或者说是框架。这个软件栈框架的工作很大一部分类似标准编译器一般——做平台无关的优化技术，如死代码消除、公共子表达式消除、常量折叠等。剩下一部分则是特定的和模型有关，或者说硬件相关优化如数据排布变换等。

为了解决上述问题，我们提出了一个给硬件厂商的统一接入框架，利用对模型统一标准处理流程，厂商只需要关注专有的代码生成部分，以缩短软件栈的开发周期。

图6 BYOC框架将ASIC硬件接入TVM

总之，这个统一的接入框架基本思想是解耦深度学习编译器（Deep Learning Compiler），将其分为通用和专用两部分：

通用部分：被CPU、GPU等AI加速器共用，以应对新模型即解决挑战1。
专用：可以由每个硬件加速器厂商维护。

该框架为提供对加速器不支持的模型部分进行切分，留在设备CPU上执行。这样解决了挑战2。本文提供的这套框架可以用于各种深度学习编译器里，厂商只需要关注自家硬件平台的优化即可，解决了挑战3。

图7 BYOC工作流程概览

这个框架主要有如下几点贡献：

面向加速器厂商提出统一接入框架，可以无痛将自家硬件（尤其是代码生成模块）接入已有的深度学习编译器，通过深度学习编译器已有的优化手段做硬件无关的优化；
异构切图接口简单好用。包含对计算图的标注与切分功能、可以应用硬件相关的图切分或等优化策略，以提升性能；
扩展新硬件工作少。引入集成多个实际的边缘加速器，代码生成模块集成进来，只需要大概两千行代码；
该框架已在多个商用加速器中使用，并已融入他们的模型编译优化流程中。

图8 目前已通过BYOC接入TVM的加速器及其编译工具链

2 Framework Design and Implementation

下图是我们框架的编译流程。该框架是基于一个深度学习编译器搭建起来的，其中白色的是编译器中已有的部分，可以被不同硬件复用，灰色的是硬件加速器特有，各家硬件针对性定制实现。

图9 框架编译流程。灰色矩形是主要模块，每个点是模型中的节点

根据上图的流程，模型读取进来后，先构图得到对应中间表示（IR，如TVM里高层级表达Relay IR），然后执行硬件无关的一系列图级别优化，如算子简化、常量折叠等。之后，优化后的图会被为两种类型的很多子图，这两种类型分别是可以在加速器上执行的以及不能在加速器上执行的（只能在主机端CPU执行）两部分。两类子图，将会进行各自的后续编译流程、针对性优化、代码生成等，最终生成可以执行的kernel，供运行时读取使用。

以RCNN模型为例，可能会有4个子图由加速器来执行：整个模型骨架，兴趣区域检测头模块、RPN部分除NMS以外的算子、后处理部分除去控制流相关的算子。主机端子图是剩下保留的算子，交给CPU或者GPU来完成计算。加速器运行和主机端运行时最后会打包到一起。

图10 TVM中的层级表示：Relay IR、图标注与图切分、代码生成、运行时

代码以Ethos-N加速器为例，图切分、标注对应的TVM实现代码见 partition_for_ethosn函数中 passes（对应上图 Annotation、 Partitioning等），该函数传入后续需要被做处理的Pass的 Module mod等，返回切分好被标记的 Module，个别代码已省略。

对上面的Python代码不太理解没关系，后面会有丰富的图例进行说明。这里的 transformation或称为 transform后面会反复出现，它主要有两个用途：① 对当前程序优化；②下沉当前表达为更贴近目标硬件的版本。transform往往包含一系列用于优化模型的 pass。

下面我们会具体讲讲上图的编译流程的每个部分。拿到Relay IR后，先从图标注与图切分开始。

2.1 图切分

加速器通常加速的是模型中占比最多的计算密集型算子，如矩阵乘法等大量需要乘加计算的操作，而对其非此类算子，即我们前文提到的控制流等算子并不易并行实现的算子可能往往得不到加速，甚至会执行失败，为了保证这些算子能成功运行和高效，通过切图模块将他们交给主机端如主机CPU来计算。因为现代编译器为了做不同层级的性能分析，通常都是多级中间表达，每个层级的中间表达的目的不同，因此切图的策略与条件也会不同。

框架中，切图前的图表达是一种高层级表示，如Conv2D会包含其属性（stride、padding、dilation等信息），使用高层级表示来做切图有两个原因：

加速器厂商的算子库也是带有一定限制的，并不会暴露过多信息。算子库通常是支持哪些算子，甚至是规格信息，因而只要在切图时候，对于建立加速器算子库与对应图中算子的类型和相应属性的映射关系即可，无需更低层级的硬件信息；
加速器厂商对每个算子可能有自己厂内内部的低级表达，有自己厂内的深度学习编译器等内置算子表达，在高层级切图也可以更灵活地连接低级表达。

下面是一个针对某硬件加速器切图的例子，切出硬件加速器支持的子图，剩下是不支持的子图交给主机端处理器计算。

图11 图切分示例，后文会提到

有两种图切分为子图的方式，分别是针对多个算子的子图划分与单算子的子图划分。

图切分：1. 基于模式匹配的子图划分（Rules for graph patterns）

加速器厂商为了能让模型中固定的算子模式跑在加速器上，需要通过图的模式匹配算法对模型原图进行分析，切分出支持加速的子图。而如果用我们的框架，硬件厂商无需实现这一过程，我们实现了模式匹配机制并提供一个易用的编程模型供加速器厂商来自定义需要匹配的算子模式。例如下面是一段匹配 Conv2D-Add-ReLU模式的代码：

图12 针对 Conv2D-Add-ReLU的模式模板

上述代码中 wildcard()表示三个数据节点， is_op用来给定算子类型做匹配，如 Conv2D， optional表明这个 Add在 Conv2D的节点上是可选的模式，也意味着此处的模式匹配 Conv2D-Add以及 Conv2D两种。最后把定义的模式，放入模式匹配表中 pattern_table。

图13 图切分

通过 pattern_table，可将上面（a）中的子图经过模式匹配分组，会被替换为（b）中的一个三角形单节点子图，在后面代码生成阶段，这些序列模式会被映射为加速器相应的指令。TVM里在匹配前是这样的：

匹配完成后，变为如下的函数 fn形式：

这个新函数 fn的输入输出即将上面三个算子的第一个输入和第三个输出，其 Composite名字是 pattern table设置的名字， PartitionedFromPattern即 Partitioned前的形式。在后续 runtime.module通过 GetFunction方法，传入 fn这个名字就可以拿到这个 PackedFunc。

图切分：2. 基于图标记的子图划分（Rules for single operators）

不仅可以对多个算子组成的模式进行匹配划分，还可以对没有复杂模式的单算子进行划分，即定义当前支持加速器的算子列表，例如下面代码为 codegen.myAcce1这个加速器的代码生成模块，注册了名为 nn.conv2d的算子，并在代码中限定了参数数据类型都必须为 float32，表示当 nn.conv2d满足符合这样的条件时，就会转为加速器支持的 Conv2D算子节点。

图14 为加速器 Acce1注册参数为 float32类型的 Conv2D算子，此外还可以添加其他限制如NCHW某维度支持的最大值

通过对一系列的子图做相关加速器的类型标记，我们可以把加速器支持的所有算子都予以支持。但是考虑一个问题，当一个原始模型的图被四分五裂为多个加速器子图，但是加速器子图之间存在一些加速器不支持的算子时，就会带来主机端与加速器间的数据传输、加速器启动开销。为此我们采用贪心地策略，将连续且加速器支持的算子进行合并为一个子图区域并交给注册给加速器的代码生成模块。在前文图中示例（c）虚线圈起来的2个部分，即表示有2个合并后的区域。

目前我们只关注一个加速设备时候的例子，对于多个加速器，我们的框架也支持用户对优先级进行设定，如某算子可以给计算更快的多个加速器中的某个。多设备情况下，让每个算子达到最佳性能的问题是一个开放问题，不在本文中讨论。

下面是TVM代码部分，以DNNL（Intel OneDNN）为例：

经过上述两种子图划分的方式，代码对应上面的BYOC Relay passes，就可以完成下面图(a)到(d)的过程。

图15 Graph Transformation on original relay graph

其中，(a)是原始的Relay Graph；(b)黑色的框相当于对算子进行多算子模式匹配分图、单算子分图后的结果，如 Conv2D_add_ReLU、 Conv2D、 ReLU，在这里黑色框被称为 region；(c)会将连在一起的 region进行合并；(d)已经是异构图了，可以看到 Fn1即(c)中合并后的一个大region，同时会有相应 Compiler生成对应加速器的代码，即 CodeGen结果，在指定加速器硬件上交给 Runtime去跑。

刚看了Intel OneDNN的情况，下面再看看Arm ethos-N后端的相关代码，从 pattern_table实现开始：

类似前文的 conv2d_pattern，不同在于这里 is_op("qnn.conv2d")这样的代码，其中 qnn表示量化的算子，与常规的 conv2d相比，二者参数接口不统一，例如卷积 qnn.conv2d相比 nn.conv2d多了scale、shift等量化参数，而在算子匹配的时候，前面可能会有 nn.pad来做硬件相关的数据行对齐，后续还有 qnn.requantize的结构等。

这是一个匹配模式，因为对一个硬件后端支持时，必然有很多算子当然可能也有很多模式。

pattern_table的返回形式元组列表，每个元组是 pattern名、 pattern子图逻辑、 pattern匹配的限制条件函数的形式，其中限制条件是非必要参数，如检查参数类型和维度、参数限制，返回的结果会作为 transform.MergeComposite(pattern_table)(mod)的参数。

而如果是单节点的算子，虽然不需要写 pattern匹配多个算子，但在注册的时候，需要在注册实现中加上对应参数在该加速器上支持情况的检查，也相当于pattern_table返回的元组列表里的第三个参数（即限制条件函数）。

当我们定义好以上 pattern规则等，接下来就是按照pattern进行切分的代码流程 partition_for_ethosn，下面这部分代码的核心部分，已经在前文出现多次：

需要注意的是，这个 partition_for_ethosn函数 relay不会调用，应该不是标准的处理流程的函数，在 tests/python/contrib/test_ethosn/infrastructure.py单元测试中 partition_for_ethosn会被调用，完成后会生成很多局部函数如F1、F2等，以上就是前端的处理过程。需要注意的是，这里的"ethos-n"，与后面进一步下沉有关。

这里代码分析不保证正确，参考。

后续， Codegen方法进而会调用 src/relay/backend/te_compiler.cc的 transform::PassLowerTE方法， LowerTE方法这是 Relay 编译器的“后半部分”，都有对 LowerTE的调用，它对 primitivefunction 进一步下沉（lowering）到 TE 表达式、调度，并 emit PrimFuncs。LowerTE会调用 LowerExternalFunctions函数，根据 module 的名字和相应编译设置 CompilationConfig，对 ProcessFn进行下沉， LowerTE该部分代码实现如下：

LowerTE里会调用 LowerTensorExpr方法，该方法又会调用位于 relay_to_tir.cc里的 RelayToTIRVisitor.Mutate，之后会调用 compiler->LowerExternalFunctions()，根据 "relay.ext."+opt_compiler.value()拼接得到外部加速器的 Codegen工具，并对其生成 runtime::Module，供后续使用。

图切分：3. 基于代价的分图（Cost-based Partitioning）

虽然我们想尽可能多地合并子图上连续部分的加速器算子，但实际情况可能会受限于加速器的资源限制而做一些切分，如有限的内存或者计算单元等。

为此，我们又提出了基于代价的切图机制，可以将上一步骤支持的连续合并为一个区域的多个算子，按照用户定义的策略来切分，如超过有限的内存则做切分等等。前文图切分的（c）到（d）就是将原本（c）中由 4N+1个算子组成的 region1，因为触发用户定义策略而被切分为只有 2N+1个算子的新的 region1以及有 2N个算子的 region3。

由硬件资源限制带来的切图，往往导致原图模型被切成由主机端计算子图连接的多个硬件子图，计算任务从主机端到加速器会有数据传输和计算任务启动开销，所以应尽量减少这种切碎的图的情况发生，再就是主机端的计算负载尽可能低，比方任务的计算密度或乘加数量，重计算任务都尽可能放到加速器上。

（e）表示每个交给加速器算的区域都会封装到一个单独的函数中，并且打上对应跑哪个加速器后端的标志。因为在对应不同加速器生成代码前，还会进行加速器专有优化处理的过程。下面来看看。

TVM里拆出的子图更像LLVM IR以Module Function Call，即表达式子函数的形式来组织IR。对应代码实现位于 python/tvm/relay/transform/transform.py，其核心函数是 transform.PartitionGraph(mod_name="default",bind_constants=True)，会将Relay program切分为在不同后端上执行的子图，其接收参数 mod_name为字符串，切出来的图会以 tvmgen_<mod_name>为前缀，第二个参数 bind_constants，表示是否将constant节点绑定到子图中，其中的实现就会走C++实现的 PassPartitionGraph(Stringmod_name,boolbind_constants)（位于 src/relay/transforms/partition_grpah.cc）。

2.2 加速器特定处理

前文图切分后，各个区域会调用相应的计算后端执行计算。

图16 加速器特定处理，菱形节点表示特定处理新插入的节点。（a）表示子图计算前后插入量化和反量化节点；（b）表示在子图前后插入数据排布变化节点

对于主机端的部分如Arm CPU，会使用深度学习编译器中预设对该CPU设定的常见标准的优化流程，而加速器部分则需要硬件特定的优化：操作合并，化简，数据排布变换，量化等。其中，量化和数据排布变换是两个常见的加速器相关的处理。

图17 TensorRT 8bit Inference

量化是指计算过程以相对 float32数据精度更低的如8bit、16bit类型来执行，将浮点，可以减少功耗，提高带宽利用率，某些加速器只支持定点，所以需要校准数据集对模型做预先的量化，让加速器的终端用户在编译和部署模型前对每个加速器子图做量化并不实际，我们的框架可以实现部分量化，即对加速器子图的整个输入做量化、输出做反量化。图中常数 S和 Z分别是量化、反量化过程中的缩放因子和相对零点位置的偏移量。

图18 图优化例子：使用TensorFlow XLA在Volta GPU对AlexNet在HLO级别上作优化，如算子融合、常量折叠、排布变换等

数据排布变换是与计算相关的另一种重要操作，加速器通常有适配自己计算的数据排布形式如 NCHW或 NHWC，在某一种排布上的算子可能会相比另一种排布会有比较好的访存特性，性能也更好。与量化一样，数据排布对于加速器子图前后的必要位置插入排布转换节点。

图19 深度学习编译器里的各种编译优化

在TVM对于加速器特定处理的代码层面，对于Ethos-N而言，量化节点的插入已经在 pattern_table的时候考虑了，如 qnn.conv2d。

2.3 代码生成

在图切分后，编译流程的最后一步就是代码生成，整个过程如下图所示，给定已经切分子图后的模型，框架会生成一个支持在异构设备上的可推理执行的模块。

图20 代码生成示例：将加速器模块各自打包，最终打包为一个整体

该模块包括：模型图结构、每个图节点的计算实现。如异构场景下，该代码生成过程，结合了不同硬件加速器厂商的代码生成工具和各自编译流程，交由不同加速器的分区函数生成装配相应“加速器子模块”即Accel sub-module1/2/3。最后，会将它们打包起来，呈现的样子就是各种sub-module和Host sub-module的集合后的一个模块，供部署时使用。

对于加速器厂商而言，代码生成就是将TVM Relay IR，生成自己加速器推理引擎支持的模型文件格式。

每个子模块里是各自加速器生成的代码，以一种约定的格式来表达，这样可以被各自加速器的执行引擎在运行时被调用。我们的框架提供如下三种生成代码格式：

标准图表达：JSON格式是默认的图表达格式，JSON包含算子名字、属性以及数据流，而且也可以被运行时引擎很好地加载。目前NVIDIA TensorRT以及Arm Compute Library都是使用JSON代码生成器，作为运行时引擎和我们框架的沟通语言；
标准C代码：尽管标准图表达易于实现和部署，但是仍然需要图引擎来包含所有实现的算子，尤其对资源优先设备带来的问题是二进制包比较大，所以能直接运行从标准C代码编译的可执行二进制文件，那再好不过。因而我们的框架为这些加速器以及主机主机设备，提供了一个标准的C代码生成器，通过发出算子计算内核库函数命令调用并链接。
因为主机代码通常与C兼容，该方案简化代码打包，允许库调用成为主机子模块的一部分。这意味着加速器厂商不必定制自己的，但可以充分利用现有的运行时系统。
定制化的图表达：方案1和2的代码格式是平衡了简单灵活、可读性。有些加速器只能加载自定义的图表达格式，例如ARM Ethos-N和Xilinx Vitis AI有各自的流式格式表达模型，对这种需求，我们也做了对定制化格式的支持，只要加速器厂商提供①编译生成代码并将其序列化为比特流；②在运行时反序列化为比特流的能力即可封装为我们框架的统一接口。其实，标准图表达的JSON格式，也可以归为定制化的图表达。这也是大多数有自己模型格式个加速器厂商，支持的代码生成格式。

这里，我们根据TVM教程，拉取ONNX ResNet50模型，并使用如下命令对模型进行 compile。

使用的 target是 llvm，得到输出为压缩包 resnet50-v2-7-tvm.tar，对其解压后包含三部分：

mod.json：TVM Relay计算图的JSON文本形式，即模型代码生成的结果；
mod.params：模型权重；
mod.so：包含模型用到的算子的C++库，可以被TVM runtime加载。

这部分是对 Runtime.Module部分代码解读，根据自己的理解，可能有错误。

下面以Arm NPU Ethos-N为例， runtime.Module的入口位于 src/relay/backend/contrib/ethosn/codgen_ethosn.h的 runtime::ModuleCompileEthosn方法，会全局注册 TVM_REGISTER_GLOBAL("relay.ext.ethos-n").set_body_typed(CompileEthosn)。需要注意，这里注册的名字 relay.ext.ethos-n，与 LowerExternalFunctions()中获取 external mod时 std::stringext_name="relay.ext."+opt_compiler.value();，二者名字是对应的， opt_compiler.value()，即 ethos-n加速器。

CompileEthosn方法会调用位于 ./src/relay/backend/contrib/ethosn/codegen.cc里的 EthosnCompiler类的 CreateRuntimeModule方法（调用该方法的是 codegen_ethosn.h里的 CompileEthosn方法）：

其中，最重要的是第12行 CompileEthosnFunc(mod,gvar,mod_func)，编译完成后得到module，而这个编译的过程 CompileEthosnFunc实现如下，大致流程为：

ConstructNetwork：从Relay function生成 Ethos-N 支持的模型，其中，会先构造一个 ConstructNetworkVisitor类对象（一个构建模型的PASS），然后该对象调用 Construct方法。
根据计算节点推理各节点的输入输出信息。 InferTensors，调用 classInferTensorsVisitor（一个推理信息的PASS）的 Infer方法，通过Relay expression来做当前加速库的张量信息的推导。并之后依次遍历 func->params参数Expr集合，根据查找表 std::map<Expr,std::vector<sl::TensorInfo>>tensor_table_;，根据传入Expr，获取到相应TensorInfo，按输入出现的顺序，添加到 network_with_ids.input_ids[tensor_and_id.operationId]中。之后会调用 InferTensorsVisitor::VisitExpr，但这个函数 VisitExpr的实现位于基类 templateclassExprFunctor的 virtualRVisitExpr方法，但是 codegen里有实现 VisitExpr_(constCallNode*cn)，入参为 constCallNode*cn的 VisitExpr_方法会调用 InferTensorsVisitor::InferCall。 InferCall会调用每个 EthosnAPI算子，做算子数据节点的信息推导与绑定，比方 QnnConv2d、 QnnFullyConnected等等算子，各个算子内会做相关各种 EthosN专有参数在这个算子的绑定，和参数检查，各式各样参数的Tvm2Npu的实现（如量化、pad、数据排布、数据位宽等等参数，而 Tvm2Npu实现位于 ./src/relay/backend/contrib/ethosn/ethosn_api.cc）
HandleCall会调用各个算子的Make方法 Make<XXX>Layer，并返回 MakeOps(Tensors)，将 callnode翻译为对应的Make方法，建立 call到 NPU的映射，对所有算子都这么操作。各个 Make<xxx>Layer(Call,sl::TensorAndId<sl::Operand>*out)如 MakeConvolutionLayer都是类 ConstructNetworkVisitor的方法，该方法将 Call转换为 Ethos-N的层，根据Call节点拿到这层的信息，拼成这一层的属性参数， Make<xxx>Layer这个过程是倒着来的，即倒序访问将 Relay function 映射为当前硬件支持的模型图。其中会用到 IsEthosnFunc方法如 IsEthosnFunc(call,"ethos-n.qnn_conv2d")，其中的 IsEthosnFunc(call,"ethos-n.qnn_conv2d")就是前文中 pattern_table里定义的字段如"ethos-n.qnnleakyrelu"。
CreateOptions：创建必要的编译选项，这些选项是供编译 Ethos-N 模型使用的；
sl::Compile：根据1和2得到的的模型、编译选项进行编译，得到 compiled_networks，其中， namespacesl=::ethosn::support_library;namespacedl=::ethosn::driver_library;；
建立TVM runtime与Ethos-N模型的输入输出映射关系。

以上就是 CodeGen的过程，从Relay IR到加速器的模型结构文件。

到此，讲了图切分、加速器的特定处理，以及本节的代码生成，模型编译过并打好的包就准备好了，下一部分会讲解如何让一个轻量级运行时系统，加载这个包并执行推理。

2.4 运行时

TVM的Runtime只是把整个图的输入和输出定义好，等待图执行的最终结果，中间的异构子图是运行在CPU上还是加速器设备上，完全由加速器工具来决定。因此，需要实现Runtime。

大多数深度学习编译器都是调用加速器自己的运行时来完成推理。代码生成部分，我们对确定在某个平台上要跑的模型，生成对应平台上的模型包，这其中就包含了该平台所需的运行时信息。

运行时负责模型中对应算子子图在对应设备上的执行，可以把图执行引擎当做一个访问图的数据流 visitor，或想象成一个可解释字节码、可处理动态可变输入及控制流的虚拟机。

图21 运行时的Metadata module与其工作流

为此我们设计了一个统一整体的运行时模块，以分层的方式来管理主机端和各个加速器的分层计算内核，而且这个整体的运行时模块，可以自己成进深度学习编译器中。我们使用下图来介绍这个框架的工作流程。

图22 运行时工作流程。(a)表示加载模型的运行时并从主机端开始执行第一个算子；（b）表示加速器执行外部函数调用；（c）表示（b）中图节点F2的内部执行逻辑

初始化原数据模块（上图黑色圈1）。模型往往包含一系列的权重参数，在推理时在运行时模块中以常数的形式体现，运行时还有模型的图以及生成的计算内核。即上图（a）部分， Metadatamodule包含了多个运行时子模块如 Hostsub-module、 Accelsub-module1等对应不同的目标设备，通过 Metadatamodule的统一分层管理，可以在 Metadatamodule初始化时：①分层地对各需要用到的运行时初始化，②相应设备平台权重参数的预先分配等。模型计算图体现了计算节点，当然也存在一个数据条目表，其对应着模型中的数据节点，包括常量、模型输入和输出以及中间结果。由于执行加速器子图计算的函数是外部函数调用，加速器子图的中间结果将不会保存在数据条目中，对实际用户是不可见的。换句话说，尽可能合并不连续的加速器子图的数量，可以让这个加速器优化内存占用。例如，如NVIDIA TensorRT拿到较多算子组合的子图时，可以在这个子图上作更多的算子融合等优化，从而优化内存使用。
在主机端开始触发计算图执行（上图黑色圈2）。当用户在给定模型输入如一张图片数据时，会触发推理过程，主机子模块 Hostsub-module加会载模型的图，并从主机上启动执行引擎。执行按照图节点的顺序来执行，上图（a）中都是计算节点， DataEntries是数据节点，对应计算节点的输入与输出；
主机端拿到输入数据并开始执行计算图（上图黑色圈3）。这个过程会调用计算内核，并将结果写到当前计算节点的输出，对应 DataEntries中的2， data entry 2也会成为下一个节点的输入。
加速器上的子图计算（上图黑色圈4）。上图（b）开始执行F1节点，执行细节见（c），分为两步：初始化该计算节点对应的加速器执行引擎，根据加速器实现不同，这个初始化时间可能从几秒甚至到分钟，如某些加速器引擎会JIT编译。初始化好的加速器引擎，根据自己的 sub-module读取输入完成计算，并将结果写到 data entry中，供后续计算使用。

当模型要执行时，有三种模式：对于简单已知维度且没有控制流的模型，可以用 graph executor后端来一层一层地执行；也可以通过 virtualmachine这一动态后端执行复杂模型；也可通过 ahead of time compilation即 AOT executor来加载整网执行，会事先将执行的整网络结构编译为一个可执行的文件，以上这三种模式都封装在统一的运行时 runtime.Module接口中。

运行时代码解读。运行时对应TVM代码部分是在 src/relay/backend/build_module.cc、 src/runtime/contrib/ethosn/ethosn_runtime.h中的两个类：classEthosnModule和 classRelayBuildModule类中，后者其成员方法 BuildRelay(IRModulerelay_module,constString&mod_name)会编译Relay IR module为Runtime module。

IRModule与runtime.Module
这两个概念很容易混淆，这里贴一下来自TVM文档的解释。
IRModule是TVM整个工作栈的重要数据结果，其字面意思为中间表达模块，模块内包含了函数的集合，其中的函数有两种主要形式：
IRModule形式1： relay::Function，一种TVM中的高层级程序表达。一个 relay.Function，通常对应一整个模型，也可以将 relay.Function理解为支持复杂操作（控制流、递归等）和数据结构的计算图；
IRModule形式2： tir::PrimFunc，一种TVM中的低层级程序表达，包含各种元素操作，如循环嵌套、多维度的读与存、线程操作、向量/张量指令。一般用于表达模型中某个算子的计算逻辑。
运行时的能力由runtime提供，其会加载编译好的模型文件并运行，runtime的重要模块是 runtime.Module，其包含了各种运行时函数，即 runtime.PackedFunc。runtime.Module作为上层模块会通过 GetFunction方法以及确定的名字（这个名字猜测是pattern_table里定义的），获取想要的 PackedFunc。
在编译过程中，一个 relay::function可能会下沉为多个 tir::PrimFunc小函数以及调用这些个 tir::PrimFunc小函数的顶层函数。

需要注意的是，BYOC框架对加速器就到这个Relay IR层级。

不过先来看一下 classEthosnModule，其继承自 classModuleNode，父类的几个重要方法：SaveToFile、 SaveToBinary将module保存为文件或二进制流， GetSource从module中获取源文件， GetFunction从module中（根据 function名，可选）获取 PackedFunc等（见下面代码），找到返回的相当于一个匿名函数。后续在调用的时候 TVMArgsargs,TVMRetValue*rv作为实际入口参数，然后根据 args， name，把需要参数给进来。

子类继承父类的实现，用来创建各种不同的运行时 module，比方子类 EthosnModule，该子类实现或者覆写了 EthosnModule构造方法、 GetFunction、 SaveToBinary、 LoadFromBinary、 SaveToFile。与父类不同主要是新增 LoadFromBinary，其会从 dmlc::Stream中加载出 function的个数，输入、输出所占空间大小，以及 OrderedCompiledNetwork的内容，最终返回 EthosnModule对象。

GetFuntion会进一步调用 Inference， Inference有两个实现，都位于 src/runtime/contrib/ethosn/ethosn_device.cc，分别对应于使用硬件和不使用硬件的时候，宏 ETHOSN_HW若打开，则会编译另一个实现的 Inference，以及 WaitForInference、 InferenceWaitStatus等方法。

当时打开 ETHOSN_HW时， Inference会执行下面代码的5个步骤：解包参数如输入输出个数相关指针、初始化输入输出8bit的 Buffer、掏出输入输出的裸指针、调用加速器自己的 Inference即 ScheduleInference做推理并等待 WaitForInference(inference.get(),60)让结果返回、最终拷贝拿到结果。

3. 案例与实验

略，我个人觉得实验部分在对性能的benchmark比较，其实意义不大。BYOC框架的目的在接入硬件，用上现有深度学习编译器的切异构子图的切图能力，以及更上层的图优化等，减少工作量。

4. 总结

本文提出一个可以将加速器厂商的代码生成模块与运行时集成到深度学习编译器的统一框架。在这一框架的加持下，对于通用的代码性能优化工作就交给标准编译器，对于深度学习的优化技术即高层级的工作则交给深度学习编译器，对硬件厂商而言只需要关注其内部算子的代码生成质量。

我们的框架通锅图切分、加速器特定过程、代码生成、运行时得以实现。现已有不少商业加速器厂商通过我们的框架，将他们的编译软件栈集成到其深度学习编译器中，如NVIDIA Jetson Xavier、Xilinx Vitis-AI。

作者：开心的派大星
文章来源：NeuralTalk

推荐阅读

更多嵌入式AI干货请关注嵌入式AI 专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。