sophgo/tpu-mlir

Machine learning compiler based on MLIR for Sophgo TPU.

What it solves

TPU-MLIR solves the problem of deploying pre-trained neural networks from various mainstream frameworks onto TPU hardware. It bridges the gap between high-level model definitions (like PyTorch or ONNX) and the specific binary format (bmodel) required for efficient execution on TPU chips.

How it works

The project uses a compilation pipeline based on MLIR (Multi-Level Intermediate Representation) to transform models through several stages:

  1. Front-end Import: Models from PyTorch, ONNX, TFLite, Caffe, or HuggingFace are imported and converted into a unified MLIR representation.
  2. Lowering: The compiler transforms the high-level representation into TPU-specific dialects (Top and Tpu) using pattern rewrites and memory planning.
  3. Quantization: It applies various quantization techniques (F32, BF16, F16, INT8) and calibration to optimize the model for hardware efficiency.
  4. Deployment: The final output is a bmodel file that can be run on the target TPU processor.

Who it’s for

It is designed for AI engineers and developers who need to deploy deep learning models (including Large Language Models and vision models) onto Sophgo TPU hardware.

Highlights

  • Broad Framework Support: Compatible with PyTorch, ONNX, TFLite, Caffe, and HuggingFace.
  • LLM Optimization: Dedicated one-shot conversion for HuggingFace LLMs with support for history KV cache and chunked inference.
  • Comprehensive Quantization: Supports symmetric/asymmetric INT8, F16, BF16, and passthrough for AWQ, GPTQ, and AutoRound.
  • Production Tooling: Includes a universal inference runner (model_runner), a bmodel inspection tool (model_tool), and accuracy validation utilities.