[論文レビュー] 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、アフィンダイアレクト、およびポリヘドラルパス（アンロールアンドジャンプ、タイリング、ベクトル化を含む）を活用することで、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 のポリヘドラルおよびアフィンインfra が、OpenBLAS や MKL のような手動で最適化されたライブラリと同等の性能を達成できるかどうかを評価すること。
• 密度行列代数カーネルに対して、手作業による最適化の代わりに、モジュラーで再利用可能かつ合成可能なコンパイラベースのアプローチが実現可能かどうかを検討すること。
• MLIR のパスインfra（タイリング、ベクトル化、スカラーリプレースを含む）が、ピーク性能に近い性能を達成する上で果たす役割を評価すること。

提案手法

• 著者たちは、行列乗算を明示的なタイルパラメータ（例：M_R=3, N_R=16, K_U=4）を伴うアフィン演算として表現するため、MLIR のアフィンダイアレクトを用いた。
• タイリング、アンロールアンドジャンプ、ベクトル化、スカラーリプレースといった一連の MLIR パスを適用し、すべてがアフィン演算上で動作するようにした。
• コード生成パイプラインは、高レベルの MLIR 演算から始まり、MLIR インfraストラクチャを段階的に低レベル化することで、ターゲット固有のコードに到達する。
• ストライド、転置、メモリアクセスパターンを表現するために、アフィンレイアウトマップを備えた memref 型が使用され、正確なデータレイアウト解析を可能にした。
• 変換シーケンスは C++ パスとしてエンコードされており、最適化選択のプログラマティックな探索と自動化を可能にした。
• 最終的なコードは、MLIR のインfra を用いて生成され、手書きの C やインラインアセンブリを回避した。

実験結果

リサーチクエスチョン

• RQ1 MLIR のポリヘドラルおよびアフィンインfra は、OpenBLAS や MKL のような手作業で最適化されたライブラリと同等の性能を達成する GEMM キャパシティを自動生成できるか？
• RQ2 アンロールアンドジャンプ、ベクトル化、タイリングといった MLIR の組み込みパスは、手動チューニングなしでどれほど高性能を達成できるか？
• RQ3 MLIR のモジュラーインfra は、高性能な線形代数における手書きアセンブリやドメイン特化最適化の専門知識の必要性をどれほど代替できるか？
• 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％以内の差にとどまった。
• Clang と GCC が -O3 および -ffast-math オプションを用いた場合、それぞれ 0.47 GFLOPS および 4.53 GFLOPS にとどまり、一般向けコンパイラの限界を浮き彫りにした。
• MLIR のポリヘドラルパス（特にアンロールアンドジャンプとベクトル化）は、適切なタイリングと組み合わせることで、高性能を達成する上で不可欠であった。
• MLIR がアフィンレイアウトマップをサポートすることで、ストライド、転置、メモリアクセスパターンの明確なモデル化が可能になり、コード生成が簡素化された。
• 本研究は、MLIR が高レベルの GEMM から高性能コードへの最適化パイプラインを完全に自動化できることを示しており、手作業による低レベルチューニングへの依存を低減できることを示している。