论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR
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。