1 of 43

PFN で開発している深層学習

コンパイラスタックについて

浜地慎一郎

MLSE夏合宿2021

2 of 43

自己紹介

  • Google (2007-2018)
    • クラスタ
    • ブラウザ
    • 分散コンパイラ
    • Android
    • 機械翻訳
  • Preferred Networks (2018-)
    • デプロイと MN-Core

3 of 43

Agenda

  • Preferred Networks の紹介
    • いろいろやってる会社
  • PFVM
    • 推論もトレーニングもサポートするグラフコンパイラ
  • MN-Core
    • MN-Core 向けソフトウェアスタック
    • 今日はこの話が多めです

4 of 43

Preferred Networks

5 of 43

Preferred Networks

Founded

March 2014

Directors

CEO　Toru Nishikawa

COO　Daisuke Okanohara

CTO　Ryosuke Okuta

Located

Tokyo, Japan (HQ) ​

Burlingame, CA., US

(Preferred Networks America, Inc.)​

Number of Employees

Over 300 Engineers & Researchers

(April, 2021)​

Manufacturing

Logistics

Transportation

Bio & Healthcare

Personal Robot

Entertainment

6 of 43

PFN のおもしろポイント

Founded

March 2014

Directors

CEO　Toru Nishikawa

COO　Daisuke Okanohara

CTO　Ryosuke Okuta

Located

Tokyo, Japan (HQ) ​

Burlingame, CA., US

(Preferred Networks America, Inc.)​

Number of Employees

Over 300 Engineers & Researchers

(April, 2021)​

Manufacturing

Logistics

Transportation

Bio & Healthcare

Personal Robot

Entertainment

決して小さな会社ではない

7 of 43

PFN のおもしろポイント

Founded

March 2014

Directors

CEO　Toru Nishikawa

COO　Daisuke Okanohara

CTO　Ryosuke Okuta

Located

Tokyo, Japan (HQ) ​

Burlingame, CA., US

(Preferred Networks America, Inc.)​

Number of Employees

Over 300 Engineers & Researchers

(April, 2021)​

Manufacturing

Logistics

Transportation

Bio & Healthcare

Personal Robot

Entertainment

そうは言っても

事業が多すぎる！

8 of 43

PFN のおもしろポイント

  • 確かに人数は多いが、事業も多い
    • プロジェクトあたりの人数は決して多くない
    • GraphCore の人「うちらDNNチップ事業だけだけど、お前らと同じくらい人数いるんだが……」
  • 事業間の連携/シナジーがなければ
    • 単に小さなベンチャーの集合になってしまう
  • →計算インフラや基本的なライブラリ整備� 事業間のノウハウの交換などを重視

9 of 43

PFVM

10 of 43

PFVM: Versatile graph compiler

  • モデルが決まった後の最適化 & 実行は全部やる
    • (……というのは心意気で、現実は若干厳しいですが)
    • まあ「エッジからスパコンまで」ではあります
  • デプロイ
    • 社内のデプロイ案件はなるべく広くサポートできるように
      • 客先デスクトップ PC の GPU からロボットのチップまで
  • MN-Core
    • トレーニング用途メインなので、自動微分も完備
  • かなり広いユースケースをカバーしている

11 of 43

PFVM architecture

  • Box: software component / Rounded corner: data

Torch

JIT

PFVM Runtime

pytorch

ONNX

inference

PFVM IR

ONNX+

training

Compiler�code gen

MN-Core�IR

Compiler

middle

end

libtorch

+α

OpenVINO

TensorRT

etc.

torch.onnx を使って Python を ONNX に変換

12 of 43

PFVM architecture

  • Box: software component / Rounded corner: data

Torch

JIT

PFVM Runtime

pytorch

ONNX

inference

PFVM IR

ONNX+

training

Compiler�code gen

MN-Core�IR

Compiler

middle

end

libtorch

+α

OpenVINO

TensorRT

etc.

PFVM が最適化/グラフ変換をする

13 of 43

PFVM architecture

  • Box: software component / Rounded corner: data

Torch

JIT

PFVM Runtime

pytorch

ONNX

inference

PFVM IR

ONNX+

training

Compiler�code gen

MN-Core�IR

Compiler

middle

end

libtorch

+α

OpenVINO

TensorRT

etc.

ここを切り離して Python 無しでデプロイ

14 of 43

