[論文レビュー] High-Performance Tensor Contraction without Transposition

Devin A. Matthews|arXiv (Cornell University)|Jul 1, 2016

Parallel Computing and Optimization Techniques参考文献 33被引用数 5

ひとこと要約

本論文では、TBLISと呼ばれる高性能テンソル乗算ライブラリを提示する。このライブラリは、BLISベースの行列乗算カーネルとテンソルから行列へのマッピングを統合することで、明示的な転置を排除する。BLISフレームワークの柔軟なアルゴリズム的構造を活用することで、著者たちは行列乗算に近いパフォーマンスを達成し、特に複雑なテンソル形状において、TTGT や BSMTC といった従来のアプローチを上回る性能を発揮する。これは、シングルコアおよびマルチコアワークロードの両方で顕著である。

ABSTRACT

Tensor computations--in particular tensor contraction (TC)--are important kernels in many scientific computing applications. Due to the fundamental similarity of TC to matrix multiplication (MM) and to the availability of optimized implementations such as the BLAS, tensor operations have traditionally been implemented in terms of BLAS operations, incurring both a performance and a storage overhead. Instead, we implement TC using the flexible BLIS framework, which allows for transposition (reshaping) of the tensor to be fused with internal partitioning and packing operations, requiring no explicit transposition operations or additional workspace. This implementation, TBLIS, achieves performance approaching that of MM, and in some cases considerably higher than that of traditional TC. Our implementation supports multithreading using an approach identical to that used for MM in BLIS, with similar performance characteristics. The complexity of managing tensor-to-matrix transformations is also handled automatically in our approach, greatly simplifying its use in scientific applications.

研究の動機と目的

• テンソル乗算における明示的転置のパフォーマンスおよびメモリオーバーヘッドを排除すること。
• コード生成やランタイムオーバーヘッドを回避する汎用的で高性能なテンソル乗算ライブラリの開発。
• 柔軟で拡張可能なフレームワークを用いて、最適化された行列乗算と同等のパフォーマンスを達成すること。
• テンソルレイアウト管理における最小限のユーザー介入で、効率的なマルチスレッド実行を可能にすること。
• 多様なテンソル形状と乗算パターンをサポートする生産環境向けオープンソースソリューション（TBLIS）を提供すること。

提案手法

• 著者たちは、マトリクスカーネル内でのテンソル要素へのイン-place アクセスを可能にする2つの新規なテンソルから行列へのマッピング手法——Scatter-Matrix Tensor Contraction（SMTC）およびBlock-Scatter-Matrix Tensor Contraction（BSMTC）——を設計した。
• これらのマッピングは、BLISフレームワークに統合され、転置とデータパッケージングがカーネル計算と統合されることで、別個の転置操作や追加のワークスペースが不要になる。
• このアプローチは、行列乗算が独立したアルゴリズム的コンponentに分解可能なBLISのモジュラー設計を活用しており、MMに用いられる同じプリミティブからカスタムテンソル乗算アルゴリズムを構築できる。
• 実装ではブロックベースのデータレイアウトと最適化されたメモリアクセスパターンを用いて、キャッシュ再利用とベクトル化を向上させ、データ移動を最小限に抑える。
• マルチスレッド化は、行列乗算と同一のスレーディングモデルを採用することで実現され、一貫性のあるパフォーマンススケーリングを確保する。
• TBLISライブラリは、テンソルから行列へのインデックス変換を自動で処理し、ユーザーが低レベルの複雑さを意識しなくてよいように抽象化する。

実験結果

リサーチクエスチョン

• RQ1 明示的な転置や追加のワークスペースが不要な状態で、テンソル乗算を行列乗算に近いパフォーマンスで実装できるか？
• RQ2 多様なテンソル形状において、BLISベースのテンソル乗算のパフォーマンスは、従来のTTGTおよびBSMTCアプローチと比べてどうか？
• RQ3 計算制限および通信制限の両状況下で、提案手法は複数コアにわたってどの程度スケーリングするか？
• RQ4 BLISフレームワークのモularityを効果的に活用して、高性能で汎用的なテンソル乗算カーネルを構築できるか？
• RQ5 非効率なデータアクセスパターンがテンソル乗算パフォーマンスに与える影響は何か？ また、それらはどのように緩和できるか？

主な発見

• BSMTCアルゴリズムは、多くのテンソル形状においてピークマシン性能の10％以内のパフォーマンスを達成し、行列乗算の効率に近づいている。
• シングルコアの計算制限ケースでは、BSMTCとBLISは類似したパフォーマンスを示し、ピークに近く、一方TTT（TTGT）は大きなパフォーマンス低下を示し、大きなテンソルでは10 GFLOPs未満に低下する。
• 通信制限ケースでは、BSMTCは約30 GFLOPsを維持するが、TTTは10 GFLOPs未満に低下し、提案手法の顕著なパフォーマンス優位性が示された。
• マルチコア実行では、BSMTCはTTTに対して最大21.1倍のスルーブロックを達成し、テンソル形状に応じて1.3倍から21.1倍の範囲でスルーブロックが変動し、強力なスケーラビリティを示した。
• いくつかのケースでBSMTCはBLISを上回るパフォーマンスを示しており、これは内部でより最適な行列乗算バリアントが使用されている可能性を示唆しており、アルゴリズムカスタマイズの利点を強調している。
• TBLISライブラリは、TTGTのワークスペースオーバーヘッドを効果的に回避し、BSDライセンス下での生産環境向けオープンソース実装を提供した。