论文标题

TPU-MLIR:使用MLIR的TPU编译器

TPU-MLIR: A Compiler For TPU Using MLIR

论文作者

Hu, Pengchao, Lu, Man, Wang, Lei, Jiang, Guoyue

论文摘要

多层中间表示(MLIR)通过提供可重复使用且可扩展的编译器基础架构来降低构建域特异性编译器的成本的巨大希望。这项工作提出了TPU-MLIR,这是一种基于MLIR的端到端编译器,该编译器将预训练的神经网络(NN)模型部署到称为张量处理单元(TPU)的自定义ASIC。 TPU-MLIR定义了两个新方言以实现其功能:1。编码深度学习图语义的张量操作(顶部)方言,独立于深度学习框架和2。TPU内核方言,以提供TPU上的标准内核计算。根据芯片的配置,将NN模型转换为顶部方言,然后将其降低到TPU方言。我们演示了如何使用MLIR PASS管道来组织和执行TPU的优化以生成机器代码。本文还提出了一个验证程序,以确保每个变换阶段的正确性。

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.

扫码加入交流群

加入微信交流群

微信交流群二维码

扫码加入学术交流群,获取更多资源