PFVM architecture

  • Box: software component / Rounded corner: data

Torch

JIT

PFVM Runtime

pytorch

ONNX

inference

PFVM IR

ONNX+

training

Compiler�code gen

MN-Core�IR

Compiler

middle

end

libtorch

+α

OpenVINO

TensorRT

etc.

ここを切り離して Python 無しでデプロイ

15 of 43

Philosophy

  • Python コードの再実装を必要としない
    • 前処理後処理などのモデルの外側も再実装必要としない
  • Flexibility
    • 動的 shape やループ、自動微分などの仕組み
  • 全てを自分たちでやろうとしない
    • 他社フレームワークにうまく乗っかる
    • 難しい計算カーネルを書くのは基本 MN-Core に対してだけ
  • 社内需要に注力

16 of 43

PFVM のおもしろポイント

  • ミドルエンド IR として ONNX をそのまま利用
    • 他社のグラフコンパイラは普通、入口で自前 IR に変換する
    • netron などの ONNX のエコシステムがそのまま利用できる
  • グラフを partitioning して他社 FW に丸投げする
    • 内部 IR が ONNX だからこそ出来ること
    • 他社 FW の ONNX の対応状況などに応じてグラフ変換
  • マイナーな処理などは libtorch などで
    • libtorch は普通に速いので、これらが律速することは少ない

17 of 43

MN-Core

18 of 43

PFVM architecture

  • Box: software component / Rounded corner: data

Torch

JIT

PFVM Runtime

pytorch

ONNX

inference

PFVM IR

ONNX+

training

Compiler�code gen

MN-Core�IR

Compiler

middle

end

libtorch

+α

OpenVINO

TensorRT

etc.

19 of 43

MN-Core とは

  • ディープラーニングに特化
  • 半精度: 524 TFLOPS
  • すごい幅の SIMD プロセッサ
    • 1命令で数百万回浮動小数点計算したりする
  • SRAM がたくさんある

20 of 43

Green 500 にて1位！

  • 2020年6月と2021年6月に1位
    • 2020年11月は惜しくも2位
  • 超省電力設計の Proof of Concept
  • 倍精度でも高効率

21 of 43

最近のアナウンス

22 of 43

3つのワークロードで優位性を確認

  • Detection/Instance Segmentation
    • 実用的/複雑なワークロード
  • Graph Conv
    • 当初の設計(CNN)外の動作
  • NAS
    • パラメータに対する柔軟性

23 of 43

発表にエンコードされてる気持ち

  • Detection/Instance Segmentation (training)
    • 実用的/複雑なワークロードで動作している
      • 逆に言うと classification は事業的重要度低め&簡単
      • 今回最大のモデルは 42 種類 10000 個程度の ONNX op
        • これにさらに backward も追加される

24 of 43

発表にエンコードされてる気持ち

  • Graph Convolution (inference)
    • 本来 CNN に向けて設計されていた MN-Core が、他の DNN ワークロードでも適用可能であることを示した
    • Conv/Gemm ヘビーでなかったりする
  • Neural Architecture Search (training)
    • 探索中に発見される広い範囲のハイパーパラメータに対応
    • 特定パラメータへのチューニングではない

cf. Neural Architecture Searchを用いたセマンティックセグメンテーションモデルの探索

25 of 43

MN-Core のおもしろポイント

  • 普通のプロセッサが勝手にやることをやらない
    • 例えばキャッシュがないのでメモリ階層の移動は自力
    • 全てソフトウェアで制御できるので、うまくソフトを書けばとても高効率に動作できるように、注意深く設計されている
  • 遅いソフトウェアをとりあえず書くのが難しい
    • 練習として 4x4 の行列積を書くのに1週間かかったり……
  • 最速を目指すのであれば、むしろ書きやすい説もある
    • 例えば Halide や TVM は CPU/GPU キャッシュサイズに応じてループ回数の探索をしたりするが MN-Core では不要

26 of 43

MN-Core のおもしろポイント

  • 普通のプロセッサが勝手にやることをやらない
    • 汎用プロセッサでは失敗したアプローチだと考えられる
      • VLIW や明示的なスクラッチパッドへのコピーなど
        • IA-64, Cell, …
    • 普通のプログラムではハードウェアを賢くして、動的な情報を元に並列性を出すのが良い
  • ディープラーニングでは、定型の処理をとにかく高効率にすれば良いので、有効なアプローチとなる
    • TPU も歴史的な趣のある systolic array を生かしていて面白い

