# TPU-MLIR: A Compiler For TPU Using MLIR

**Authors:** Pengchao Hu, Man Lu, Lei Wang, Guoyue Jiang  
**arXiv:** [2210.15016](https://arxiv.org/abs/2210.15016)  
**Date:** 2022-10-23  
**Summary page:** https://www.summarizepaper.com/en/paper/2210.15016v1/

## Key points

- Developing domain-specific compilers can be costly and time-consuming
- Multi-level intermediate representations (MLIR) offer a reusable and extensible compiler infrastructure
- TPU-MLIR is an end-to-end compiler based on MLIR that deploys pre-trained neural network models to a custom ASIC called a Tensor Processing Unit (TPU)
- TPU-MLIR defines two new dialects: Tensor operation (TOP) dialect and TPU kernel dialect
- NN model is translated into the TOP dialect and then lowered to the TPU dialect for different TPUs according to their chip configuration
- MLIR pass pipeline is used to organize and perform optimization on TPU to generate machine code
- Verification procedure ensures correctness of each transform stage
- TPU-MLIR offers an efficient solution for designing AI compilers for ASIC chips using MLIR, reducing costs and improving performance.

## Summary

The development of domain-specific compilers can be a costly and time-consuming process. Multi-level intermediate representations (MLIR) offer a promising solution by providing a reusable and extensible compiler infrastructure. In this paper, the authors present TPU-MLIR, an end-to-end compiler based on MLIR that deploys pre-trained neural network (NN) models to a custom ASIC called a Tensor Processing Unit (TPU). To implement its functionality, TPU-MLIR defines two new dialects: 1. a Tensor operation (TOP) dialect that encodes the deep learning graph semantics and is independent of the deep learning framework; and 2. a TPU kernel dialect to provide standard kernel computation on TPU. The NN model is translated into the TOP dialect and then lowered to the TPU dialect for different TPUs according to their chip configuration. The authors demonstrate how to use the MLIR pass pipeline to organize and perform optimization on TPU to generate machine code. They also present a verification procedure to ensure the correctness of each transform stage. Overall, TPU-MLIR offers an efficient solution for designing AI compilers for ASIC chips using MLIR. With its ability to reduce costs and improve performance, it has significant potential in advancing the field of deep learning hardware acceleration.