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

5148cbb852e980b3b3d0e71d579ddbf8?s=47 Idein
June 08, 2022

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

5148cbb852e980b3b3d0e71d579ddbf8?s=128

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. 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 • 2つ問題がある – 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. サブグラフでの評価 • Resnet50,Bert, Mobilenetのsubgraphで 評価

  28. End-to-Endネットワークでの評価 • CCE optはResnet50向けのみライブラリ提供 がある

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

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