27 of 43

PFVM architecture

  • Box: software component / Rounded corner: data

Torch

JIT

PFVM Runtime

pytorch

ONNX

inference

PFVM IR

ONNX+

training

Compiler�code gen

MN-Core�IR

Compiler

middle

end

libtorch

+α

OpenVINO

TensorRT

etc.

この箱に PFVM の3倍くらいのコード量があったりする

28 of 43

l1ir

l2ir

l3ir

PFVM

PyTorch Code

l3ir graph

MN-Core�(executable)

ONNX

ONNX+training

Generic Conv

MNTensor

Generic Move

PEVector

Layers

MN-Core�Instructions

Graph simplification

Backprop generation

Code generation for MN-Core

Graph-to-instruction conversion

Instruction scheduling

Memory allocation

l1ir graph

Layout

Planner

l3ir Scheduler

scheduler

l1merge

Address

Planner

29 of 43

l1ir

l2ir

l3ir

PFVM

PyTorch Code

l3ir graph

MN-Core�(executable)

ONNX

ONNX+training

Generic Conv

MNTensor

Generic Move

PEVector

Layers

MN-Core�Instructions

Graph simplification

Backprop generation

Code generation for MN-Core

Graph-to-instruction conversion

Instruction scheduling

Memory allocation

l1ir graph

Layout

Planner

l3ir Scheduler

scheduler

l1merge

Address

Planner

謎の独自用語が

発明されまくっている！

30 of 43

自社開発部分が非常に多い

  • 非常に独自性の高いチップなので……
    • このチップの性能を出しきれるコンパイラスタックを考えていくと、やはり独自性の高いものにならざるを得なかった
  • LLVM などの汎用コンパイラスタックはほぼ自明に無用
  • TVM などの DNN 用コンパイラスタックは検討したが
    • TVM が頑張っている領域と我々が困っているところがあまりオーバーラップしない

31 of 43

レイアウトの話

  • 普通のCPUでのSIMDだと
    • NCHWc とかが良かったりする (c は SIMD 幅)
  • MN-Core の階層は、この c の軸が複数あるのに相当
    • n2h1w1h0w0CHWc�のような非常に複雑な�レイアウトを扱う
  • TensorFlow XLA の shape
    • 似ている気がするけど、たぶん�我々の要件の方が複雑に見える

32 of 43

レイアウトの話の例

shape=(12) を 4 分割の 1 階層に配置する例

​

​

​

​

PE0

PE1

PE2

PE3

0 1 2

3 4 5

6 7 8

9 10 11

PE0

PE1

PE2

PE3

0 4 8

1 5 9

2 6 10

3 7 11

33 of 43

レイアウトの話の例

shape=(4, 6) を 4 分割の 1 階層に配置する例

​

​

2軸を2分割ずつする、こういう分割もありえる

​

PE0

PE1

PE2

PE3

(0,0) (0,1) … (0,5)

(1,0) (1,1) … (1,5)

(2,0) (2,1) … (2,5)

(3,0) (3,1) … (3,5)

PE0

PE1

PE2

PE3

(0,0) (0,1) (0,2)

(1,0) (1,1) (1,2)

(0,3) (0,4) (0,5)

(1,3) (1,4) (1,5)

(2,0) (2,1) (2,2)

(3,0) (3,1) (3,2)

(2,3) (2,4) (2,5)

(3,3) (3,4) (3,5)

34 of 43

レイアウトの話の例

  • 単項 elementwise op (例えば ReLU) ではどちらも同じ
    • 各 PE が 6 回 ReLU をするだけ

​

​

​

​

PE0

PE1

PE2

PE3

(0,0) (0,1) … (0,5)

(1,0) (1,1) … (1,5)

(2,0) (2,1) … (2,5)

(3,0) (3,1) … (3,5)

PE0

PE1

PE2

PE3

(0,0) (0,1) (0,2)

(1,0) (1,1) (1,2)

(0,3) (0,4) (0,5)

(1,3) (1,4) (1,5)

(2,0) (2,1) (2,2)

(3,0) (3,1) (3,2)

(2,3) (2,4) (2,5)

(3,3) (3,4) (3,5)

35 of 43

