[論文レビュー] 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

ひとこと要約

本論文は、NVIDIA A100およびRTX 3090 GPUのテンソルコアを対象として、高性能な行列同士の乗算のためのMLIRベースの自動コード生成パイプラインを提示する。最適化されたダイアレクト変換を通じて高レベル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を用い、matmulを高レベル、中レベル、低レベルの抽象化レベルで表現し、体系的な変換を可能にする。
• Warp Matrix Multiply Accumulate (WMMA) 操作を扱う新しいMLIRダイアレクトを導入し、LLVM/NVPTX IRへの下位変換を実施する。
• タイリング、共有メモリへのタイリング、レジスタレベルの最適化を含む、ダイアレクト下位変換および最適化パスのシーケンスを適用する。
• パイプライン化されたグローバルメモリロードやタイルサイズの選定といった最適化を段階的に適用し、その性能への影響を検証する。
• 高レベルのmatmul演算から始まり、テンソルコアをターゲットにした効率的なCUDAコードにコンパイルする、エンドツーエンドのコード生成パイプラインを実装する。
• 性能評価は、標準的な問題サイズと混合/半精度設定を用いて、NVIDIAのAmpereアーキテクチャ搭載RTX 3090 GPU上で実施する。

実験結果

リサーチクエスチョン

• RQ1 IRベースのコード生成パイプラインは、テンソルコア上で行列同士の乗算に、CuBLASのような手動最適化ライブラリと同等の性能を達成できるか？
• RQ2 タイルサイズやメモリアクセスのパイプライン化といった個々の最適化が、モジュラーかつIR駆動のパイプラインにおいて性能に与える影響はいかほどか？
• RQ3 自動コード生成が、異なる精度形式において、ベンダー最適化ライブラリの性能をどれほど模倣または上回ることができるか？
• RQ4 MLIRのような統一されたIRインfraは、テンソルコアのような専用アクセラレータ向けの高レベル計算抽象を効果的に表現・変換できるか？
• RQ5 自動生成されたカーネルの性能は、さまざまな問題サイズとデータ型において、CuBLASと比べてどの程度か？

主な発見

• 自動生成されたカーネルは、RTX 3090におけるFP32の混合精度matmulで、CuBLAS性能の95–119％を達成し、デバイスの理論ピーク性能の95.4％にまで到達する高い効率性を示した。
• FP16の混合精度matmulでは、性能が常にCuBLAS性能の95–119％の範囲に収まり、小さな問題サイズではより高いスレッドブロックタイルによる高いスループットのおかげで、CuBLASを上回った。
• FP16専用のmatmulでは、生成コードがCuBLAS性能の80–160％を達成し、CuBLASが不規則な性能を示し、非効率なタイル選択をしていた大きな問題サイズでは、CuBLASを上回った。
• 研究により、CuBLASはより小さなタイルサイズ（例：128×128×32）と5段階のパイプライン化を採用しているが、より高いグローバルメモリスタールが発生しており、遅延隠蔽が不十分であることが判明した。
• 小さなタイルサイズ（例：64×64×64）は、小さな問題サイズで高いスループットを実現するが、大きな行列ではより良いデータ再利用が可能になるため、より大きなタイルが効果的である。
• 段階的な最適化の検証により、パイプライン化とタイルサイズ選定が性能に顕著な影響を与えることが確認され、最終的な最適化パイプラインはピーク性能に近い性能を達成した。