Multi-level intermediate representations (MLIR) show great promise for reducing the cost of building domain-specific compilers by providing a reusable and extensible compiler infrastructure. This work presents 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). TPU-MLIR defines two new dialects to implement its functionality: 1. a Tensor operation (TOP) dialect that encodes the deep learning graph semantics and independent of the deep learning framework and 2. a TPU kernel dialect to provide a standard kernel computation on TPU. A NN model is translated to the TOP dialect and then lowered to the TPU dialect for different TPUs according to the chip's configuration. We demonstrate how to use the MLIR pass pipeline to organize and perform optimization on TPU to generate machine code. The paper also presents a verification procedure to ensure the correctness of each transform stage.

翻译：多级中间表示（MLIR）通过提供可复用且可扩展的编译器基础设施，在降低构建领域专用编译器的成本方面展现出巨大潜力。本工作提出TPU-MLIR——一种基于MLIR的端到端编译器，用于将预训练神经网络（NN）模型部署到名为张量处理单元（TPU）的定制专用集成电路（ASIC）上。TPU-MLIR定义了两种新方言（dialect）以实现其功能：1. 张量操作（TOP）方言，用于编码深度学习图语义且独立于深度学习框架；2. TPU内核方言，用于提供TPU上的标准内核计算。神经网络模型首先被转换为TOP方言，随后根据芯片配置降级（lower）至面向不同TPU的TPU方言。我们展示了如何使用MLIR pass管线在TPU上组织并执行优化以生成机器码。本文还提出了一种验证流程，以确保每个变换阶段的正确性。