[論文レビュー] Minimum Cost Loop Nests for Contraction of a Sparse Tensor with a Tensor Network

Raghavendra Kanakagiri, Edgar Solomonik|arXiv (Cornell University)|Jul 11, 2023

Tensor decomposition and applicationsMathematics被引用数 3

ひとこと要約

本論文は、スパーステンソル結合（SpTTN）における最も効率的なループネストを自動的に特定・実行するフレームワーク、SpTTN-Cyclopsを提示する。動的計画法を用いてバッファサイズやキャッシュミスのコスト指標を最小化する。一般向けライブラリを上回り、テンソル分解および補完ワークロードにおけるデータに依存しないスパarsityパターンを活用することで、特化型コードと同等の性能を達成する。

ABSTRACT

Sparse tensor decomposition and completion are common in numerous applications, ranging from machine learning to computational quantum chemistry. Typically, the main bottleneck in optimization of these models are contractions of a single large sparse tensor with a network of several dense matrices or tensors (SpTTN). Prior works on high-performance tensor decomposition and completion have focused on performance and scalability optimizations for specific SpTTN kernels. We present algorithms and a runtime system for identifying and executing the most efficient loop nest for any SpTTN kernel. We consider both enumeration of such loop nests for autotuning and efficient algorithms for finding the lowest cost loop-nest for simpler metrics, such as buffer size or cache miss models. Our runtime system identifies the best choice of loop nest without user guidance, and also provides a distributed-memory parallelization of SpTTN kernels. We evaluate our framework using both real-world and synthetic tensors. Our results demonstrate that our approach outperforms available generalized state-of-the-art libraries and matches the performance of specialized codes.

研究の動機と目的

• テンソル分解および補完におけるパフォーマンスボトルneckを解消するため、スパーステンソル結合と密度の高いテンソルネットワーク（SpTTN）が主要な計算カーネルを占める。
• ユーザーの指導なしに、可能なループ順序の組み合わせの爆発的増加を克服し、SpTTNカーネルの最も効率的なループネストの選択を自動化する。
• 既存の高性能テンソルライブラリと統合された汎用的で自動チューニング可能なフレームワークを通し、高性能でスケーラブルかつポータブルなSpTTN実行を実現する。
• データに依存しないスパarsityとBLASカーネルの効率的利用を活用することで、手で最適化された特化型コードの性能を再現または上回る。

提案手法

• 各ノードがループで、その子がネストされたループである木構造としてループネストを表現し、構造的なコストモデリングを可能にする。
• 指数的空間のループ順序を効率的に探索するため、動的計画法を適用し、計算量をO((m!)^N)からO(N^3 2^m m)に削減する。
• ループネスト木構造に基づき、バッファサイズとキャッシュミスといった主要なパフォーマンスボトルneckのコストモデルを構築する。
• CTFランタイムシステムと統合し、SpTTNカーネルの分散メモリ並列化を可能にする。
• コスト関数を最小化するループ順序を自動的に選択するとともに、高性能なBLAS-1およびBLAS-2カーネルの高レベルな利用を可能にする。
• Cyclopsテンソルフレームワークを拡張し、1つのスパーステンソルを、テンソルネットワーク内の複数の密度の高いテンソルと同時に結合する機能をサポートする。

実験結果

リサーチクエスチョン

• RQ1 与えられたSpTTNカーネルに対して、メモリフットプリントとキャッシュミスを最小化する最適なループネスト構造は何か？
• RQ2 全列挙を避けながら、SpTTNカーネルの膨大なループ順序の置換空間を効率的に探索するために、動的計画法を用いることができるか？
• RQ3 自動生成されたループネストのパフォーマンスは、テンソル分解および補完ワークロードにおける手で最適化された特化型実装と比べてどうか？
• RQ4 中間テンソルの次元の上限が、ループネストの効率性とBLASカーネルの利用に与える影響は何か？
• RQ5 汎用的なフレームワークは、多様なSpTTNカーネルにわたってポータブルでありながら、特化型コードと同等のパフォーマンスを達成できるか？

主な発見

• スモールサイズで極めてスパースなテンソル（N=40、スパarsity 0.1%）において、TACOに対して534倍の高速化を達成し、低次元領域における顕著なパフォーマンス向上を示した。
• 本フレームワークは、TACO や SparseLNR といった一般向けライブラリを、実世界および合成的なSpTTNカーネルの多様なスケールで上回った。
• スカラの中間テンソルよりも、サイズTおよびS×Tに制限された中間テンソルを有するループネストの方が、より高いメモリ使用量を伴いながらも、BLAS-1およびBLAS-2カーネルの利用度が高いため、パフォーマンスに優れた。
• SpTTN-Cyclopsが選択したループ順序は、25%のランダムに抽出されたループ順序の中でも最高のパフォーマンスを発揮しており、効果的な自動チューニングであることが示された。
• MTTKRP、TTMc、TTTPカーネルにおいて、本フレームワークは特化型コードと同等のパフォーマンスを達成し、その汎用性と効率性を検証した。
• 動的計画法の使用により、探索空間の計算量が階乗的から多項式的へと削減され、SpTTNループネストの実用的自動チューニングが可能になった。