https://arxiv.org/pdf/2210.15016.pdf

本文将对TPU-MLIR的论文进行讲解，介绍其编译流程及内部工作机制，帮助读者在之后的开发流程中更好地开源。

概览

在这里插入图片描述

TPU-MLIR概括有如下四个特点：

多框架统一处理 ，支持多种神经网络框架（pytorch/onnx/tensorflow/paddlepaddle等），且处理流程完全统一；
分层设计 ，将网络模型抽象成2层处理：
1. TOP层，与芯片无关层，包括图优化、量化、推理等等；
2. TPU层，与芯片相关层、包括权重重排、算子切分、地址分配、推理等等；
正确性保证，
1. TOP层推理结果与ONNX结果对比验证，确保一致；
2. TPU层推理结果与TOP层对比，确保精度可靠；
3. Cmodel推理结果与TPU层对比，确保硬件一致性；
  4.过程可跟踪，每一步的转换可以生成MLIR，用于跟踪和调试。

模型转换

在这里插入图片描述

ONNX
Caffe
TFLite
Pytorch/Tensorflow/PaddlePaddle --> ONNX
ONNX提供了丰富的算子和一些基础数据类型，其通过多个node来定义计算图，每个node定义了输入（input）、输出（output）、类型和属性等，输入输出都是符号（symbol），通过这些符号的关系来确定计算依赖，完成计算图的定义。

Pytorch/Tensorflow/PaddlePaddle --> ONNX过程由外部工具完成。

Translation

在这里插入图片描述
ONNX/Caffe/TFlite --> TOP MLIR
TOP DIalect接近于原始计算图，Op的定义与ONNX和Pytorch近似，表示高层的抽象计算，与具体硬件无关。
代码位置：

./tpu-mlir/python/transform
|--BaseConverter.py
|--CMakeLists.txt
|--CaffeConverter.py
|--MLIRImporter.py
|--OnnxConverter.py
...

举例：resnet18.onnx --> resnet18_opt.onnx --> final_opt.onnx --> resnet18_origin.mlir

Canonicalize

包含算子融合，计算化简等。
使用MLIR自带的class。
在这里插入图片描述
代码位置：

./tpu-mlir/lib/Dialect/Top/Canonicalize
|--Add.cpp
|--BatchNorm.cpp
|--Compare.cpp
|--Concat.cpp
|--Conv.cpp
...
|--Scale.cpp

scale函数
左边为原始mlir文件，右边为优化后；左边多个scale合并为一个scale，scale又转化为一个conv的过程。

tpuc-opt --init --canonicalize --mark-FLOPs --save-weight --mlir-print-debuginfo Scale_original.mlir -o Scale.mlir

Lowering

在这里插入图片描述
将TOP Dialect转化为TPU Dialect。
TPU Dialect是用于表示TPU芯片的Kernel库，与具体的设备有关。TPU Dialect可以表示内存分配，软件流水，计算和数据读写并行等与最终指令生成相关的信息。
该过程包含：

Operation Conversion
Type Conversion

代码位置：

./tpu-mlir/lib/Conversion/TopToTpue
|--BM1684
|--BM1684X
|--CMakeLists.txt
|--CV18xx
|--LoweringBM1684.cpp
|--LoweringBM1684X.cpp
|--LoweringCV18xx.cpp
|--TopLowering.cpp
|--TopToTpuPass.cpp

LayerGroup + Bufferization

在这里插入图片描述

Tile+Fuse (layer group)
用重复计算代替部分数据搬运
计算与数据搬运并行
Memory分配优化

代码位置：

./tpu-mlir/lib/Dialect/Tpu/Transforms
|--AddressAssign.cpp
|--BM168X
|--CV18xx
|--DynamicLayer.cpp
|--DynamicNetlr.cpp
...
|--LayerGroup
|--LayerGroup.cpp
|--StripIOQuant.cpp
|--SubnetDivide.cpp
|--WeightReorder.cpp

LayerGroup过程：
在这里插入图片描述
通过深层次的依赖关系来进行更为合理的切割：越深依赖关系越复杂，不限于一个op层，不是在一层conv上做切割，而是计算一个conv输出被下一个conv利用的相关性，做整个group的相关性的切割，用计算来代替搬运。

Calibration + Quantization

对经过优化的TOP MLIR进行多次前向推理，获取每个中间Tensor的数据，并计算它们的统计信息，通过KL方法得到初步的阈值，然后使用误差/余弦相似度方式再微调阈值，是的INT8的计算结果和FP32的结果尽量相似。
在这里插入图片描述
代码位置：

./tpu-mlir/python/calibration
|--data
|--data_selector.py
|--gen_data_list.py
|--kid_calibrator.py
|--mix_precision.py

在这里插入图片描述
得到量化表：

run_calibration.py resnet18.mlir --dataset ./test_img/ --input_num 100 --tune_num 5 -o resnet18_cali_table

在这里插入图片描述
将量化表导入到tpu mlir中去：

... --import-calibration-table=“file ...

Correctness Check

在这里插入图片描述
TPU-MLIR提供对TOP和TPU Dialect的Inference。通过比较对应数据的相似性，来确定整个转化/编译过程的正确性。同时由于可以比较每个中间Tensor的结果，开发者可以快速地定位错误点，便于Debug。

代码位置：

./tpu-mlir/lib/Dialect/Top/
|--Interfaces
|--Transforms

./tpu-mlir/lib/Dialect/Top/
|--BM684
|--BM684X
|--CV18xx
|--Common

./tpu-mlir/lib/Support
|--Dnnl
|--ModuleInterpreter.cpp
...

在这里插入图片描述