[論文レビュー] SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning

Zihao Ye, Ruihang Lai+3 more|arXiv (Cornell University)|Jul 11, 2022

Advanced Neural Network Applications被引用数 5

ひとこと要約

SparseTIR は、スパースディープラーニングコンパイルのための合成可能な中間表現を導入し、合成可能なフォーマットと変換によって効率的な最適化を可能にする。スパース演算子においてベンダー・ライブラリを 1.05–7.45× 、GNN のエンド・ツー・エンド推論において最大 40.18× の高速化を達成しており、動的に最適なスパースフォーマットとハードウェアに最適化された最適化を組み合わせることで実現している。

ABSTRACT

Sparse tensors are rapidly becoming critical components of modern deep learning workloads. However, developing high-performance sparse operators can be difficult and tedious, and existing vendor libraries cannot satisfy the escalating demands from new operators. Sparse tensor compilers simplify the development of operators, but efficient sparse compilation for deep learning remains challenging because a single sparse format cannot maximize hardware efficiency, and single-shot compilers cannot keep up with latest hardware and system advances. In this paper, we observe that the key to addressing both these challenges is to leverage composable formats and composable transformations. We propose SparseTIR, a sparse tensor compilation abstraction that offers composable formats and composable transformations for deep learning workloads. SparseTIR constructs a search space over these composable components for performance tuning. With these improvements, SparseTIR obtains consistent performance speedups vs vendor libraries on GPUs for single operators: 1.20-2.34x for GNN operators, 1.05-2.98x for sparse attention operators, and 0.56-7.45x for sparse convolution operators. SparseTIR also accelerates end-to-end GNNs by 1.08-1.52x for GraphSAGE training, and 4.20-40.18x for RGCN inference.

研究の動機と目的

• 新たな多様なワークロードが登場する中で、既存のベンダー・ライブラリが対応していないスパース演算子の高性能開発の課題に取り組む。
• 異種のハードウェアや進化を続けるディープラーニングワークロードに適応できない単一フォーマットのスパースコンパイラの限界を克服する。
• フォーマットとスケジューリングの意思決定を分離することで、スパース演算子のための効率的で自動化されたパフォーマンスチューニングを可能にする。
• 合成可能なコンパイル抽象化を通じて、GNN やスパーストランスフォーマーのような複雑なディープラーニングワークロードのエンド・ツー・エンド加速を実現する。
• 既存の密度テンソルコンパイラ最適化を再利用しながら、新しいスパース固有の変換を可能にするスケーラブルで拡張可能なフレームワークを提供する。

提案手法

• スパースフォーマットの定義と計算記述を分離する、合成可能な中間表現（IR）である SparseTIR を導入する。
• 局所的なスパarsityパターンに基づいて、スパーステンソルの異なる部分を最適なフォーマットに保存できるように、フォーマットの合成性を実現する。
• コンパイルを再利用可能で合成可能なプログラム変換のシーケンスに分解することで、変換の合成性を実現する。
• 既存の密度テンソルコンパイラ（例：テンソル化、タイリング）のループレベルの抽象化を活用し、スパースデータ処理に拡張する。
• 再利用可能なフォーマットと変換の組み合わせを探索するための探索空間を構築し、構成可能な探索システムを用いて自動パフォーマンスチューニングを支援する。
• 既存のコンパイラからハードウェア固有の最適化（例：テンソルコア、GPUマッピング）を統合し、現代のアクセラレータ上でスパースカーネルを高速化する。

実験結果

リサーチクエスチョン

• RQ1 合成可能なフォーマットは、多様なスパースディープラーニングワークロードにおけるパフォーマンスの移植性とハードウェア利用度を向上させることができるか？
• RQ2 合成可能な変換は、スパースコンパイルにおいて密度テンソルコンパイラ最適化の効率的で再利用可能な実装を可能にするか？
• RQ3 合成可能なコンponentを対象とした探索ベースのチューニングシステムは、新興のスパース演算子においてベンダー最適化ライブラリを上回る性能を発揮できるか？
• RQ4 SparseTIR は、ベースラインライブラリと比較して、エンド・ツー・エンドの GNN 学習および推論をどの程度高速化できるか？
• RQ5 フォーマットの合成性は、スパースカーネルにおけるテンソルコア やメモリ階層といったハードウェア機能のより良い活用を可能にするか？

主な発見

• GNN SpMM 演算子において、SparseTIR はベンダー・ライブラリを 1.20–2.34× 速度向上させ、特に不規則なスパarsityパターンの高い場合に顕著な向上を示した。
• スパースアテンション演算子では、SparseTIR が 1.05–2.98× の速度向上を達成し、中程度のスパースで構造化されたスパarsityに対しても有効であることが示された。
• スパース畳み込み演算子では、SparseTIR が最大 7.45× の速度向上を達成し、ポイントクラウドベースのワークロードにおいてベンダー・ライブラリを大きく上回った。
• エンド・ツー・エンドの GraphSAGE 学習は 1.08–1.52× 速度向上し、フルモデルの学習パイプライン全体で一貫した向上が得られた。
• RGCN 推論では、4.20× から 40.18× の顕著な高速化が達成され、複雑で多様なスパarsityを持つワークロードにおける合成最適化の利点が顕著に現れた。
• アーティファクトは、異なるGPU（RTX 3070、V100）で再現性を示しており、L2キャッシュフラッシュを有効にした場合、論文で報告された数値と非常に近い結果が得られた。