论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR

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

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

概览

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

1. 多框架统一处理 ，支持多种神经网络框架（pytorch/onnx/tensorflow/paddlepaddle等），且处理流程完全统一；
2. 分层设计 ，将网络模型抽象成2层处理：
   1. TOP层，与芯片无关层，包括图优化、量化、推理等等；
   2. TPU层，与芯片相关层、包括权重重排、算子切分、地址分配、推理等等；
3. 正确性保证，
   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

左边为原始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可以表示内存分配，软件流水，计算和数据读写并行等与最终指令生成相关的信息。
该过程包含：

1. Operation Conversion
2. Type Conversion

代码位置：

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

LayerGroup + Bufferization

1. Tile+Fuse (layer group)
2. 用重复计算代替部分数据搬运
3. 计算与数据搬运并行
4. 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
...

相关资料

1. 按照Readme操作，了解运行过程：https://github.com/sophgo/tpu-mlir
2. TPU-MLIR的设计思路：https://arxiv.org/abs/2210.15016
3. 开发计划：https://github.com/sophgo/tpu-mlir/wiki/Roadmap%5BCN%5D
4. 工程结构：https://github.com/sophgo/tpu-mlir/wiki/Tutorial%5BCN%5D
5. 技术细节可以参考：https://tpumlir.org/docs/deverloper_manual/index.html
6. TPU-MLIR官网https://tpumlir.org/获得更多信息，包括文档和视频资料

欢迎大家一起参与学习和开发TPU-MLIR。