论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR
论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR
- https://arxiv.org/pdf/2210.15016.pdf
-
- 概览
- 模型转换
- Translation
- Canonicalize
- Lowering
- LayerGroup + Bufferization
- Calibration + Quantization
- Correctness Check
- 相关资料
https://arxiv.org/pdf/2210.15016.pdf
本文将对TPU-MLIR的论文进行讲解，介绍其编译流程及内部工作机制，帮助读者在之后的开发流程中更好地开源。
概览
TPU-MLIR概括有如下四个特点：
- 多框架统一处理 ，支持多种神经网络框架（pytorch/onnx/tensorflow/paddlepaddle等），且处理流程完全统一；
- 分层设计 ，将网络模型抽象成2层处理：
- TOP层，与芯片无关层，包括图优化、量化、推理等等；
- TPU层，与芯片相关层、包括权重重排、算子切分、地址分配、推理等等；
- 正确性保证，
- TOP层推理结果与ONNX结果对比验证，确保一致；
- TPU层推理结果与TOP层对比，确保精度可靠；
- 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可以表示内存分配，软件流水，计算和数据读写并行等与最终指令生成相关的信息。
该过程包含：
- 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
...
相关资料
- 按照Readme操作，了解运行过程：https://github.com/sophgo/tpu-mlir
- TPU-MLIR的设计思路：https://arxiv.org/abs/2210.15016
- 开发计划：https://github.com/sophgo/tpu-mlir/wiki/Roadmap%5BCN%5D
- 工程结构：https://github.com/sophgo/tpu-mlir/wiki/Tutorial%5BCN%5D
- 技术细节可以参考：https://tpumlir.org/docs/deverloper_manual/index.html
- TPU-MLIR官网https://tpumlir.org/获得更多信息，包括文档和视频资料
欢迎大家一起参与学习和开发TPU-MLIR。