[论文解读] High Performance GPU Code Generation for Matrix-Matrix Multiplication using MLIR: Some Early Results

Navdeep Katel, Vivek Khandelwal+1 more|arXiv (Cornell University)|Aug 23, 2021

Parallel Computing and Optimization Techniques参考文献 1被引用 6

一句话总结

本文提出了一种基于MLIR的自动代码生成流水线，用于在NVIDIA A100和RTX 3090 GPU的张量核心上实现高性能矩阵-矩阵乘法。通过在优化的方言转换下降低高层IR，该方法在FP32工作负载上达到CuBLAS性能的95–119%，在FP16混合精度工作负载上达到80–160%，表明基于IR的编译可与手工优化的库相媲美，且无需手动编码。

ABSTRACT

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.

研究动机与目标

• 开发一种基于MLIR的模块化、自动代码生成流水线，用于在NVIDIA GPU上实现矩阵-矩阵乘法。
• 在无需手动内核调优的情况下高效利用张量核心，减少对专家优化库的依赖。
• 评估基于IR的编译是否能够实现与手工调优库（如CuBLAS）相当的性能。
• 通过在模块化IR流水线中系统性地进行A/B测试，探究单个优化的影响。
• 建立一个统一中间表示基础，用于自动生成领域专用、高性能库的框架。

提出的方法

• 该方法使用MLIR的多级IR表示矩阵乘法在高层、中层和低层抽象级别，支持系统性变换。
• 引入一种新的MLIR方言用于Warp Matrix Multiply Accumulate（WMMA）操作，并将其降低为LLVM/NVPTX IR。
• 流水线应用一系列方言降低和优化过程，包括分块、共享内存分块以及寄存器级别优化。
• 逐步应用如流水化全局内存加载和分块大小选择等优化，以研究其对性能的影响。
• 代码生成流水线端到端运行，从高层矩阵乘法操作开始，编译为针对张量核心优化的高效CUDA代码。
• 在NVIDIA的Ampere架构RTX 3090 GPU上，使用标准问题规模和混合/半精度配置评估性能。

实验结果

研究问题

• RQ1 基于IR的代码生成流水线是否能在张量核心上实现与CuBLAS等手工优化库相竞争的性能，用于矩阵-矩阵乘法？
• RQ2 在模块化、基于IR的流水线中，分块大小和内存访问流水化等单个优化如何影响性能？
• RQ3 在不同精度格式下，自动代码生成在多大程度上可以匹配或超越厂商优化库的性能？
• RQ4 像MLIR这样的统一IR基础设施是否能有效表达并转换针对专用加速器（如张量核心）的高层计算抽象？
• RQ5 在不同问题规模和数据类型下，自动生成内核的性能与CuBLAS相比如何？

主要发现

• 在RTX 3090上，自动生成的内核在FP32混合精度矩阵乘法中达到CuBLAS性能的95–119%，峰值效率达到设备理论峰值的95.4%。
• 在FP16混合精度矩阵乘法中，性能始终维持在CuBLAS的95–119%之间，且在小规模问题上表现优于CuBLAS，这是由于更小的线程块分块带来了更高的占用率。
• 在FP16纯精度矩阵乘法中，生成代码达到CuBLAS性能的80–160%，在大矩阵规模上表现更优，因为CuBLAS在此类规模下表现出性能不一致和次优的分块选择。
• 研究发现，CuBLAS使用更小的分块大小（如128×128×32）和五级流水线，但存在更高的全局内存停顿，表明其延迟隐藏效果较差。
• 较小的分块大小（如64×64×64）在小规模问题上提升性能，因占用率更高；而更大的分块在大矩阵上更有效，因数据重用更优。
• 增量优化研究证实，流水化和分块大小选择显著影响性能，最终优化流水线实现了接近峰值性能。