[論文レビュー] SPORES: Sum-Product Optimization via Relational Equality Saturation for Large Scale Linear Algebra

Yisu Remy Wang, Shana Hutchison+3 more|arXiv (Cornell University)|Feb 19, 2020

Parallel Computing and Optimization Techniques参考文献 27被引用数 11

ひとこと要約

SPORES は、等価飽和と E-Graph を用いて同等の式を探索することで、大規模な線形代数ワークロードのための完全で関係代数ベースの最適化フレームワークを導入する。SystemML に統合され、既存の最適化を再発見するとともに新たな最適化を発見し、スパarsity と代数的同値性を体系的かつ効果的に活用することで、実際の機械学習ワークロードで最大 5 倍の高速化を達成した。

ABSTRACT

Machine learning algorithms are commonly specified in linear algebra (LA). LA expressions can be rewritten into more efficient forms, by taking advantage of input properties such as sparsity, as well as program properties such as common subexpressions and fusible operators. The complex interaction among these properties' impact on the execution cost poses a challenge to optimizing compilers. Existing compilers resort to intricate heuristics that complicate the codebase and add maintenance cost but fail to search through the large space of equivalent LA expressions to find the cheapest one. We introduce a general optimization technique for LA expressions, by converting the LA expressions into Relational Algebra (RA) expressions, optimizing the latter, then converting the result back to (optimized) LA expressions. One major advantage of this method is that it is complete, meaning that any equivalent LA expression can be found using the equivalence rules in RA. The challenge is the major size of the search space, and we address this by adopting and extending a technique used in compilers, called equality saturation. We integrate the optimizer into SystemML and validate it empirically across a spectrum of machine learning tasks; we show that we can derive all existing hand-coded optimizations in SystemML, and perform new optimizations that lead to speedups from 1.2X to 5X.

研究の動機と目的

• スパarsity、共通部分式、統合可能な演算子といった複数の最適化次元を考慮した、ヒューリスティックベースのコンパイラが複雑な線形代数式を最適化する際の限界を克服すること。
• 関係代数を中間表現として用いることで、同等の最適化可能な線形代数式の完全かつ体系的な探索を実現すること。
• 等価式の探索空間が非常に大きくなる問題を、等価飽和と E-Graph を用いて克服し、完全性を損なわずスケーラブルな探索を可能にすること。
• SystemML に最適化器を統合し、既存の最適化を再発見できるか、および性能を向上させる新たな最適化を発見できるかを実証的に検証すること。

提案手法

• 和、積、転置演算子を関係演算子（例：ジョイン、ユニオン、集約）として表現する正式なマッピングを用いて、線形代数式を関係代数式に変換する。
• 数値を伴う被覆集合（bag semantics）を特徴とする K-関係に基づく関係代数フレームワークを用い、同値性が決定可能であり、結合的クエリの同型性によって完全性が達成可能であることを保証する。
• 標準的な関係恒等式の集合を用いて等価飽和を実行し、E-Graph を構築することで、探索空間内に存在するすべての同等式を効率的に表現する。
• スキーマに依存するクラスインvariantを活用して属性情報を取得し、ルール適用中に定数畳み込みと正確なコスト推定を可能にする。
• 大規模プログラムにおける完全飽和を避けるために、等価飽和の過程でリライトマッチをサンプリングし、探索の深さと収束時間のバランスを取る。
• コストベースの制約ソルバーを用いて E-Graph 内から最適な式を抽出し、その後、実行可能となる線形代数式に変換する。

実験結果

リサーチクエスチョン

• RQ1 関係代数ベースの中間表現は、ヒューリスティックルールが見逃す可能性のある同等の線形代数式を含め、完全に探索可能か？
• RQ2 ユニオン、集約、共通部分式を含む和積式において生じる巨大な探索空間を、等価飽和と E-Graph を用いてどのように管理できるか？
• RQ3 本手法は、SystemML のような本番環境コンパイラに実装された既存の手作業最適化をどの程度再発見し、それを上回れるか？
• RQ4 実世界の機械学習ワークロードにおいて、体系的な探索によって発見された非自明な新規最適化により、どの程度の性能向上が達成可能か？

主な発見

• SPORES は SystemML に実装されたすべての既存の手作業最適化を再発見し、完全性と実用的適用性を実証した。
• 本最適化器は、スパarsity や代数的同値性を活用した、以前は未知であった新たな最適化を発見し、多様な機械学習ワークロードで 1.2× から 5× の性能向上を達成した。
• 被覆集合の関係代数と正規形同値性を用いることで、SPORES はすべての同等式を体系的に探索し、理論的完全性を達成した。
• クラスインvariantの統合により、正確なコスト推定と効果的な定数畳み込みが可能となり、最適化意思決定の質が向上した。
• 等価飽和の過程でリライトマッチをサンプリングすることで、完全飽和を回避しつつも十分な探索効率を達成でき、大規模プログラムへのスケーラビリティが実現した。
• 実験的評価により、制約ソルバーは理論的には最適であるが、大きなオーバーヘッドを伴うことが判明した。したがって、実際にはサンプリングに基づく探索がより効果的であった。