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

TL;DR

This paper demonstrates early results in automatically generating high-performance GPU code for matrix multiplication using MLIR, achieving near-optimized performance on NVIDIA tensor cores and highlighting MLIR's potential for domain-specific library development.

Contribution

The paper introduces an MLIR-based pipeline for automatic code generation targeting NVIDIA GPU tensor cores, showing promising performance results and advancing compiler infrastructure for deep learning libraries.

Findings

Achieved 95-119% of CuBLAS performance for FP32 on RTX 3090.

Achieved 80-160% of CuBLAS performance for FP16 on RTX 3090.

Demonstrated MLIR's potential for automatic, high-performance code generation.

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.