README.md

July 31, 2026 · View on GitHub

TPU-MLIR

TPU-MLIR

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

PyPI Python Stars Issues arXiv

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.

TPU-MLIR compilation workflow

🚀 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 toolingmodel_runner, model_tool, accuracy validation, simulator, visualizer.
  • Bilingual docs & active community — English / 中文 manuals, papers, and video tutorials.

📚 Table of Contents


🔧 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
# 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 \times $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:

ParameterShortRequiredDescription
model_pathmPath to the model weights
seq_lengthsMaximum sequence length
max_input_lengthMaximum input length; defaults to seq_length (-s) when omitted
use_history_kvEnable history KV cache (adds the block_kv_ prefill-with-history instruction group)
chunk_lengthSegment length for chunked inference; only meaningful with use_history_kv
chipcTarget platform: bm1684x / bm1688 / cv186ah
out_diroOutput directory
dynamicDynamic-shape compilation; recommended for all new conversions (Qwen3.5 forces it)
do_sampleEnable random sampling
max_pixelsMax image size for VLMs; leave unset to use the per-model defaults
embedding_diskStore word embeddings in a .bin file and run them on CPU (saves on-chip memory)
lora_max_rankMax LoRA rank; setting it compiles a LoRA-enabled version

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:

Qwen demo

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=(xmean)×scaley = (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:

ArgumentRequiredDescription
model_nameModel name
model_defModel definition file (.onnx, .pt, .tflite, .prototxt)
model_dataCaffe weight file (.caffemodel)
input_shapesInput shape, e.g. [[1,3,640,640]] — supports multiple inputs
resize_dimsImage resize size before feeding into the model
keep_aspect_ratioKeep aspect ratio (pads with 0). Off by default
meanPer-channel mean (default 0,0,0)
scalePer-channel scale (default 1,1,1)
pixel_formatrgb / bgr / gray / rgbd
output_namesOutput tensor names. Defaults to model outputs
test_inputValidation input (image / npy / npz). Skipped if not specified
test_resultOutput file for validation
exceptsComma-separated list of layers excluded from validation
debugKeep intermediate files
mlirOutput MLIR file path

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:

ArgumentRequiredDescription
mlirInput MLIR file
quantizeF32 / BF16 / F16 / INT8
processorTarget chip
calibration_tableCalibration table (required for INT8)
toleranceMin similarity between MLIR-quantized and MLIR-fp32 inference
correctnessMin similarity between simulator and MLIR-quantized inference (default 0.99,0.90)
exceptsComma-separated layers excluded from validation
debugKeep intermediate files
modelOutput bmodel path
dynamicDynamic codegen for dynamic shapes

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:

YOLOv5 results


🛠️ 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

TypeLink
PaperTPU-MLIR (arXiv 2210.15016)
PaperAn MLIR-Based Compilation Method for Large Language Models (arXiv 2607.15865)
ManualTechnical Reference Manual
GuideQuick Start

Talks

Video tutorials

Click to expand video index
#TopicLinks
01What is a Deep Learning Compiler?Intro
02MLIR IntroSyntax 1 · Syntax 2 · Syntax 3 · Dialect Conversion · Pattern Rewriting
03TPU-MLIR IntroOverview · Front-end · Lowering
04QuantizationOverview · Formula · Calibration · QAT
05TPU MemoryEp1 · Ep2
06TPU-MLIR PracticeTo ONNX · Graph Optimization · Operator Support · Model Support · Fuse Preprocess · Accuracy Validation

📝 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.