tpu-mlir

TPU-MLIR

An open-source, MLIR-based machine-learning compiler for TPUs.

English · 简体中文 · Quick Start · Docs · Issues

✨ Overview

TPU-MLIR converts pre-trained neural networks from mainstream frameworks into bmodel files that run efficiently on TPUs. Built on top of MLIR, it provides a unified IR, a clean lowering pipeline, and a rich set of tools for quantization, calibration, and deployment.

🚀 Highlights

• Multi-framework front-ends — PyTorch, ONNX, TFLite, Caffe (other frameworks via ONNX).
• LLM-ready — one-shot conversion of HuggingFace LLMs (Qwen, MiniCPM-V, …) via llm_convert.py.
• Full quantization toolchain — F32 / BF16 / F16 / INT8 (symmetric & asymmetric), AWQ / GPTQ / AutoRound passthrough, calibration, QAT.
• MLIR-based pipeline — clean dialects (Top / Tpu), pattern rewrites, layer-group memory planning.
• Production tooling — model_runner, model_tool, accuracy validation, simulator, visualizer.
• Bilingual docs & active community — English / 中文 manuals, papers, and video tutorials.

📚 Table of Contents

• Installation
• Quick Start (LLM — Qwen)
• Quick Start (Vision — YOLOv5)
• Auxiliary Tools
• Resources
• Citation
• Contributing
• License

🔧 Installation

TPU-MLIR runs inside a prebuilt Docker image. After the container is running you can either install the Python wheel or build from source.

1. Pull the Docker image

docker pull sophgo/tpuc_dev:latest

If the pull fails, download the tarball and load it manually:

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

Create and enter the container:

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

2a. Install the prebuilt wheel

Requires Python ≥ 3.10 on Ubuntu 22.04 (already satisfied inside the Docker image).

pip install tpu_mlir

2b. Build from source (recommended)

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

🤖 Quick Start — LLM (Qwen)

Convert and run a HuggingFace LLM (here: Qwen) on a TPU.

Click to expand the full LLM walkthrough

1. Download the model

A pre-quantized AWQ / GPTQ / AutoRound build is recommended.

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

2. Compile to bmodel

LLM compilation falls into two scenarios, controlled by --use_history_kv.

Without history KV — compiles two instruction groups: block_ (prefill) and block_cache_ (decode). Recommended for single-turn conversations with short context (e.g. within 4K):

# -s sets the max total sequence length;
# --max_input_length sets the max input length (defaults to -s if omitted).
llm_convert.py \
  -m /workspace/Qwen3.5/Qwen3.5-2B-int4-AutoRound \
  -s 2048 \
  --max_input_length 1024 \
  -c bm1684x \
  -o qwen3.5_4b

With history KV — compiles three instruction groups: block_ (prefill), block_kv_ (prefill with history), and block_cache_ (decode). Recommended for multi-turn conversations, long contexts (e.g. 8K), or when unsure — it is more flexible with good overall performance:

llm_convert.py \
  -m /workspace/Qwen3.5/Qwen3.5-2B-int4-AutoRound \
  -s 8192 \
  --use_history_kv \
  --chunk_length 1024 \
  -c bm1684x \
  -o qwen3.5_4b_history

--chunk_length sets the segment length for chunked inference: with 1K chunks, a 7K input runs prefill as block_ + 7 × block_kv_; decode is also segmented by KV-cache length, so performance differs at 1K / 2K / 4K / 8K lengths.

Main arguments of llm_convert.py:

3. Run on PCIe / SoC

Copy the cpp_demo folder onto your device and build it:

mkdir build && cd build
cmake ..
make
mv pipeline ..
cd ..

Then run the bmodel:

./pipeline -m xxxx.bmodel -c config

Sample output:

The demo accepts slash commands (e.g. /exit, /clear) and uses @ to attach files — what is the image about? @./test.jpg for images, or what is it talking about? @./story.txt for text files (.txt / .md).

🖼️ Quick Start — Vision (YOLOv5)

Compile and run yolov5s.onnx on the BM1684X TPU. The model is bundled in regression/model/yolov5s.onnx.

Click to expand the full YOLOv5 walkthrough

1. Prepare the working directory

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. Convert the model to MLIR

If the model takes images as input, the preprocessing must be specified. The preprocessing formula is:

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

YOLOv5's official input is RGB scaled by 1/255, so mean = 0,0,0 and 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

Main arguments of model_transform.py:

A ${model_name}_in_f32.npz file containing the preprocessed input is generated after this step.

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

Main arguments of model_deploy.py:

4. MLIR → INT8 bmodel

Run calibration first (typically 100–1000 images). Prefer symmetric quantization unless accuracy demands asymmetric.

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. Verify the results

The sample script lives at 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

Comparison of outputs:

🛠️ Auxiliary Tools

model_runner.py — universal inference runner

Supports 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 — inspect & edit bmodel

model_tool
  --info     model_file                                : show brief model info
  --print    model_file                                : show detailed model info
  --extract  model_file                                : split a multi-net bmodel into single-net bmodels
  --combine  file1 .. fileN -o new_file                : merge bmodels by file path
  --combine_dir dir1 .. dirN -o new_dir                : merge bmodels by directory
  --dump     model_file start_offset byte_size out_file: dump raw bytes from a bmodel

model_tool --info resnet18_1684x_f32.bmodel

📖 Resources

Documentation & Papers

Talks

• TPU-MLIR Paper Walkthrough
• LayerGroup

Video tutorials

Click to expand video index

📝 Citation

If TPU-MLIR helps your research, please cite:

@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}
}

@misc{llmtpu2026,
  title         = {An MLIR-Based Compilation Method for Large Language Models},
  author        = {HuPengchao and XinZhibin and ChenYifan and ZhouYangyang and WangLiang and ZhangXin},
  year          = {2026},
  eprint        = {2607.15865},
  archivePrefix = {arXiv},
  primaryClass  = {cs.PL}
}

🤝 Contributing

Bug reports, feature requests and pull requests are welcome! Before you start:

1. Search existing issues to avoid duplicates.
2. For non-trivial changes, open an issue first to discuss the design.
3. Run the regression tests under regression/ before sending a PR.

📄 License

This project is licensed under the terms of the LICENSE file in the root of this repository.