sophgo/tpu-mlir

TPU-MLIR

面向 TPU 的开源 MLIR 机器学习编译器。

✨ 项目简介

TPU-MLIR 提供完整的工具链，将主流框架下预训练的神经网络转换为可在 SOPHGO TPU 上高效运行的 bmodel 文件。基于 MLIR 构建，提供统一 IR、清晰的下降流水线，以及完善的量化、校准与部署工具。

┌──────────────────┐    model_transform.py    ┌──────────┐    model_deploy.py    ┌──────────┐
│ ONNX / PyTorch / │ ───────────────────────► │   MLIR   │ ────────────────────► │  bmodel  │
│  TFLite / Caffe  │      (前端导入)          │  (Top →  │  (lowering、量化、    │  TPU 部署 │
│   HuggingFace    │                          │   Tpu)   │   layer-group …)      │          │
└──────────────────┘                          └──────────┘                       └──────────┘

🚀 主要特性

多框架前端 —— PyTorch、ONNX、TFLite、Caffe（其他框架请先导出为 ONNX）。
LLM 一键转换 —— 通过 llm_convert.py 直接编译 HuggingFace 大模型（Qwen、MiniCPM-V 等）。
完整量化工具链 —— F32 / BF16 / F16 / INT8（对称 & 非对称），兼容 AWQ / GPTQ / AutoRound，支持校准与 QAT。
基于 MLIR 的流水线 —— 清晰的 Top / Tpu 方言、模式重写、layer-group 内存规划。
完善的辅助工具 —— model_runner、model_tool、精度验证、模拟器、可视化工具。
中英双语文档与活跃社区 —— 提供论文、技术手册与系列视频教程。

🔧 安装

TPU-MLIR 在指定的 Docker 镜像内运行。容器启动后，既可以直接安装预编译 wheel，也可以从源码编译。

1. 拉取 Docker 镜像

docker pull sophgo/tpuc_dev:latest

如果拉取失败，可手动下载并加载镜像包：

wget https://sophon-assets.sophon.cn/sophon-prod-s3/drive/25/04/15/16/tpuc_dev_v3.4.tar.gz
docker load -i tpuc_dev_v3.4.tar.gz

创建并进入容器：

docker run --privileged --name tpu-mlir -v $PWD:/workspace -it sophgo/tpuc_dev:latest

2a. 安装预编译 wheel

要求 Python ≥ 3.10，Ubuntu 22.04（推荐直接使用上述 Docker 镜像）。

pip install tpu_mlir

2b. 从源码编译（推荐）

# git clone with --depth 1
cd /workspace/tpu-mlir
pip install -r requirements.txt
source ./envsetup.sh
./build.sh

🤖 快速上手 — LLM (Qwen)

将 HuggingFace 上的 LLM 编译为 TPU 上的 bmodel。

点击展开完整 LLM 流程

1. 下载模型

推荐使用预量化 (AWQ / GPTQ / AutoRound) 版本：

git lfs install
git clone https://huggingface.co/Intel/Qwen3.5-2B-int4-AutoRound

2. 编译为 bmodel

# 如遇 transformers / torch 版本问题：
#   pip3 install transformers torchvision -U
# --max_input_length 设置最大 prefill 长度，缺省时取 -s。
llm_convert.py \
  -m /workspace/Qwen3.5-2B-int4-AutoRound \
  --max_input_length 1024 \
  -s 2048 \
  -c bm1684x \
  --max_pixels 768,768 \
  -o qwen3.5_2b

llm_convert.py 主要参数：

参数	简写	必选	说明
`model_path`	`m`	✅	模型权重路径
`seq_length`	`s`	✅	最大序列长度
`max_input_length`	—	—	最大单次输入长度，缺省时取 `seq_length`（`-s`）
`quantize`	`q`	✅	量化方式：`auto` / `w4bf16` / `w4f16` / `bf16` / `f16` 等（如源模型已量化可省略）
`q_group_size`	`g`	—	量化分组大小（默认 64）
`chip`	`c`	✅	目标芯片：`bm1684x` / `bm1688` / `cv186ah`
`max_pixels`	—	—	多模态最大分辨率 `width,height`，按 `model_type` 取默认值（qwen2_5_vl: `672,896`；minicpmv: `980,980`；其他：`768,768`）
`out_dir`	`o`	✅	输出目录

