Quartz 5

Vectorization for Digital Signal Processors via Equality Saturation

Properties2
tagspaper/egraph, fhe/vectorize
aliases

7 min read

Vectorization for Digital Signal Processors via Equality Saturation

egraph vectorize

@inproceedings{vanhattum2021vectorization,
  title={Vectorization for digital signal processors via equality saturation},
  author={VanHattum, Alexa and Nigam, Rachit and Lee, Vincent T and Bornholt, James and Sampson, Adrian},
  booktitle={Proceedings of the 26th ACM International Conference on Architectural Support for Programming Languages and Operating Systems},
  pages={874--886},
  year={2021},
  url={https://www.cs.cornell.edu/~avh/diospyros-asplos-2021-preprint.pdf}
}

Equality Saturation を用いて, 自動ベクトル化を行う

言語定義とか最適化フローとか参考に出来る部分たくさんありそう

実装: GitHub - cucapra/diospyros: Search-based compiler for high-performance DSP programming

Abstract

  • デジタル信号処理(DSP)のアプリケーションでは線形代数計算の高速な実装が重要
  • 最適化のための最適なデータ移動を考えるのは大変
  • Diospyros は効率的なベクトル化を自動で見つけるコンパイラ
  • avg x3.1 早い

1. Introduction

  • DSPアプリケーションは
    • ベクトル幅より遥かに大きいデータ並列化
    • ベクトル幅に収まる小さな並列化
  • コンパイルをプログラムの候補空間における探索問題とした　
    • シャッフル命令やセレクト命令を用いて不規則な最適化を行う
  • コンパイルフロー
    • scalarなプログラムをvector DSLに変換
    • Equality Saturation でDSLを最適化
    • lowering
  • 貢献

2. Motivating Example

  • 従来の自動並列化は境界条件が並列化に対する弊害になっていた
  • 逐次コードをベクトル化するための一般的なコンパイラ技術
  • 2dconvでは if が並列化の邪魔に
    • 従来なら定数回数ループなのでunrollしてSLP
int indices[4] = {1, 2, 0, 5};
xb_vecMx32 vec3 = PDX_SEL_MX32(vec1, vec2, indices);
xb_vecMx32 vec4 = PDX_MUL_MX32(vec1, vec3);
  • 例えば vec_multiply_accumulate を使用して, 各要素の積にfusionされる
(VecMAC (...)
(Vec (Get I 6) (Get I 7) (Get I 8) (Get I 9))
(Vec (Get F 0) (Get F 0) (Get F 0) (Get F 0)))
  • + と * の可換性と結合性により可能なshuffleはたくさんある

3. Rewriting for Vectorization

3.1 Defining and Lifting Specifications

  • RacketDSL: 行列, ベクトルと演算があるスカラープログラム

  • ベクトルの要素和の例
(define (vector-add-spec A B n)
  (vec-decl 'A n 'input)
  (vec-decl 'B n 'input)
  (define C (make-vector n))
  (for ([i n])
    (vector-set! C i
      (add (vector-ref A i)
           (vector-ref B i))))
  C)

  • 探索を単純化するためにまず入力プログラムを “数学的な表現” (???) に変換する

  • RacketDSL の拡張である Rosette で記号評価をする

  • n=2 のベクトル要素和の変換例

(List
(+ (Get a 0) (Get b 0))
(+ (Get a 1) (Get b 1)))

3.2 Rewriting Strategy

  • 固定サイズのベクトルの連結とみなす
  • ex. でベクトル幅が 2
(List (+ (Get a 0) (Get b 0))
(+ (Get a 1) (Get b 1))
(+ (Get a 2) (Get b 2))
(+ (Get a 3) (Get b 3)))
⇝
(Concat 
  (Vec (+ (Get a 0) (Get b 0))
       (+ (Get a 1) (Get b 1)))
  (Vec (+ (Get a 2) (Get b 2))
       (+ (Get a 3) (Get b 3)))
)

(Vec (+ a b) (+ c d)) ⇝ (VecAdd (Vec a c) (Vec b d))
  • チャンク化することでこのルールが適用できる

3.3 Searching for Rewrites

  • ex. fused multiply-accumulate
(VecAdd a (VecMul b c)) ↭ (VecMAC a b c)

  • レーンが空でも動作する書き換えルールを用意している

    • ex. ベクトル幅4で 3x3 の行列演算を行う時等

  • egg のカスタム書き換え規則を作った

    • (vec ...) で ... が <op> ... or 定数かをチェック

  • 可換性, 結合性を適用すると E-Graph のサイズが劇的に大きくなる

    • 2項が可換性, 結合性において同値か(AC-matching problem)はNP完全
    • なので, 無効に出来る

  • 例: 4要素のVec

(Vec (+ a0 (* b0 c0))
(+ a1 (* b1 c1))
(+ a2 (* b2 c2))
(+ (* b3 c3) a3))
  • 4要素目は可換性が無いと統一出来ない
    • 対策として各レーンで独立にマッチングしてその結果を統合するCustom Searcherを使用している
  • ルールをapplyする度に可換性, 結合性を再計算する
    • 大規模カーネルで512GB使い果たしていたのが改善した

3.4 Extraction

  • DSLの各演算子に固定コストを割り当てたコストモデルを用いて抽出
    • 効率的に(状態数ではなくE-graphノード数に比例)するためにはコストモデルが単調でなければならない
  • データ移動のコストは高め -> シャッフルを使うようになった
  • オプションで検証のために生成されたプログラムが全ての入力で正しいことをチェックする

4. Lowering & Code Generation

  • IRに変換
    • シャッフル: インデックスをどこに移動させるか
    • ループネストを展開する
      • 冗長な項を排除するためにLVN(Local Value Numbering)パスを実装している
        • Quarterinoベンチでは10万+行から500-行になった

5. Evaluation

  • 4.8k loc Racket

  • 1.4k loc Rust

  • egg

    • timeout: 180s
    • node limit: 10000000

  • benchmark

    • 2Dconv
    • MatMul
    • Quaternion product
    • QRDecomp

6. Limitations & Portablitiy

Graph View

Backlinks