$30 off During Our Annual Pro Sale. View Details »

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

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

    View Slide

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

    View Slide

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

    View Slide