3. 在 PCIe / SoC 环境运行

将 python_demo 拷贝到设备上并编译：

mkdir build && cd build
cmake ..
make
cp *cpython*.so ..
cd ..

执行 bmodel：

python3 pipeline.py -m xxxx.bmodel -c config

运行示例：

Qwen demo

🖼️ 快速上手 — 视觉 (YOLOv5)

以 yolov5s.onnx 为例，在 BM1684X TPU 上完成编译与部署。模型已置于 regression/model/yolov5s.onnx。

点击展开完整 YOLOv5 流程

1. 准备工作目录

mkdir model_yolov5s && cd model_yolov5s
cp ${REGRESSION_PATH}/model/yolov5s.onnx .
cp -rf ${REGRESSION_PATH}/dataset/COCO2017 .
cp -rf ${REGRESSION_PATH}/image .
mkdir workspace && cd workspace

2. 模型转 MLIR

如果模型以图像为输入，需指定预处理；输入为 npz 文件时则无需预处理。预处理公式为：

$y = (x - \text{mean}) \times \text{scale}$

YOLOv5 官方输入为 RGB，每像素乘 1/255，因此 mean = 0,0,0、scale = 0.0039216,0.0039216,0.0039216。

model_transform.py \
  --model_name yolov5s \
  --model_def ../yolov5s.onnx \
  --input_shapes [[1,3,640,640]] \
  --mean 0.0,0.0,0.0 \
  --scale 0.0039216,0.0039216,0.0039216 \
  --keep_aspect_ratio \
  --pixel_format rgb \
  --output_names 350,498,646 \
  --test_input ../image/dog.jpg \
  --test_result yolov5s_top_outputs.npz \
  --mlir yolov5s.mlir

model_transform.py 主要参数：

参数	必选	说明
`model_name`	✅	模型名称
`model_def`	✅	模型定义文件（`.onnx` / `.pt` / `.tflite` / `.prototxt`）
`model_data`	—	Caffe 权重文件（`.caffemodel`）
`input_shapes`	—	输入形状，例如 `[[1,3,640,640]]`，支持多输入
`resize_dims`	—	图像 resize 尺寸；不指定时按模型输入大小
`keep_aspect_ratio`	—	是否保持长宽比（不足部分补 0），默认关闭
`mean`	—	各通道均值，默认 `0,0,0`
`scale`	—	各通道缩放，默认 `1,1,1`
`pixel_format`	—	`rgb` / `bgr` / `gray` / `rgbd`
`output_names`	—	输出名；不指定时使用模型默认输出
`test_input`	—	用于验证的输入文件（图像 / npy / npz），不指定则跳过验证
`test_result`	—	验证结果保存路径
`excepts`	—	不参与验证的层名，逗号分隔
`debug`	—	保留中间文件
`mlir`	✅	输出 MLIR 文件路径

完成后会生成预处理后的 ${model_name}_in_f32.npz。

3. MLIR → F16 bmodel

model_deploy.py \
  --mlir yolov5s.mlir \
  --quantize F16 \
  --processor bm1684x \
  --test_input yolov5s_in_f32.npz \
  --test_reference yolov5s_top_outputs.npz \
  --model yolov5s_1684x_f16.bmodel

model_deploy.py 主要参数：

