tpu-mlir

把 PyTorch/ONNX 模型一键编译到算能 TPU 上跑起来

开源 🏷️完全免费 🔧 仅开源项目，无在线服务，需自行部署基础设施 🍴 复刻 230 🐛 待处理 114 👀 关注 25

GitHub ↗

GitHub 星标 ★ 1016

维护状态活跃

是否开源开源

定价模式 🏷️完全免费开源免费，无在线服务

项目预览

项目数据

分类基础设施

开发团队sophgo

GitHubhttps://github.com/sophgo/tpu-mlir

定价模式🏷️完全免费

价格说明开源免费，无在线服务

访问状态🔧 仅开源项目，无在线服务，需自行部署

是否开源是

开源协议NOASSERTION

主要语言C++

技术栈/模型compiler,llm,mlir,tpu

GitHub 星标★ 1016

上线时间2022-06-03 00:00:00

代码更新2026-11-03

仓库体积5.89 GB

本站数据更新2026-11-03

维护状态活跃

收录时间2026-09-16

浏览次数1

项目介绍

tpu-mlir 是算能（Sophgo）开源的机器学习编译器，基于 MLIR 框架构建，专门面向 Sophgo TPU 硬件做模型编译与优化。它解决的问题是：主流深度学习框架（PyTorch、ONNX、TFLite 等）训练出的模型无法直接在 Sophgo TPU 上高效运行，需要一套从高层算子到 TPU 指令的完整转换链路。核心能力包括多框架模型导入与图优化、算子融合与量化（支持 INT8 等低精度）、内存与调度规划，以及最终生成可在 TPU 上执行的二进制。项目用 C++ 实现，围绕 MLIR 的 dialect 体系组织编译流程，同时覆盖 LLM 相关算子，适合边缘 AI 推理场景。对使用算能 BM1684X、CV186 等芯片做部署的团队来说，它是打通模型到硬件的关键工具链组件。

Star 增长趋势

本站每日记录该项目 GitHub 星标变化，观测越久曲线越完整

数据来源：GitHub API 逐日快照 · 收录于 2026-09-16 · 当前星标 ★ 1016

技术标签

compiler llm

使用教程

难度：进阶约 20 分钟部署方式：Docker 7 步

环境要求

Docker 环境（用于拉取 sophgo/tpuc_dev 镜像）
Python ≥ 3.10，Ubuntu 22.04（Docker 镜像内已满足）
一台 Sophgo TPU 设备（如 BM1684X，PCIe/SoC）用于实际运行 bmodel
编译 cpp_demo 需要 cmake 与 make

安装与启动步骤

1拉取开发镜像

从 Docker Hub 拉取官方预编译镜像，后续编译与运行都在该容器内进行。
```
docker pull sophgo/tpuc_dev:latest
```

2离线加载镜像

若拉取失败，下载镜像 tar 包后用 docker load 导入，效果与拉取一致。

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

3启动容器

以特权模式创建名为 tpu-mlir 的容器，并把当前目录挂载到 /workspace。
```
docker run --privileged --name tpu-mlir -v $PWD:/workspace -it sophgo/tpuc_dev:latest
```
4安装 Python 包

在容器内安装预编译 wheel，即可使用 tpu_mlir 相关工具与 Python 接口。
```
pip install tpu_mlir
```
5编译推理 demo

把 cpp_demo 拷到设备上，用 cmake + make 编译出 pipeline 可执行文件。
```
mkdir build && cd build
cmake ..
make
mv pipeline ..
cd ..
```
6运行 bmodel

用编译好的 pipeline 加载 bmodel 与配置，在 TPU 设备上执行推理。
```
./pipeline -m xxxx.bmodel -c config
```
7通用推理脚本

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
```

关键配置

配置项	必填	说明	示例
`-m`	是	pipeline 运行时指定的 bmodel 模型文件路径	`xxxx.bmodel`
`-c`	是	pipeline 运行时指定的配置文件（目录）	`config`
`--model`	是	model_runner.py 的模型文件，支持 bmodel/mlir/onnx 等	`resnet18_1684x_f32.bmodel`
`--input`	是	model_runner.py 的输入数据文件（npz）	`resnet18_in_f32.npz`
`--output`	是	model_runner.py 的输出结果保存文件（npz）	`resnet18_output.npz`

如何确认成功

容器内 pip show tpu_mlir 能看到版本号；执行 ./pipeline 后终端打印出推理结果（README 的 Sample output）。

常见问题

Q：docker pull 拉取镜像失败怎么办？

A：改用离线方式：wget 下载 tpuc_dev_v3.4.tar.gz，再执行 docker load -i tpuc_dev_v3.4.tar.gz 导入本地镜像。

Q：对 Python 版本有要求吗？

A：需要 Python ≥ 3.10、Ubuntu 22.04；官方 Docker 镜像内已满足，无需额外配置。

Q：示例模型在哪里获取？

A：YOLOv5 示例 yolov5s.onnx 已内置在仓库的 regression/model/ 目录下，也可从 ultralytics 发布页下载。

Q：model_runner.py 支持哪些模型格式？

A：支持 bmodel、mlir、PyTorch、ONNX、TFLite、Caffe 六种输入。

Q：为什么必须加 --privileged？

A：README 的容器启动命令使用了特权模式，用于访问 TPU 等设备资源，去掉可能导致设备不可用。

注意事项

容器启动需加 --privileged，否则可能无法访问 TPU 设备。
用 -v $PWD:/workspace 挂载当前目录，编译产物与模型放在该目录下便于宿主机访问。
bmodel 需在真实 Sophgo TPU 设备（PCIe/SoC）上运行，编译阶段可在容器内完成。
README 未给出完整的模型编译（转 bmodel）命令，具体流程请参考官方 Quick Start 文档 https://tpumlir.org/quick_start_en/index.html。