880

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

June 08, 2022

## Transcript

2. ### どんな論文？ • NPC（Huawei Ascend 910）向け グラフコンパイラの実装のレポート – 既出の技術をどう組み合わせたのか・どう改良 したのか詳しく記載されている –

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

Engine （Img2colできる • TPUのような階層型メモリ構造

を実施
5. ### おさらい：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
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 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
7. ### おさらい：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の係数が求まりそう！

9. ### 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

11. ### Tiling • tile shapeをどうやって決めるか – reverse strategy （筆者らの提案）に従う • conservative

clustering によりtiling可能性を高める • Intermediate なノードを見つける – トップレベルのSequence ノードの Filter ノード • Tlingは Live-out iteration space だけを対象とする – 計算が終わったあとに参照される、メモリに書き出す Stmtを対象

14. ### Tiling • tile sizeをどうやって決めるか – 従来のコンパイラはデフォルト値決め打ち だった • TVMはtile sizeを指定できる

• 込み入ったアーキに対してより洗練された方法が 必要 – Tile-Size Specification Languageを新たに作る • Auto-Tuningも実装：データの移動についてヒュー リスティックなコストモデルを作成し、Greedy Search

16. ### Fusion 2 • dot-productをconvolution演算とみなし、 Cube-unitを使うものはL1、Vector/Scala UnitつかうものはUBラベルを付ける • Post-Tiling Fusionを行う –

従来のグラフコンパイラでは不可能らしい （？） – Scheduling TreeにExtension Nodeを追加し追加 情報を付加

20. ### コード生成 • スケジューラツリーはHalideIRに変換 • Ascned 910の機能をフルに使える、CCEと 呼ばれるC-like な言語に変換 • ベクトル化はislが処理する範囲ではないが、

Polyhedral IRのうちに実施 – ここがISLの弱点だね、と言っている

22. ### 特定ハードウエア向け最適化 • 多くのNPCは計算ユニットとデータ転送ユニッ トが独立なので、データ依存を守るために Syncが必要 • メモリ遅延隠蔽はPolyhedral 変換のあとにやる が、TVMはheterogeneous 計算機を考慮しない

• 最初に異なる計算ユニット間ではsyncを挿入 – その後syncを省いていく
23. ### 自動チューニング戦略 • 機械学習ベースの性能チューニング – Tiling Sizeなど探索が必要なパラメータあり • 最初はランダムにサンプリングし、2番目 のフェーズとしてN-Best Performing

Sample を確立pでサンプルし1-pの確率でランダム にパラメータを与え計測
24. ### 評価 • 比較対象 – ﾆﾝｼﾞｬが手動で最適化したコード（CCE opt） – TVMにDaVinciアーキの情報を移植したもの – TVMの手動スケジューリングテンプレートで

自動チューニングさせたもの
25. ### 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

29. ### 読んだ感想 • すごく新しいアイディアが提案されている わけではないが、System Descriptionのよう にTensor Compilerを作る上で必要な技術が 詳細に記述されていた • 技術詳細は今までの論文をリファーして

いる感じなので詳しくはほか読んでとい う感じ