arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2108.13191cs.DCcs.PF
收藏

基于MLIR的矩阵-矩阵乘法的高性能GPU代码生成:一些早期结果

High Performance GPU Code Generation for Matrix-Matrix Multiplication using MLIR: Some Early Results

Navdeep Katel, Vivek Khandelwal, Uday Bondhugula

摘要:

本报告介绍了使用MLIR编译器基础设施针对NVIDIA GPU上的张量核心进行代码生成的一些早期结果。当今最先进的高性能深度学习主要由手动优化的高度调优库驱动。开发这种库的方法通常不是模块化的,也不是像LLVM这样的编译器基础设施那样可重用的。手动优化通常不使用标准的中间表示(IR),尽管执行的优化可以被编码为IR上的一系列转换步骤和定制传递。手动调优还可能错过仅可通过自动代码生成轻松实现的设计点的探索。我们认为,直到最近推出的MLIR(多级中间表示),IR基础设施不适合以有效的方式解决自动生成特定领域的库的问题。特别是,使用单个IR很难在高、中、低级别上表示和转换计算抽象。通过MLIR中的适当抽象,我们构建了一个实验性的降低流水线,能够在NVIDIA GPU上自动生成针对其张量内核的矩阵-矩阵乘法代码。在我们评估的一组问题大小上,初始性能结果表明,在NVIDIA基于Ampere微架构的Geforce 3090 RTX上,FP 32和FP 16的性能分别达到CuBLAS的95-119%和80-160%。我们相信,这些结果可以作为进一步研究和开发自动代码和库生成使用类似的专业加速器的IR基础设施的动机。

英文摘要:

This report presents some early results on code generation targeting tensor cores on NVIDIA GPUs using the MLIR compiler infrastructure. The state-of-the-art in high-performance deep learning today is primarily driven by manually optimized highly tuned libraries. The approach to develop such libraries is often not modular or reusable to the same extent that compiler infrastructure like LLVM is. Manual optimization typically does not use a standard intermediate representation (IR), although the optimizations performed can be encoded as a sequence of transformation steps and customized passes on an IR. Hand tuning may also miss exploration of design points only reachable easily by automatic code generation. We believe that until the recent introduction of MLIR (Multi-level intermediate representation), IR infrastructure was not geared to tackle the problem of automatic generation of domain-specific libraries in an effective manner. In particular, it was hard to represent and transform compute abstractions at high, middle, and low levels using a single IR. With suitable abstractions in MLIR, we build an experimental lowering pipeline that is able to automatically generate code for matrix-matrix multiplication on NVIDIA GPUs targeting its tensor cores. On a set of problem sizes we evaluated, initial performance results show that we are able to attain performance that is 95-119% and 80-160% of CuBLAS for FP32 and FP16 accumulate respectively on NVIDIA's Ampere microarchitecture-based Geforce 3090 RTX. We believe that these results could be used as motivation for further research and development on automatic code and library generation using IR infrastructure for similar specialized accelerators.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