论文讲解——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: A Compiler For TPU Using MLIR_第1张图片

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: A Compiler For TPU Using MLIR_第2张图片
代码位置:

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

论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第3张图片
左边为原始mlir文件,右边为优化后;左边多个scale合并为一个scale,scale又转化为一个conv的过程。

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

Lowering

论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第4张图片
将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

论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第5张图片

  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过程:
论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第6张图片
通过深层次的依赖关系来进行更为合理的切割:越深依赖关系越复杂,不限于一个op层,不是在一层conv上做切割,而是计算一个conv输出被下一个conv利用的相关性,做整个group的相关性的切割,用计算来代替搬运。

Calibration + Quantization

对经过优化的TOP MLIR进行多次前向推理,获取每个中间Tensor的数据,并计算它们的统计信息,通过KL方法得到初步的阈值,然后使用误差/余弦相似度方式再微调阈值,是的INT8的计算结果和FP32的结果尽量相似。
论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第7张图片
代码位置:

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

论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第8张图片
得到量化表:

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

论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第9张图片
将量化表导入到tpu mlir中去:

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

Correctness Check

论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第10张图片
TPU-MLIR提供对TOP和TPU Dialect的Inference。通过比较对应数据的相似性,来确定整个转化/编译过程的正确性。同时由于可以比较每个中间Tensor的结果,开发者可以快速地定位错误点,便于Debug。
论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第11张图片
代码位置:

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

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

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

论文讲解——TPU-MLIR: A Compiler For TPU Using MLIR_第12张图片

相关资料

  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。

你可能感兴趣的:(mlir)