[論文レビュー] Sparse Tensor Algebra Optimizations with Workspaces

Fredrik Kjølstad, Willow Ahrens+2 more|arXiv (Cornell University)|Feb 28, 2018

Tensor decomposition and applications参考文献 24被引用数 3

ひとこと要約

本論文は、スパーステンソル代数のための新しい中間言語であるコンcrete index notationと、スパース計算の高速化に寄与するデュアルテンソル（ワークスペース）を用いた最適化を導入する。ループ不変演算の持ち上げと、高コストなスパース挿入およびマージループの削除により、MKL や SPLATT といった手動最適化ライブラリと同等のパフォーマンスを達成し、MTTKRP のようなカーネルで従来不可能だったスパース出力も可能にする。

ABSTRACT

This paper shows how to optimize sparse tensor algebraic expressions by introducing temporary tensors, called workspaces, into the resulting loop nests. We develop a new intermediate language for tensor operations called concrete index notation that extends tensor index notation. Concrete index notation expresses when and where sub-computations occur and what tensor they are stored into. We then describe the workspace optimization in this language, and how to compile it to sparse code by building on prior work in the literature. We demonstrate the importance of the optimization on several important sparse tensor kernels, including sparse matrix-matrix multiplication (SpMM), sparse tensor addition (SpAdd), and the matricized tensor times Khatri-Rao product (MTTKRP) used to factorize tensors. Our results show improvements over prior work on tensor algebra compilation and brings the performance of these kernels on par with state-of-the-art hand-optimized implementations. For example, SpMM was not supported by prior tensor algebra compilers, the performance of MTTKRP on the nell-2 data set improves by 35%, and MTTKRP can for the first time have sparse results.

研究の動機と目的

• スパーステンソル代数コンパイラにおけるパフォーマンスギャップを解消するため、一時的な密テンソル（ワークスペース）を活用する最適化を可能にする。
• SpMM や MTTKRP のようなスパース出力を伴うカーネルを、従来不可能または遅慢だったものを含め、一般化されたスパーステンソルコード生成を可能にする。
• 計算中に条件付きマージループや高コストなスパース挿入を削除することでパフォーマンスを向上させる。
• 強力な最適化が可能な形式的かつ合成可能な中間表現を提供し、スパースコード生成の前段階で積極的な最適化を可能にする。

提案手法

• ループ順序、部分計算の配置、ワークスペース使用を明示する、テンソルインデックス表記の拡張であるコンクリートインデックス表記を導入する。
• スパース式を再書き換えして、一時的な密テンソル（ワークスペース）で部分式を事前に計算することで、実行時のコストを低減するワークスペース最適化を適用する。
• スケジューリング構造を用いてワークスペース挿入を要求し、ループ不変コードの移動を可能にするとともに、繰り返しのスパース挿入を回避する。
• 先行研究（Kjolstad et al., 2017）の反復グラフにコンクリートインデックス表記を低レベル化することで、効率的なスパースコードへのコンパイルを可能にする。
• SpMM、SpAdd、MTTKRP といった主要なカーネルに最適化を適用し、一般性とパフォーマンス向上を実証する。
• マルチウェイマージループをランダムアクセスパターンに置き換えることで、ワークスペースを活用し、条件分岐を削減し、データ局所性を向上させる。

実験結果

リサーチクエスチョン

• RQ1 スパーステンソル代数を、一時的な密テンソル（ワークスペース）を用いてどのように最適化し、パフォーマンスを向上させることができるか？
• RQ2 このワークスペース最適化を、部分計算が発生するタイミングと場所を正確に制御できる高レベルの中間言語で表現・コンパイルすることは可能か？
• RQ3 この最適化により、従来不可能だった MTTKRP や SpMM のようなテンソルカーネルでスパース出力を実現できるか？
• RQ4 最適化されたコードのパフォーマンスは、MKL や SPLATT といった手動最適化の最先端ライブラリと比較してどの程度か？
• RQ5 この最適化は、調査対象のケースを超えて広範なスパーステンソルカーネルのクラスに一般化可能か？

主な発見

• ワークスペース最適化により、nellt-2 データセットにおける MTTKRP のパフォーマンスが、従来のテンソル代数コンパイル手法と比較して 35% 向上した。
• 従来のテンソル代数コンパイラではサポートされていなかった SpMM が、今や効率的にコンパイル可能となり、手動最適化実装と競争力のあるパフォーマンスを達成した。
• 本手法により、MTTKRP に対して初めてスパース出力を可能にし、サポートされるテンソル演算の範囲を拡大した。
• 高コストなスパース挿入と条件付きマージループを削減し、代わりに効率的な密ワークスペースへのランダムアクセスに置き換えた。
• 得られたコードは、MKL や Eigen、SPLATT といった最先端のライブラリと同等のパフォーマンスを達成しており、本手法の有効性を示している。