https://mlxse.connpass.com/event/102910/

1. Diesel: DSL for Linear Algebra and
Neural Net Computations on GPUs
XX for ML 論文読み会 #1
@no_maddo
注：本文中の図はURLが注記されているものを除き、すべて本論文からのものです

2. 私は誰
• コンパイラとか好きです

3. 論文概要
• MAPL 2018 (PLDIの併設ワークショップ）
• DSLをPolyhedralを用いたコンパイラ最適化で
NVIDIA GPU向け数値計算カーネルを高速化

4. 発表の構成
• モチベーションの説明
• 技術解説
– Polyhedral Model
– Diesel Compilerのながれ
• Scheduler
• 実験結果
• ほか研究との比較

5. 発表の構成
• モチベーションの説明
• 技術解説
– Polyhedral Model
– Diesel Compilerのながれ
• Scheduler
• 実験結果
• ほか研究との比較

6. NN向けコンパイラ技術研究の
モチベーション
• 現在の主流：GPUの高速な線形代数カーネル
はアセンブラで手書きれている
• 複数の演算を跨いだ最適化は重要だが、
この組み合わせは膨大
– 演算を融合するとグローバルメモリへの読み書き
が減らせる場合
• 全自動で高速な実装を手に入る方法が必要

7. 論文でやろうとしていること
• Diesel言語から自動的にNVIDIA GPU向けの
高速なカーネルを生成すること
• 簡単にパラメータをチューニングできること
• 高速な演算が融合されたカーネルを生成
できること

8. 発表の構成
• モチベーションの説明
• 技術解説
– Polyhedral Model
– Diesel Compilerのながれ
• Scheduler
• 実験結果
• ほか研究との比較
– TensorCompression

9. Polyhedral Model: 1
• プログラム最適化のフレームワーク
• Static Control Partsに対して最適化可能
– forループに定数ストライド
– if文
– loop-boundがaffine

10. Polyhedral Model: 2
• 記法：
– Set: [n] -> {S[i]: 0 <= i < n}
– Map: [n] -> {S1[i] -> S2[i]: 0 <= i, j<n and j=i+1}
定義に必要な
パラメータのリスト １次元の集合

11. Polyhedral Model: 3
for (i=0; i<N; i++)
for (j=0; j<N; j++)
for (k=0; k<N; k++)
S1: C[i][j] += A[i][k] * B[k][j]
• Iteration domain
– 𝐷𝑆1 ≔ 𝑁 → {𝑆1 𝑖, 𝑗, 𝑘 : 0 ≤ 𝑖, 𝑗, 𝑘 < 𝑁}
• Access function
– ≔ 𝑁 → {𝑆1 𝑖, 𝑗, 𝑘 → 𝐴 𝑖, 𝑘 : 0 ≤ 𝑖, 𝑗, 𝑘 < 𝑁}

12. Polyhedral Model: 4
for (i=0; i<N; i++)
for (j=0; j<N; j++)
for (k=0; k<N; k++)
S1: C[i][j] += A[i][k] * B[k][j]
• Scheduling Function
– Sのインスタンス𝑠1が𝑠2に先行 𝑖𝑓𝑓 𝜙 𝑆 𝑠1 < 𝜙 𝑆 𝑠2 ;
辞書的で比較
– 𝜙 𝑆1 ≔ 𝑁 → {𝑆1 𝑖, 𝑗, 𝑘 → 𝑖, 𝑗, 𝑘 ∶ 0 ≤ 𝑖, 𝑗, 𝑘 < 𝑁}
• Data dependencies
– ≔ 𝑁 → {𝑆1 𝑖, 𝑗, 𝑘 →
𝑆1 𝑖, 𝑗, 𝑘 + 1 : 0 ≤ 𝑖, 𝑗 < 𝑁 𝑎𝑛𝑑 0 ≤ 𝑘 < 𝑁 − 1 }

13. Dieselの処理の流れ
• Diesel言語をパース
• DAGに変換
• Polyhedral IRに変換
• 各アーキごとに変換

14. Dieselの処理の流れ
• Diesel言語をパース
– 宣言
– 演算
– CodeGen宣言
• DAGに変換
• Polyhedral IRに変換
• 各アーキごとに変換

15. Dieselの処理の流れ
• Diesel言語をパース
• DAGに変換
– Nodeに形状情報を付加
• dataset
• datamap
• Polyhedral IRに変換
• 各アーキごとに変換

16. DAG: Datasets
• それぞれのノードの形状、polyhedron
– Datasetsと呼ぶ
– 例：A(M, K) は
𝑀, 𝐾 → 𝐴 𝑖, 𝑗 : 0 ≤ 𝑖 < 𝑀 𝑎𝑛𝑑 0 <= 𝑗 < 𝐾

17. DAG: Datamaps
• datasetsのメモリレイアウトを表すAffile関数
– 例えば式𝐴と𝐴 𝑇
は同じメモリ空間を使い回せる
Datamap 𝐴 𝑇
𝑖, 𝑗 → 𝐴 𝑗, 𝑖 で表す
– Broadcastにも使う

18. Dieselの処理の流れ
• Diesel言語をパース
• DAGに変換
• Polyhedral IRに変換
– Iteration Domainを計算
• 同じく形状情報を伝搬させて計算
– Access functionを求める
• Iteration Vector[i,j,k]のときにRead/Writeでどこを触るか
– データ依存関係を解析
• 各アーキごとに変換

19. Scheduler 1
• PlutoのAlgorithmで最初のスケジュール
– ある種の“距離”を最小化し、Data再利用を促進
• Tiling
– データ局所性改善のため
– threadとthreadblockの分散のため
• Memory Promotion
– global memoryを使っているところを
shared memory, registerに変更
• Software Pipelining

20. Scheduler 2
• Software Pipelining
– ループを跨いだスケジューリング
– load_to_shared(); sync(); compute(); sync()
の繰り返しを展開して無駄なsyncをへらす
• Code Generation
– 最後にschedule treeでスケジューリング
– schedule tree:
ASTからコード生成するときに凸多面体解析の
情報を木にして扱う手法

21. http://www.pollylabs.org/slides/polyhedral-AST-generation.pdf

22. 発表の構成
• モチベーションの説明
• 技術解説
– Polyhedral Model
– Diesel Compilerのながれ
• Scheduler
• 実験結果
• ほか研究との比較

23. 実験結果： SGEMM/DGEMM

24. 実験結果: GEMM + bias + ReLU
• 90%以上で
baseline超え

25. 関連研究
• アセンブラ手書きカーネル
– CuBlas/CuDNN/OpenAI
• PPCG
– 凸多面体解析を用いるC to CUDAコンパイラ
– これはより一般のプログラムが対象
• CUTLASS
– NVIDIA, 線形代数の計算を翻訳するC++ template
• Tensor Comprehension
– Facebook, ネットワーク全体が最適化対象

26. 読んだ感想
• “堅い”技術で作られたコンパイラ
– Workshop paperだしそんなに新規性はなさそう？
• Tensor Compressionと共通の部分が多い
– Range Inferenceして最後は Schedule Tree
• Kernel fusionの結果は驚き
– bias, ReLUの融合はいかにもメモリ帯域を節約で
きる

27. Kernel Fusionに対する言及少ない
• Plutoで融合できたら融合するくらいの話？
• [1][2]はもっとKernel Fusionを真面目に扱ってる
– 線形代数演算の融合対象の候補探索アルゴリズム

28. まとめ
• NVIDIA GPU向けに凸多面体解析を用いた
コンパイラ最適化を行うDSLコンパイラを提案
• 実際に手書きアセンブラのカーネルを超える
場合もある
– 融合演算の結果は驚くべき結果

29. 参考文献
• [1] On Optimizing Operator Fusion Plans for LargeScale
Machine Learning in SystemML
• [2] On Optimizing Machine Learning Workloads via Kernel
Fusion
• 本論文で引用されていたものは省略しました