使用MLIR进行矩阵-矩阵乘法的高性能GPU代码生成：一些早期结果
分布式、并行与集群计算
2021-08-31 v1 性能
摘要
本报告展示了使用MLIR编译器基础设施针对NVIDIA GPU上张量核进行代码生成的一些早期结果。当今高性能深度学习的最先进水平主要由手动优化的高度调优库驱动。开发此类库的方法通常不如LLVM等编译器基础设施那样模块化或可复用。手动优化通常不使用标准的中间表示（IR），尽管所执行的优化可以编码为IR上的一系列变换步骤和自定义遍。手工调优也可能错过仅能通过自动代码生成轻易到达的设计点探索。我们认为，直到近期引入MLIR（多级中间表示），IR基础设施才得以有效应对自动生成领域特定库的问题。特别是，使用单一IR在高、中、低层级表示和变换计算抽象此前十分困难。借助MLIR中合适的抽象，我们构建了一个实验性的 lowering 流水线，能够自动生成面向NVIDIA GPU张量核的矩阵-矩阵乘法代码。在我们评估的一组问题规模上，初始性能结果表明，在基于NVIDIA Ampere微架构的Geforce 3090 RTX上，对于FP32和FP16累加，我们分别达到了CuBLAS性能的95-119%和80-160%。我们相信这些结果可作为使用IR基础设施针对类似专用加速器进行自动代码与库生成进一步研究与开发的动力。
引用
@article{arxiv.2108.13191,
title = {High Performance GPU Code Generation for Matrix-Matrix Multiplication using MLIR: Some Early Results},
author = {Navdeep Katel and Vivek Khandelwal and Uday Bondhugula},
journal= {arXiv preprint arXiv:2108.13191},
year = {2021}
}