参数	必选	说明
`mlir`	✅	输入 MLIR 文件
`quantize`	✅	`F32` / `BF16` / `F16` / `INT8`
`processor`	✅	目标芯片
`calibration_table`	—	量化校准表（INT8 必填）
`tolerance`	—	MLIR 量化与 fp32 推理结果的最小相似度
`correctness`	—	模拟器与 MLIR 量化推理结果的最小相似度（默认 `0.99,0.90`）
`excepts`	—	不参与验证的层名，逗号分隔
`debug`	—	保留中间文件
`model`	✅	输出 bmodel 路径
`dynamic`	—	动态 codegen，支持动态 shape

4. MLIR → INT8 bmodel

先做校准（一般 100~1000 张样本），优先使用对称量化。

run_calibration.py yolov5s.mlir \
  --dataset ../COCO2017 \
  --input_num 100 \
  -o yolov5s_cali_table

model_deploy.py \
  --mlir yolov5s.mlir \
  --quantize INT8 \
  --calibration_table yolov5s_cali_table \
  --processor bm1684x \
  --test_input yolov5s_in_f32.npz \
  --test_reference yolov5s_top_outputs.npz \
  --tolerance 0.85,0.45 \
  --model yolov5s_1684x_int8.bmodel

5. 验证结果

示例脚本位于 python/samples/detect_yolov5.py：

# ONNX
detect_yolov5.py --input ../image/dog.jpg --model ../yolov5s.onnx          --output dog_origin.jpg
# F16 bmodel
detect_yolov5.py --input ../image/dog.jpg --model yolov5s_1684x_f16.bmodel --output dog_f16.jpg
# INT8 bmodel
detect_yolov5.py --input ../image/dog.jpg --model yolov5s_1684x_int8.bmodel --output dog_int8.jpg

不同模型的输出对比：

YOLOv5 results

🛠️ 辅助工具

`model_runner.py` —— 通用推理工具

支持 bmodel / mlir / PyTorch / ONNX / TFLite / Caffe。

model_runner.py \
  --input  resnet18_in_f32.npz \
  --model  resnet18_1684x_f32.bmodel \
  --output resnet18_output.npz

`model_tool` —— bmodel 查看与编辑

model_tool
  --info     model_file                                 : 查看简要信息
  --print    model_file                                 : 查看详细信息
  --extract  model_file                                 : 将多网络 bmodel 拆分为多个单网络 bmodel
  --combine  file1 .. fileN -o new_file                 : 按文件路径合并 bmodel
  --combine_dir dir1 .. dirN -o new_dir                 : 按目录合并 bmodel
  --dump     model_file start_offset byte_size out_file : 从 bmodel dump 二进制数据

model_tool --info resnet18_1684x_f32.bmodel

📖 资源

文档与论文

类型	链接
论文	TPU-MLIR (arXiv 2210.15016)
手册	技术参考手册
入门	快速入门

分享会

视频教程

点击展开视频目录

#	主题	链接
01	什么是深度学习编译器？	入门
02	MLIR 入门	语法 1 · 语法 2 · 语法 3 · Dialect Conversion · Pattern Rewriting
03	TPU-MLIR 概览	总览 · 前端转换 · Lowering
04	量化	总览 · 公式推导 · 校准 · QAT
05	TPU 内存	Ep1 · Ep2
06	TPU-MLIR 实战	转 ONNX · 图优化 · 算子支持 · 模型支持 · 融合预处理 · 精度验证

📝 引用

如果 TPU-MLIR 对您的研究有帮助，请按下列格式引用：

@misc{tpumlir2022,
  title         = {TPU-MLIR: A Compiler For TPU Using MLIR},
  author        = {HuPengchao and LuMan and WangLei and JiangGuoyue},
  year          = {2022},
  eprint        = {2210.15016},
  archivePrefix = {arXiv},
  primaryClass  = {cs.PL}
}

🤝 贡献

欢迎提交 Bug、功能请求与 Pull Request！

提交前请先在 Issues 中检索，避免重复。
较大改动建议先开 Issue 讨论设计方案。
PR 前请运行 regression/ 下的回归测试。

📄 许可证

本项目采用根目录下 LICENSE 文件中规定的开源许可证。