SPORES

SPORES: Sum-Product Optimization via Relational Equality Saturation for Large Scale Linear Algebra

egraph

@article{wang2020spores,
  title={SPORES: sum-product optimization via relational equality saturation for large scale linear algebra},
  author={Wang, Yisu Remy and Hutchison, Shana and Leang, Jonathan and Howe, Bill and Suciu, Dan},
  journal={arXiv preprint arXiv:2002.07951},
  year={2020},
  url={https://arxiv.org/abs/2002.07951}
}

Equality Saturation を用いて密行列を疎行列に変換する

Abstract

• 機械学習は線形代数で記述される
  • Sparse性や 共通部分式除去 , fusible operators などの特性を利用することでより効率の良い形式に書き換えられる
    • がコンパイラの最適化には困難が伴う
• 線形代数を関係代数(RA)式に変換して最適化して線形代数に戻す
  • 探索空間の大きさが課題だが Equality Saturation を使うことで対処
• ルールサンプリングと制約ソルバにより最適なプログラムを抽出
• 実行時間最大x10 faster
• SystemML: Reference Guide for Python Users - SystemML 1.2.0

Introduction

• $\sum (X-U{V}^T{)}^2$ は典型的な $X$ と低ランク行列 $U{V}^T$ の誤差
• $X$ が疎行列で 1M * 500k で, $U,V$ が密行列だったら掛け算の回数でかい
  • 幸い $\sum X^2-2U^{T}XV+U^{T}U\ast V^{T}V$ と等価で $U^{T}XV$ がスカラなので 1M + 500k multで済む
• 13ruleのみproveする
• RAは K-relations 上
• タプル $X(i,j)$ はboolではなく多重度を持つ

• SystemML上に SPORES を統合し, SystemMLのすべてのハンドコード化されたルールを導出できることを示した
• SPORESを様々な機械学習タスクで評価した結果、より成熟したヒューリスティックベースのオプティマイザと比較して、競争力のある性能向上を示すことができました
• 後者の最適化をすべて発見し, さらに10倍以上の高速化に寄与する新たな最適化も発見した。本論文では、以下の貢献を行う
  • 複雑な線形代数表現を関係代数に変換して最適化する新しいアプローチを説明し、書き換えルールの完全性を証明する. (2 REPRESENTING THE SEARCH SPACE)
  • メモリを再利用しながら広い探索空間を探索することができる等価飽和に基づく探索アルゴリズムを発表する (3 EXPLORING THE SEARCH SPACE)
  • 実世界の機械学習タスクを用いた実証評価を行い、SystemMLのヒューリスティック駆動型オプティマイザに対する優位性を実証する

2. REPRESENTING THE SEARCH SPACE

• LAからRAに変換してまた戻すのを $R_{LR}$ と書く

$$\begin{array}{rl}& \sum (X-U{V}^T{)}^2\\ =& \sum (X-U{V}^T)(X-U{V}^{T)}\\ =& \sum X^2-2X\ast U{V}^T+(U{V}^T{)}^2\\ =& \sum X^2-2\sum X\ast U{V}^T+\sum (U{V}^T{)}^2\end{array}$$

添字を用いた表現はRAの式を表す, ベクトル, 行列, テンソルは実数上の K-relations として解釈.
$X_{ij}$ を $X(i,j)$ として捉えて多重度を行列要素の実数値にする.

要素積は自然結合として解釈し??, additionはunionとする. sumは集約として解釈.
行列積はjoin上の集約
(実装では、ポイント単位の乗算と加算に外部結合を使用し、それに応じて行列のエントリを乗算、加算しています。本論文では、表示を簡単にするためにjoinとunionを使用しています)

表記の話.

$A\otimes x^T$ はRAではjoinで表される $A(i,j)\bowtie x(j)$

2.3 Completeness of the Optimization Rules

• LA式の書き換えが行き詰まってもRAに変換することで進むことがある
• このアプローチの完全性
  • 2つのLAが等価であれば等価性は $R_{EQ}$ を用いることで証明できることを証明する

証明の流れ

1. $R_{EQ}$ は任意のRA式 $e$ を正規形 $\mathcal{C}(e)$ と相互変換することが出来る
2. RA式 $e,e^{\prime }$ は正規形 $\mathcal{C}(e),\mathcal{C}(e^{\prime })$ が等価であれば等価

定義: Relations

半環 $S$ を固定して, $S$-relationは関数 $A:D^a\to S$, $a$ は関係のアリティ(演算の引数の個数)

定義: Expressions

atom $r=R(x_1,\cdots ,x_a)$

• $R$: 関係の名前, タプルはand/or定数の変数??

集約演算の順番は関係ないので ${\sum }_{\{x_1,\cdots ,x_n\}}$ とかく
式 $e$ で ${\sum }_x{e}^{\prime }$ が含まれているとき $x$ は拘束されているという
そうでない時 自由

$\mathrm{vars}(e)=fv(e)\cup bv(e)$
ラムダ式内の任意のatom $R(x_1,\cdots ,x_a)$ はある $s\in S$ に評価される, $+,\times ,\sum$ は $S$ 上で計算する
RA式が等価とは同じ入力で同じ結果になることをいう

定義: Equivalence of Expressions

relation symbols $R_1,\cdots ,R_n$ 上の $e_1,e_2$ を固定.
2値が $S$ 上等価とは同じ自由変数を持ち, 全てのinterpretations $I=(R_1^I,\cdots ),R_j^I:D^{a_j}\to S$ を2値は同じ結果を返す.

定義: R-monomials, Terms, R-polynomials

R-monomial $m:=r_1\times \cdots r_n$ atom の積
terms $t:={\sum }_{x}m$
R-polynomials: $e:=c_0+c_1{t}_1+\cdots +c_n{t}_n$

3. EXPLORING THE SEARCH SPACE

Constraint Solving and Cost Function

• 抽出するためにILPでencodingする
  • グラフの各演算子に対して bool $B_{op}$ を生成し, 各EClassに対して $B_c$ を生成, ルート: $B_r$
  • $F(op)$ は $B_{op}$ が選択されるならば子も選択する
  • $G(c)$ ソルバがEClassを選択するときメンバの少なくとも1つを選択しなければならない
• これで入力式との等価性を表せる
• ILP Solverとして Gurobi

$$\begin{array}{rl}\mathrm{Constraints}& \equiv B_r\wedge \underset{op}{\bigwedge }F(op)\wedge \underset{c}{\bigwedge }G(c)\\ F(op)& \equiv B_{op}\to \underset{c\in op.children}{\bigwedge }{B}_c\\ G(c)& \equiv B_c\to \underset{op\in c.nodes}{\bigvee }{B}_{op}\end{array}$$ $$\begin{gathered} \min \sum _{op}{B}_{op}\cdot C_{op} \\ s.t\mathrm{Constraints} \end{gathered}$$

4. EVALUATION

2022-11-02

• System ML, Tensorflow関数をXLAに変換

4.1 Completeness of Relational Rules

• w.r.t: に関して

• rules

  • UnnecessaaryOuterProduct: x * (y @ 1) -> x * y if colvec?(y)
  • ColwiseAgg: colsums(x) -> sum(x) if colvec?, x if rowvec?