[論文レビュー] Polyhedral Specification and Code Generation of Sparse Tensor Contraction with Co-Iteration

Tuowen Zhao, Tobi Popoola+3 more|arXiv (Cornell University)|Aug 25, 2022

Parallel Computing and Optimization Techniques被引用数 6

ひとこと要約

本論文は、配置と計算空間を分離することで、柔軟な配置指定とSMTで合成されたfindアルゴリズムを介した効率的な共反復を可能にする、スパーステンソル畳み込みのポリヘドラフレームワーク拡張を提示する。TACOより平均1.63倍の高速化を達成し、最適化されたfindアルゴリズムを用いることで最大2.72倍の高速化を実現。BCSRや対角計算など、従来不可能だった配置もサポートする。

ABSTRACT

This paper presents a code generator for sparse tensor contraction computations. It leverages a mathematical representation of loop nest computations in the sparse polyhedral framework (SPF), which extends the polyhedral model to support non-affine computations, such as arise in sparse tensors. SPF is extended to perform layout specification, optimization, and code generation of sparse tensor code: 1) we develop a polyhedral layout specification that decouples iteration spaces for layout and computation; and, 2) we develop efficient co-iteration of sparse tensors by combining polyhedra scanning over the layout of one sparse tensor with the synthesis of code to find corresponding elements in other tensors through an SMT solver. We compare the generated code with that produced by a state-of-the-art tensor compiler, TACO. We achieve on average 1.63$ imes$ faster parallel performance than TACO on sparse-sparse co-iteration and describe how to improve that to 2.72$ imes$ average speedup by switching the find algorithms. We also demonstrate that decoupling iteration spaces of layout and computation enables additional layout and computation combinations to be supported.

研究の動機と目的

• 計算と配置空間が強く結合されているため、BCSR や対角計算のような複雑な配置をサポートできない、既存のスパーステンソルコンパイラの制限を解消すること。
• 論理的座標空間と物理的配置反復空間を分離することで、複数のスパーステンソルの一般化された共反復を可能にすること。
• 未解釈関数とSMTソルバを用いて、正式な配置仕様と共反復サポートを拡張したスパースポリヘドラフレームワーク（SPF）を拡張すること。
• 配置変換と共反復パターンの両方をサポートする、アーキテクチャに配慮した最適化コードを生成すること。今後は論理的マージ（disjunctive merge）も含む。
• 異なるハードウェア、特にSIMDやテンソルコアを含む環境において、データ配置と計算の共最適化を可能にすることで、パフォーマンスの移植性を向上させること。

提案手法

• 配置を物理的位置と論理的座標の数学的関係としてモデル化することで、配置と計算の反復空間を分離し、独立した最適化を可能にする。
• ポリヘドラモデルにおける実行時依存のインデックス配列を未解釈関数で表現することで、スパース計算における非アフィンアクセスをサポートする。
• SMTソルバを用いて、1つのテンソルの配置から別のテンソルの非ゼロ要素に対応するインデックスをマッピングするfindアルゴリズムを合成する。これにより共反復が可能になる。
• 1つのテンソルの配置に対するポリヘドラスキャンと、合成されたコードを組み合わせることで、他のテンソルの一致する非ゼロ要素を特定し、明示的なマージラティスを回避する。
• 1つの配置レベルと計算インデックスの1対1対応を必要としないため、ブロック形式（例：BCSR）や対角計算などの複雑な配置をサポートする。
• 配置マッピングと計算変換を合成することで最適化コードを生成し、ブロッキング、再順序付け、ベクトル化などの最適化と直交的に組み合わせることで、現代のアーキテクチャ上でパフォーマンスを向上させる。

実験結果

リサーチクエスチョン

• RQ1 ポリヘドラフレームワークは、配置と計算反復空間を分離しつつ、スパーステンソル配置の正式な仕様をサポートできるか？
• RQ2 形式的かつ数学的に根拠のあるフレームワークを用いて、複数のスパーステンソルの共反復を効率的に表現・コンパイルできるか？
• RQ3 SMTで合成されたfindアルゴリズムは、手動で書かれたものやテンプレートベースのアプローチを上回る共反復パフォーマンスを達成できるか？
• RQ4 配置と計算を、BCSR や対角演算など、従来サポートされていなかった形式を含めて、どのように共最適化できるか？
• RQ5 本フレームワークで生成されたコードのパフォーマンスは、TACO などの最先端のテンソルコンパイラと比較して、多様なスパーステンソルワークロードにおいてどの程度優れているか？

主な発見

• 提示されたフレームワークは、ベンチマークスイート上での並列スパース-スパース共反復において、TACOに対して平均1.63倍の高速化を達成した。
• 最適化されたfindアルゴリズムに切り替えることで、フレームワークはTACOに対して最大2.72倍の平均高速化を達成し、アルゴリズム合成のパフォーマンスへの影響を実証した。
• 配置と計算空間の分離により、TACOが表現できない、ブロック圧縮スパース行（BCSR）のような従来不可能だった配置もサポート可能になった。
• フレームワークは、複数のレベルにまたがる共通のループインデックスを必要とする対角計算やその他のパターンもサポートでき、TACOのレベルフォーマットモデルでは不可能だった。
• SMTで合成されたfindアルゴリズムの使用により、マージラティスや事前計算されたインデックスマッピングを必要とせず、効率的な共反復が可能になった。
• 本アプローチにより、ブロッキングやベクトル化などの他の最適化と直交的に組み合わせることが可能となり、現代のアーキテクチャ上でパフォーマンスを向上させることができる。