レイアウトの話の例

  • 二項 elementwise op (例えば A+B) に異なるレイアウトの 2 つの Tensor が来ると計算できない
    • 例えば下の例では、 (0,5) は同じ PE にない

​

​

​

​

PE0

PE1

PE2

PE3

(0,0) (0,1) (0,2)

(1,0) (1,1) (1,2)

(0,3) (0,4) (0,5)

(1,3) (1,4) (1,5)

(2,0) (2,1) (2,2)

(3,0) (3,1) (3,2)

(2,3) (2,4) (2,5)

(3,3) (3,4) (3,5)

PE0

PE1

PE2

PE3

(0,0) (0,1) … (0,5)

(1,0) (1,1) … (1,5)

(2,0) (2,1) … (2,5)

(3,0) (3,1) … (3,5)

A

B

36 of 43

レイアウトの話の例

  • A を B のレイアウトに変換するか、 B を A のレイアウトに変換してから A+B を計算することになる

​

​

​

​

PE0

PE1

PE2

PE3

(0,0) (0,1) (0,2)

(1,0) (1,1) (1,2)

(0,3) (0,4) (0,5)

(1,3) (1,4) (1,5)

(2,0) (2,1) (2,2)

(3,0) (3,1) (3,2)

(2,3) (2,4) (2,5)

(3,3) (3,4) (3,5)

PE0

PE1

PE2

PE3

(0,0) (0,1) … (0,5)

(1,0) (1,1) … (1,5)

(2,0) (2,1) … (2,5)

(3,0) (3,1) … (3,5)

A

B

37 of 43

レイアウトの話の例

  • レイアウト変換の時に、遠くの PE に送る必要があると階層をかけ上がって下りてくる感じになって、パフォーマンスへのペナルティが大きい

​

​

​

​

38 of 43

レイアウトの話の例

  • 任意軸数 Tensor と MN-Core の階層数の組みあわせでこういうことを考えないといけない
    • NCHWc のような固定の表記では表現できない
  • レイアウト変換で大移動が発生すると、非常に大きなパフォーマンスペナルティを喰らう
    • 普通のプロセッサでもキャッシュミスという形で現れているが
  • モデル全体を大域的に見て、なるべく大きなレイアウト変換を減らしたいが、これはとても難しい問題

39 of 43

他にも難しいソフトウェア課題がもりだくさん

  • ホストCPUとMN-Coreの効率的な並列実行
  • スーパースカラをなるべく生かす命令列の詰め込み
  • PE 内に複数種類あるレジスタ的な存在の割りつけ
  • 広いパラメータをサポートする効率的な op 実装
    • 特に Conv/Gemm は広いカバレージを求められる
  • MN-Coreの特性を生かした高効率レイアウト変換
  • DRAM<=>SRAM の転送を最小化する計算スケジューラ
  • PyTorch とのシームレスかつ効率を損なわない連携

40 of 43

このあたりは実装が大変

  • ホストCPUとMN-Coreの効率的な並列実行
  • スーパースカラをなるべく生かす命令列の詰め込み
  • PE 内に複数種類あるレジスタ的な存在の割りつけ
  • 広いパラメータをサポートする効率的な op 実装
    • 特に Conv/Gemm は広いカバレージを求められる
  • MN-Coreの特性を生かした高効率レイアウト変換
  • DRAM<=>SRAM の転送を最小化する計算スケジューラ
  • PyTorch とのシームレスかつ効率を損なわない連携

41 of 43

このあたりはアルゴリズム的に難しい

  • ホストCPUとMN-Coreの効率的な並列実行
  • スーパースカラをなるべく生かす命令列の詰め込み
  • PE 内に複数種類あるレジスタ的な存在の割りつけ
  • 広いパラメータをサポートする効率的な op 実装
    • 特に Conv/Gemm は広いカバレージを求められる
  • MN-Coreの特性を生かした高効率レイアウト変換
  • DRAM<=>SRAM の転送を最小化する計算スケジューラ
  • PyTorch とのシームレスかつ効率を損なわない連携

42 of 43

まとめ

  • Preferred Networks
    • 事業がたくさん。共通インフラを重視
  • PFVM
    • 「エッジからスパコンまで」のグラフコンパイラ
    • 他社が解いてる問題はうまく使う
  • MN-Core
    • ソフトウェアをちゃんと書けば高効率な計算ができる
    • 他社が解いてくれてないので、自分たちで頑張っている

43 of 43