文章目录
- BYOC架构
-
- 算子标注
-
- 单算子标注
- 复合算子标注
- Cost-based Partition
- Codegen
-
- Codegen for C
-
- 代码生成流程概览
- 代码生成工程实现
-
- 实现CodegenC
- 实现CSourceCodegen
- Codegen for JSON
-
-
- 实现JsonCodegen
-
- Runtime
-
-
- JSONRuntime
-
- 参考
随着后端设备数量激增,为达到较高的效果在这些设备上,对应的知识要求也同步增加。为缓解这些压力,硬件厂商会提供算子库如OneDNN、cuDNN或者推理引擎TensorRT以指定的方式描述模型以达到较高性能,但是这种方式需要用户学习新的编程接口。BYOC框架提供了方法使硬件厂商较为容易的实现自己的codegen并注册为Relay后端编译器支持自己的硬件或者算子库。
自己的ASIC加速器必须有自己的编译流,一般分为两类:
其一,生成一个图表示然后投喂到图执行引擎中。加速器上需要一个可以执行计算图的图执行引擎,可以降低算子间内存传输或者做算子融合等操作。这两种优化操作一般是在编译期间在图上做处理的,比如conv2D和bias add在TVM中是两个单独的算子,但在专属加速器中可能是一个算子,需要在优化图上通过新算子替换调原始算子。
其二,生成汇编代码或C代码,编译为可执行程序。对于没有端到端执行框架的平台需要提供一个编译器来编译目标平台ISA汇编代码表示的程序。对应的就需要codegen模块将Relay转化为C代码,这种厂商会提供对应的算子库来调用。
下图为BYOC框架对应后端的类型已经相应的数据结构。
BYOC架构
上图是BYOC(Bring Your Own Codegen)框架整体结构。其主要逻辑是对算子按照支持的执行平台进行划分,然后按照不同的编译流进行处理。详细介绍如下,首先加载模型,转化为Relay格式,然后进行一些后端无关的优化操作。第三步就是分割计算图,这里分为四个小步骤,后面会详细介绍。主要职责就是把计算图分为两类子图,一类是包含加速器支持的算子构成的子图,一类是不支持的,由默认的TVM编译流处理。在加速器这条编译流上可以添加加速器专属的优化Pass,可以自定义codegen。加速器中的算子会被包装为一个外部函数调用,算子执行由加速器处理,结果返回给TVM的runtime。
算子标注
流程图
在得到优化后的Relay计算图后,框架需要指定哪些算子会在指定的计算平台上执行,因此需要给算子做标注。BYOC支持单一算子和复合算子标注。在单一算子中,可以通过python的包装器为算子添加编译器属性确定执行平台,对以复合算子,需要先设计匹配pattern在计算图中识别子图以生成复合算子,复合算子的格式是一个composite复合函数。composite函数不用单独添加编译器属性,因为在匹配pattern时,pattern的名称中第一个“.”前面的字符串就是编译平台名称,这点可以在源码中看到。
单算子标注
标注规则在python/tvm/relay/op/contrib/your_codegen_name.py目录下。可以通过包装器@tvm.ir.register_op_attr(relayOp, target.yourTarget)直接将算子与指定运行平台绑定。该包装器是给算子添加了新属性target.yourTarget,BYOC框架会对每个算子调用target.yourTarget()检查是否支持该硬件。如下helper函数可以批量处理算子标注。
复合算子标注
添加的后端往往会将一组算子转化为一个优化后的指令或者一个高级API,因此需要先对计算图做划分。如上流程图所示,Conv2D+Bias Add+ReLU组合为一个复合算子。对于复合算子,可以设计匹配模板,然后通过@register_pattern_table进行注册。pattern可以通过TVM提供的Relay pattern language编写。
对于生成的组合算子,包含两个重要属性即Composite,表示该算子在target上满足的模式;PartitionedFromPattern表示满足的匹配模板。因为设计的模板是类似于通配符,具体匹配的是哪个pattern是不确定的,而该属性就显示了实际匹配的pattern是那个。比如conv+*的模式既可以匹配conv+relu又可以匹配conv+add等,而PartitionedFromPattern的就是实际的匹配模板如 nn.conv2d_add_或 nn.conv2d_nn.relu_。
对于复合算子,加速器可能会有对应的指令或接口。因此可以实现一个映射函数将复合算子与接口对应。如下,conv+relu可以通过DNNLConv2d(false, true)匹配,conv+add+relu可以通过DNNLConv2d(true, true)匹配。
Cost-based Partition
经过标注后的计算图被划分为若干个子图,可以进一步进行融合,组合一个大的算子,封装为一个函数,这样若干算子整体只做一次内存调用,一次kernel启动,减少了内存负载和内核调用的耗时。具体来说,BYOC采用MergeCompilerRegions Pass把在相同执行平台上的连续地算子进行了合并,组合为一个大的region,封装为一个外部函数调用。如上图,融合了三个函数,就有三个外部函数调用Fun1、Fun2和Fun3。函数有编译器的属性来确定执行平台。函数具体执行由加速器负责,调用由TVM编译器负责。
MergeCompilerRegions默认的融合策略是贪心的,他会尽可能多地把所有满足条件的算子合并。但是由于硬件平台的内存受限,因此需要控制融合后算子大小。这里论文里简单介绍一种方式就是直接控制算子的数量,但是其他框架会进一步考虑内存的信息。并且这些区域可