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

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. AKGの大枠 • TVMを改造 – スケジューラ は isl を用いる • 最初にFusion/Tiling を実施
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の係数が求まりそう！
8. isl • PPCGというSource-to-Souce コンパイラの 内部で使われたスケジューラ • Schedule Trees　という形式を出力 – ループ変化はこの形式で扱う
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
10. Tiling • ２つ問題がある – tile shapeをどうやって決めるか – tile sizeをどうやって決めるか
11. Tiling • tile shapeをどうやって決めるか – reverse strategy （筆者らの提案）に従う • conservative clustering によりtiling可能性を高める • Intermediate なノードを見つける – トップレベルのSequence ノードの Filter ノード • Tlingは Live-out iteration space だけを対象とする – 計算が終わったあとに参照される、メモリに書き出す Stmtを対象
12. conservative clusteringの実行例
13. Live-out iteration Spaceをタイル 化
14. Tiling • tile sizeをどうやって決めるか – 従来のコンパイラはデフォルト値決め打ち だった • TVMはtile sizeを指定できる • 込み入ったアーキに対してより洗練された方法が 必要 – Tile-Size Specification Languageを新たに作る • Auto-Tuningも実装：データの移動についてヒュー リスティックなコストモデルを作成し、Greedy Search
15. Fusion 1 • Loop-FusionはPolyhedral compilationの一つ として実装可能だが、従来のものは複数 方向のデータやり取りをするアーキを考 慮せず
16. Fusion 2 • dot-productをconvolution演算とみなし、 Cube-unitを使うものはL1、Vector/Scala UnitつかうものはUBラベルを付ける • Post-Tiling Fusionを行う – 従来のグラフコンパイラでは不可能らしい （？） – Scheduling TreeにExtension Nodeを追加し追加 情報を付加
17. Fusion 3 • Intra-tile reschedulingを導入 – UBラベルsubtreeのみLoop Distributionを行い ベクトル化する
18. Convの最適化 • DaVinchiアーキにはMTEがあるのでim2col をハードがやってくれる • 今までのTensor Compilerはim2col自体は扱 おうとしていなかった
19. Manual Scheduling and Debugging • 複雑なメモリ階層を持つアーキのために メモリ階層記述言語を定義 • 手動最適化の余地があるが、ベンチでは 使っていない
20. コード生成 • スケジューラツリーはHalideIRに変換 • Ascned 910の機能をフルに使える、CCEと 呼ばれるC-like な言語に変換 • ベクトル化はislが処理する範囲ではないが、 Polyhedral IRのうちに実施 – ここがISLの弱点だね、と言っている
21. ベクトル化 • 最適化されたHalideIRを受け取り、それに SIMD化に必要な全ての情報がある • ベクトル化の際にデータレイアウトの最 適化をする余地あり • 一緒にLoop Unrollingもやる
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
26. Single Operatorの開発労力 • 手動で最適化・TVMの手動テンプレートを 描くのと比較してコード量が少なくてす む
27. サブグラフでの評価 • Resnet５０，Bert, Mobilenetのsubgraphで 評価
28. End-to-Endネットワークでの評価 • CCE optはResnet50向けのみライブラリ提供 がある
29. 読んだ感想 • すごく新しいアイディアが提案されている わけではないが、System Descriptionのよう にTensor Compilerを作る上で必要な技術が 詳細に記述されていた • 技術詳細は今までの論文をリファーして いる感じなので詳しくはほか読んでとい う感じ