Upgrade to Pro — share decks privately, control downloads, hide ads and more …

PLDI '21論文読み会: AKG: Automatic Kernel Generation for Neural Processing Units using Polyhedral Transformations

Idein
June 08, 2022

PLDI '21論文読み会: AKG: Automatic Kernel Generation for Neural Processing Units using Polyhedral Transformations

Idein

June 08, 2022
Tweet

More Decks by Idein

Other Decks in Research

Transcript

  1. AKG: Automatic Kernel Generation for Neural Processing Units using Polyhedral

    Transformations PDLI 2021 論文読み会 @nomaddo
  2. どんな論文? • NPC(Huawei Ascend 910)向け グラフコンパイラの実装のレポート – 既出の技術をどう組み合わせたのか・どう改良 したのか詳しく記載されている –

    System Descriptionという感じ • NPU向けのスケジューリング手法 – メモリ階層 – 特殊なデータフローを持つハード – NPU向けハード固有最適化が今までの手法で サポートとされず
  3. Ascend のDaVinci アーキテク チャ • 特殊なハードウエア – MTE: Memory transfer

    Engine (Img2colできる • TPUのような階層型メモリ構造
  4. おさらい:Polyhedral optimization • 2つの使い方:ループから制約を取り出し て – Analysis: ループ変形が可能かどうかチェック – Scheduling:

    制約を満たす変形を最適化問題に して見つける for i in [0, 5]: P: A[i] = input[i] + 1 for j in [0, 5]: C: B[j] = A[j] * 2
  5. おさらい:Polyhedral optimization のSchedulingの見つけ方 for i in [0, 5]: P: A[i]

    = input[i] + 1 for j in [0, 5]: C: B[j] = A[j] * 2 Read After Writeの制約 - P[i] <= C[j] (C[i] が P[j] のあとに実行) Iteration Domain - 0 <= i, j <= 5 ∀ 0 <= i, j <= 5 && i == j. NewSched_P(i) <= NewSched_C(j) NewSched_P(i) = sp * i + dp NewSched_C(j) = sc * j + dc とおくと ∀ 0 <= i, j <= 5 && i == j && sp * i + dp <= sc * j + dc
  6. おさらい:Polyhedral optimization のSchedulingの見つけ方 for i in [0, 5]: P: A[i]

    = input[i] + 1 for j in [0, 5]: C: B[j] = A[j] * 2 ∀ 0 <= i, j <= 5 && i == j && sp * i + dp <= sc * j + dc ここでLocalityを考えて、w の最小化問題を考えてみる sp * i – sc * j + dp - c <= w ILPの制約問題で、Locality を最大にするSchedulingの係数が求まりそう!
  7. AKG最初の変形 - Bias を足す - 行列積 - 絶対値を取る - ReLUを計算

    for h in [0, H), w in [0, W): A[h, w] += bias // S0 for h in [0, H-KH], w in [0, W–KW]: C[h, w] = 0 // S1 for kh in [0, KH), kw in [0, KW): C[h, w] += A[h+kh,w+kw]*B[kh,kw] // S2 for h in [0, H-KH], w in [0,W–KW]: C[h, w] = abs(C[h,w]) // S3 for h in [0, H-KH], w in [0,W–KW]: C[h, w] = ReLU(C[h,w]) // S4
  8. Tiling • tile shapeをどうやって決めるか – reverse strategy (筆者らの提案)に従う • conservative

    clustering によりtiling可能性を高める • Intermediate なノードを見つける – トップレベルのSequence ノードの Filter ノード • Tlingは Live-out iteration space だけを対象とする – 計算が終わったあとに参照される、メモリに書き出す Stmtを対象
  9. Tiling • tile sizeをどうやって決めるか – 従来のコンパイラはデフォルト値決め打ち だった • TVMはtile sizeを指定できる

    • 込み入ったアーキに対してより洗練された方法が 必要 – Tile-Size Specification Languageを新たに作る • Auto-Tuningも実装:データの移動についてヒュー リスティックなコストモデルを作成し、Greedy Search
  10. Fusion 2 • dot-productをconvolution演算とみなし、 Cube-unitを使うものはL1、Vector/Scala UnitつかうものはUBラベルを付ける • Post-Tiling Fusionを行う –

    従来のグラフコンパイラでは不可能らしい (?) – Scheduling TreeにExtension Nodeを追加し追加 情報を付加
  11. Single Operator • op1: Convolution • op2: matrix maliplication •

    op3: ReLU • op4: batched matrix multiplication • op5: cast • op6: transpose • op7: one-hot • op8: tensor addition • op9: BatchNormed Training reduction • op10: BatchNormed training update