はじめに

スパース性を（面倒な）実装の詳細ではなくプロパティとして扱うことで、コンパイラのスパース化ツール1 がスパースコードを自動的に生成できるようになります。このコンセプトは、MT1 の [Bik96] によって線形代数で初めて導入され、Sparse Tensor Algebra Compiler （TACO）プロジェクトの [Kjolstad17,Kjolstad20] によってテンソル代数に正式に導入されました。

SparseTensor 言語は、MLIR コンパイラ インフラストラクチャ内でスパース テンソル型をファーストクラスの市民にするために必要なすべての属性、型、オペレーション、パスをサポートしています。この言語は、スパース テンソル型の高レベル オペレーションと、位置、座標、値で構成される実際のスパース ストレージ スキームの低レベル オペレーションの間の橋渡しをします。下位レベルのサポートは、完全に生成されたコードで構成される場合もあれば、小さなスパース ランタイム サポート ライブラリによって提供される場合もあります。

MLIR 実装 [Biketal22] は、TACO の基盤となる「スパース反復理論」に厳密に準拠しています。コンパイラには、Linalg 言語（MLIR のテンソル インデックス表記）の各テンソル式に対する書き換えルールがあります。テンソルの各レベルのスパース性は、レベルタイプ（密、圧縮、シングルトンなど）とレベルの順序の指定によって示されます（これらのレベルタイプの詳細な説明と拡張の可能性については、[Chou18] を参照してください）。コンパイラは、入力の各テンソル式に対して、まず各テンソルのレベルに関する座標の順序を反映したトポロジカル ソートされた反復グラフを構築します。これにより、すべてのテンソルが自然なレベル座標順序でアクセスされるようになります。次に、トポロジカル順序で各インデックスの反復格子が構築されます。各反復格子点は、テンソル座標の連言と、その連言に対して評価する必要があるテンソル（サブ）式で構成されます。ラティス内では、座標の使い果たし方に応じて反復点が順序付けられます。そのため、これらの反復格子は実際のスパースコード生成を駆動します。反復格子から for ループ、while ループ、if ステートメントの組み合わせへの比較的簡単な 1 対 1 のマッピングです。初期化されていない疎テンソル出力は、すべての並列ループが最外ループである場合は直接挿入で処理され、そうでない場合は 1 次元アクセス パターン拡張（ワークスペース）を介して間接的に挿入されます（可能な場合）[Gustavson72、Bik96、Kjolstad19]。

  • [Bik96] Aart J.C. Bik。スパース行列計算のコンパイラ サポート。1996 年 5 月発行、ライデン大学博士論文。

  • [Biketal22] Aart J.C. Bik、Penporn Koanantakool、Tatiana Shpeisman、Nicolas Vasilache、Bixia Zheng、Fredrik Kjolstad。MLIR でのスパース テンソル計算のコンパイラ サポート。ACM Transactions on Architecture and Code Optimization、2022 年 6 月。参照: https://dl.acm.org/doi/10.1145/3544559

  • [Chou18] Stephen Chou、Fredrik Berg Kjolstad、Saman Amarasinghe。 スパース テンソル代数コンパイラの形式の抽象化。Proceedings of the ACM on Programming Languages, October 2018.

  • [Chou20] Stephen Chou、Fredrik Berg Kjolstad、Saman Amarasinghe。 効率的なスパース テンソル形式変換ルーチンの自動生成。Proceedings of the 41st ACM SIGPLAN Conference on Programming Language Design and Implementation, June, 2020.

  • [Gustavson72] Fred G. Gustavson。疎な連立一次方程式を解くための基本的な手法。Sparse Matrices and Their Applications（スパース行列とその応用）、41 ～ 52 ページ。Plenum Press, New York, 1972.

  • [Kjolstad17] Fredrik Berg Kjolstad、Shoaib Ashraf Kamil、Stephen Chou、David Lugato、Saman Amarasinghe。Tensor Algebra Compiler。Proceedings of the ACM on Programming Languages、2017 年 10 月。

  • [Kjolstad19] Fredrik Berg Kjolstad、Peter Ahrens、Shoaib Ashraf Kamil、Saman Amarasinghe。Tensor Algebra Compilation with Workspaces, Proceedings of the IEEE/ACM International Symposium on Code Generation and Optimization, 2019.

  • [Kjolstad20] Fredrik Berg Kjolstad. スパース テンソル代数コンパイル。MIT の博士論文、2020 年 2 月。

ご協力いただける場合は、最初に取り組むのに適した問題をいくつかご紹介します。

  1. なお、このようなパスを指す用語として、一般的に使用されている「スパース コンパイラ」ではなく「スパース化ツール」という用語を使用することをおすすめします。これは、スパース化ツールパスが個別のコンパイラではなく、MLIR コンパイラ インフラストラクチャで構築されたコンパイラ パイプラインの不可欠な部分であることを明確にするためです。 ↩