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
   

   View Slide

2. どんな論文？
   • NPC（Huawei Ascend 910）向け
   グラフコンパイラの実装のレポート
   – 既出の技術をどう組み合わせたのか・どう改良
   したのか詳しく記載されている
   – System Descriptionという感じ
   • NPU向けのスケジューリング手法
   – メモリ階層
   – 特殊なデータフローを持つハード
   – NPU向けハード固有最適化が今までの手法で
   サポートとされず
   

   View Slide

3. Ascend のDaVinci アーキテク
   チャ
   • 特殊なハードウエア
   – MTE: Memory transfer Engine （Img2colできる
   • TPUのような階層型メモリ構造
   

   View Slide

4. AKGの大枠
   • TVMを改造
   – スケジューラ は isl
   を用いる
   • 最初にFusion/Tiling
   を実施
   

   View Slide

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
   

   View Slide

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
   

   View Slide

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の係数が求まりそう！
   

   View Slide

8. isl
   • PPCGというSource-to-Souce コンパイラの
   内部で使われたスケジューラ
   • Schedule Trees　という形式を出力
   – ループ変化はこの形式で扱う
   

   View Slide

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
   

   View Slide

10. Tiling
    • ２つ問題がある
    – tile shapeをどうやって決めるか
    – tile sizeをどうやって決めるか
    

    View Slide

11. Tiling
    • tile shapeをどうやって決めるか
    – reverse strategy （筆者らの提案）に従う
    • conservative clustering によりtiling可能性を高める
    • Intermediate なノードを見つける
    – トップレベルのSequence ノードの Filter ノード
    • Tlingは Live-out iteration space だけを対象とする
    – 計算が終わったあとに参照される、メモリに書き出す
    Stmtを対象
    

    View Slide

12. conservative clusteringの実行例
    

    View Slide

13. Live-out iteration Spaceをタイル
    化
    

    View Slide

14. Tiling
    • tile sizeをどうやって決めるか
    – 従来のコンパイラはデフォルト値決め打ち
    だった
    • TVMはtile sizeを指定できる
    • 込み入ったアーキに対してより洗練された方法が
    必要
    – Tile-Size Specification Languageを新たに作る
    • Auto-Tuningも実装：データの移動についてヒュー
    リスティックなコストモデルを作成し、Greedy
    Search
    

    View Slide

15. Fusion 1
    • Loop-FusionはPolyhedral compilationの一つ
    として実装可能だが、従来のものは複数
    方向のデータやり取りをするアーキを考
    慮せず
    

    View Slide

16. Fusion 2
    • dot-productをconvolution演算とみなし、
    Cube-unitを使うものはL1、Vector/Scala
    UnitつかうものはUBラベルを付ける
    • Post-Tiling Fusionを行う
    – 従来のグラフコンパイラでは不可能らしい
    （？）
    – Scheduling TreeにExtension Nodeを追加し追加
    情報を付加
    

    View Slide

17. Fusion 3
    • Intra-tile reschedulingを導入
    – UBラベルsubtreeのみLoop Distributionを行い
    ベクトル化する
    

    View Slide

18. Convの最適化
    • DaVinchiアーキにはMTEがあるのでim2col
    をハードがやってくれる
    • 今までのTensor Compilerはim2col自体は扱
    おうとしていなかった
    

    View Slide

19. Manual Scheduling and Debugging
    • 複雑なメモリ階層を持つアーキのために
    メモリ階層記述言語を定義
    • 手動最適化の余地があるが、ベンチでは
    使っていない
    

    View Slide

20. コード生成
    • スケジューラツリーはHalideIRに変換
    • Ascned 910の機能をフルに使える、CCEと
    呼ばれるC-like な言語に変換
    • ベクトル化はislが処理する範囲ではないが、
    Polyhedral IRのうちに実施
    – ここがISLの弱点だね、と言っている
    

    View Slide

21. ベクトル化
    • 最適化されたHalideIRを受け取り、それに
    SIMD化に必要な全ての情報がある
    • ベクトル化の際にデータレイアウトの最
    適化をする余地あり
    • 一緒にLoop Unrollingもやる
    

    View Slide

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

    View Slide

23. 自動チューニング戦略
    • 機械学習ベースの性能チューニング
    – Tiling Sizeなど探索が必要なパラメータあり
    • 最初はランダムにサンプリングし、2番目
    のフェーズとしてN-Best Performing Sample
    を確立pでサンプルし1-pの確率でランダム
    にパラメータを与え計測
    

    View Slide

24. 評価
    • 比較対象
    – ﾆﾝｼﾞｬが手動で最適化したコード（CCE opt）
    – TVMにDaVinciアーキの情報を移植したもの
    – TVMの手動スケジューリングテンプレートで
    自動チューニングさせたもの
    

    View Slide

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
    

    View Slide

26. Single Operatorの開発労力
    • 手動で最適化・TVMの手動テンプレートを
    描くのと比較してコード量が少なくてす
    む
    

    View Slide

27. サブグラフでの評価
    • Resnet５０，Bert, Mobilenetのsubgraphで
    評価
    

    View Slide

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

    View Slide

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

    View Slide