AKG: Automatic Kernel Generation for Neural Processing Units using Polyhedral Transformations PDLI 2021 論文読み会 @nomaddo

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

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

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

おさらい：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

おさらい：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

おさらい：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の係数が求まりそう！

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

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

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

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

conservative clusteringの実行例

Live-out iteration Spaceをタイル 化

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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