[論文レビュー] A High-Performance Sparse Tensor Algebra Compiler in Multi-Level IR

Ruiqin Tian, Luanzheng Guo+3 more|arXiv (Cornell University)|Feb 9, 2021

Parallel Computing and Optimization Techniques参考文献 65被引用数 10

ひとこと要約

COMET は、MLIR に基づく高性能なスパーステンソル代数コンパイラであり、アインシュタイン記法に類似した抽象化を備えたドメイン固有言語（DSL）を用いて、混合スパース・ドライブテンソル演算の最適化されたカーネルを自動生成する。4つの次元属性を介して複数のストレージフォーマットをサポートし、データ再配置を実施して局所性を向上させることで、SpMV で TACO より最大 20.92x、SpMM で 6.39x、TTM で 13.9x の高速化を達成している。

ABSTRACT

Tensor algebra is widely used in many applications, such as scientific computing, machine learning, and data analytics. The tensors represented real-world data are usually large and sparse. There are tens of storage formats designed for sparse matrices and/or tensors and the performance of sparse tensor operations depends on a particular architecture and/or selected sparse format, which makes it challenging to implement and optimize every tensor operation of interest and transfer the code from one architecture to another. We propose a tensor algebra domain-specific language (DSL) and compiler infrastructure to automatically generate kernels for mixed sparse-dense tensor algebra operations, named COMET. The proposed DSL provides high-level programming abstractions that resemble the familiar Einstein notation to represent tensor algebra operations. The compiler performs code optimizations and transformations for efficient code generation while covering a wide range of tensor storage formats. COMET compiler also leverages data reordering to improve spatial or temporal locality for better performance. Our results show that the performance of automatically generated kernels outperforms the state-of-the-art sparse tensor algebra compiler, with up to 20.92x, 6.39x, and 13.9x performance improvement, for parallel SpMV, SpMM, and TTM over TACO, respectively.

研究の動機と目的

• 多様なストレージフォーマットとアーキテクチャを対象としたスパーステンソルカーネルの手動最適化の課題に対処すること。
• アインシュタイン記法に類似した DSL を用いて、テンソル代数演算の高生産性プログラミングを可能にすること。
• MLIR フレームワークとの統合を通じて、パフォーマンスのポータビリティと拡張性を提供すること。
• 幅広いストレージフォーマットに対応した、混合スパース・ドライブテンソル演算の効率的なコード生成を可能にすること。
• コンパイルパイプライン内での自動データ再配置技術を活用し、局所性とパフォーマンスを向上させること。

提案手法

• アインシュタイン記法に類似した構文を用いた高水準なテンソル代数式をサポートするドメイン固有言語（DSL）を設計する。
• ドメイン固有の最適化とアーキテクチャ固有の最適化を分離するため、MLIR の多段階中間表現（IR）をコンパイラの基盤に構築する。
• テンソルのストレージフォーマットを、密度（dense）、圧縮ユニーク（compressed unique）、圧縮非一意（compressed non-unique）、シングルトン（singleton）の4つの次元属性で表現する。
• 次元属性の分析に基づき、非ゼロ要素の反復処理に適したフォーマットに依存する効率的なカーネルを生成するコード生成アルゴリズムを実装する。
• カーネル生成段階で空間的・時間的局所性を向上させるために、データ再配置アルゴリズムを統合する。
• MLIR の階層構造内において、上位の IR 層でアーキテクチャに依存しない最適化を適用し、下位の IR 層でレジスタ／メモリ最適化を実施する。

実験結果

リサーチクエスチョン

• RQ1 どのような設計により、幅広いスパーステンソルストレージフォーマットをサポートしつつ、高いパフォーマンスを維持できるスパーステンソル代数コンパイラを実現できるか？
• RQ2 アインシュタイン記法に基づく DSL は、スパーステンソル計算におけるプログラマの生産性をどの程度向上できるか？
• RQ3 MLIR を用いたコンパイルアーキテクチャは、多様なアーキテクチャとフォーマットに対し、パフォーマンスのポータビリティと拡張性を実現できるか？
• RQ4 自動化されたデータ再配置は、スパーステンソル演算におけるデータ局所性とカーネルパフォーマンスの向上にどの程度効果的か？
• RQ5 SpMV、SpMM、TTM といった主要なスパーステンソルカーネルにおいて、TACO よりもどの程度のパフォーマンス向上が達成できるか？

主な発見

• COMET は、CPU におけるスパース行列-ベクトル乗算（SpMV）で TACO より最大 20.92x のパフォーマンス向上を達成した。
• COMET は、CPU におけるスパーステンソル-行列乗算（SpMM）で TACO より 6.39x のパフォーマンス向上を示した。
• COMET は、CPU におけるテンソル変換（TTM）演算で TACO より 13.9x のパフォーマンス向上を達成した。
• データ再配置の活用により、データ局所性が顕著に向上し、観察されたパフォーマンス向上に寄与した。
• 4つの次元属性によるフォーマットに依存しない設計により、COO、CSR、DCSR、ELLPACK、CSF およびモード汎用フォーマットをハードコードせずに対応可能となった。
• MLIR を基盤とするアーキテクチャにより、モジュラーで拡張可能かつポータブルなコンパイルが実現され、GPU を含む異種アーキテクチャへの将来的な対応が容易になった。