[论文解读] High Performance Code Generation in MLIR: An Early Case Study with GEMM

Uday Bondhugula|arXiv (Cornell University)|Mar 1, 2020

Algorithms and Data Compression参考文献 6被引用 27

一句话总结

本文展示了一项早期案例研究，表明MLIR（一种具备仿射和多面体抽象的编译器基础设施）能够自动生成与手工优化库（如OpenBLAS和MKL）性能相当的高性能GEMM内核。通过利用MLIR的memref、仿射方言及多面体优化传递——包括展开并合并（unroll-and-jam）、分块（tiling）和向量化——该方法在Intel Skylake CPU上实现了峰值性能的92%，表明模块化、自动化的代码生成在高性能线性代数内核中是可行的。

ABSTRACT

This article is primarily meant to present an early case study on using MLIR, a new compiler intermediate representation infrastructure, for high-performance code generation. Aspects of MLIR covered in particular include memrefs, the affine dialect, and polyhedral utilities and pass infrastructure surrounding those. This article is also aimed at showing the role compiler infrastructure could play in generating code that is competitive with highly tuned manually developed libraries, albeit in a more modular, reusable, and automatable way.

研究动机与目标

• 展示MLIR能够自动生成功能接近手工汇编的高性能GEMM内核，而无需手动编写汇编代码。
• 评估MLIR的多面体与仿射基础设施是否能够实现与OpenBLAS和MKL等手工调优库相媲美的性能。
• 探索使用模块化、可重用且可组合的基于编译器的方法，替代手工优化，以实现密集线性代数内核的可行性。
• 评估MLIR的传递基础设施（包括分块、向量化和标量替换）在实现接近峰值性能中的作用。

提出的方法

• 作者使用MLIR的仿射方言将矩阵乘法表示为具有显式分块参数（例如，M_R=3, N_R=16, K_U=4）的仿射操作。
• 应用一系列MLIR传递：分块、展开并合并、向量化和标量替换，所有操作均基于仿射操作。
• 代码生成流水线从高层MLIR操作开始，通过MLIR基础设施逐步降低至目标特定代码。
• 使用带有仿射布局映射的memref类型来表达步长、转置和内存访问模式，从而实现精确的数据布局分析。
• 变换序列以C++传递的形式编码，支持对优化选择的程序化探索与自动化。
• 最终代码通过MLIR的基础设施生成，避免了手动编写C代码或内联汇编。

实验结果

研究问题

• RQ1 MLIR的多面体与仿射基础设施是否能够自动生成功能接近OpenBLAS和MKL等手工优化库的GEMM内核？
• RQ2 MLIR内置的传递（如展开并合并、向量化和分块）在无需手动调优的情况下，实现高性能的有效性如何？
• RQ3 MLIR的模块化基础设施在多大程度上可以替代高性能线性代数中对手工汇编代码和领域特定优化专业知识的需求？
• RQ4 当分块尺寸不能整除矩阵尺寸时（例如非均匀问题规模），MLIR如何处理而不牺牲性能？
• RQ5 能否将用于分块参数的分析模型有效集成到MLIR的变换流水线中，以指导自动性能优化？

主要发现

• MLIR生成的GEMM内核在Intel i7-8700K CPU上实现了67.49 GFLOPS的性能，达到理论峰值性能（75.2 GFLOPS）的92%。
• MLIR生成内核的性能与OpenBLAS（67.49 GFLOPS）相当，且与MKL（67.70 GFLOPS）的差距仅0.3%。
• 使用- O3和-ffast-math选项的Clang和GCC分别仅实现0.47 GFLOPS和4.53 GFLOPS，凸显了通用编译器的局限性。
• MLIR中的多面体传递（包括展开并合并和向量化）在实现高性能方面至关重要，尤其在与恰当的分块结合时。
• MLIR对仿射布局映射的支持使得步长、转置和内存访问模式的清晰建模成为可能，简化了代码生成。
• 本研究表明，MLIR能够自动完成从高层GEMM到高性能代码的完整优化流水线，显著减少了对手动低层调优的依赖。