相互作用計算カーネルジェネレータPIKGの開発とN体およびSPHカーネルの富岳向け最適化
野村 昴太郎 / Kentaro Nomura
神戸大学理学研究科惑星科学研究センター

1

目次
• 背景
• PIKGについて• 富岳向けの最適化ついて• まとめ/富岳の感想など

2

背景• 分子動力学やN体，SPHなどの粒子系シミュレーション
コードの性能は(多くの場合)粒子間相互作用計算の実装(最適化)に大きく依存• ユーザはそれぞれの計算機向けに最適化された相互作用カーネルを書かないといけなく，できあがるものは(大抵の場合)再利用できない• 単一のコードでさまざまな計算機上で十分な性能を引き出したい！
3

FDPS (Framework for Developing Particle Simulators)•

粒子同士が相互作用しながら時間発展するシミュレーションの開発支援フレームワーク•FDPSはノード間通信・系の分割・相互作用リスト作成など面倒なことを引き受けてくれる•

ユーザはユーザー定義クラス，時間発展(main関数)と相互作用カーネルを記述•

粗密にあわせて動的に領域を分割(右図)•

ツリー構造を使った相互作用リストの作成•Fortran/Cインターフェース•

アクセラレータの利用
ユーザが記述ヘッダーをインクルード
ユーザが記述
main

FDPS
APIs

kernelユーザが記述user deﬁned class 4

白い点ひとつひとつは星，黒い部分はなにもない空間
Iwasawa+, PASJ(2016)

FDPSの目指すところ• ユーザは粒子の情報と時間発展と粒子間相互作用を記述するだけ• 用意されたAPIを使いながらシリアルコードに近いものを書くと勝手に並列化される• ユーザは(ほとんど)並列化については考えない• MPI/OpenMPの利用はコンパイルオプションだけで変更• アクセラレータのあるヘテロジニアスな構成でも利用可能• さまざまな言語から利用可能
5

FDPSの計算の流れ
１．領域分割ロードバランスを考慮しながら計算領域をサブ領域に分割して各MPIプロセスに粒子を割り当て

２．粒子情報の交換

並列化
ノード間通信粒子の移動などにあわせてそれぞれのプロセスが持つ粒子の情報を交換

３．相互作用情報の交換相互作用に必要な情報を隣接するプロセスから取得

APIを提供相互作用リストを提供し計算を高速に実行

４．相互作用の計算
各MPIプロセス内で相互作用を計算

５．粒子の時間発展それぞれのプロセスが担当するの相互作用を計算
6

ユーザーが自由に記述

並列化性能• 様々なスパコンで高い並列化効率• 相互作用カーネルの速度が性能に大きく影響• これまでは似鳥さんらエキスパートの手によっ
Iwasawa+, PASJ(2016)て素晴らしいカーネルが作られてきた
7• カーネルひとつくらいなら頑張ってもいいが…

PIKG• (Paricle-particle Interaction Kernel Generatorの略)• 相互作用計算をDSLで簡単に記述• ジェネレータにDSLコードとオプション(アーキテクチャの指定とか)を与えると最適化された粒子間相互作用カー
ネルコードがでてくる• ユーザーは(ほとんど)最適化については考えない• 単一コードから複数アーキテクチャの最適化コードを生成

8

まずは使い方から
ai =

mj
ai =
rij
3
|rij |

EPI
rx, ry, rz

mj
rij
3
|rij |

EPJrx, ry, rz, m

FORCEax, ay, az• ユーザは相互作用を及ぼされる粒子(EPI)と及ぼす粒子(EPJ)とEPIにかかる相互作用(FORCE)の変数を定義• 定義した関数などを用いて2粒子間相互作用の計算を記述• オプション等を設定してジェネレータにかけるとカーネルができる

9例(N体モノポール)
mj
ai =rij
3(ei + ej + |rij |)
ui
mj
=
mi
ei + ej + |rij |EPIrx, ry, rz, εEPJrx, ry, rz, m, ε
FORCEax, ay, az, uVariable deﬁnition:[class_type] type varname [: member_name]
Function deﬁnition:function name(variables…)[statements…]
return val
endKernel deﬁnition:
variable (=|+=|-=) expression10生成コード(non-SIMD)

11PIKGの中身に入る前に…
• Single Instruction Multiple Data(SIMD)• 複数のデータに対して同じ命令を実行• 近年プログラムのシングルコア性能をピークに近づけるためにはコードはSIMD化され(るように書か)なくてはならない• ロード/ストアもまとめて行うのでデータ形状が大事• Array of Structure(AoS)とStructure of Array(SoA)• 例えば各粒子が(x,y,z,w)の4要素を持っているとき• (AoS)x0y0z0w0x1y1z1w1… (SoA)x0x1x2x3y0y1y2y3…
• 同じ命令を各データに当てはめるのでSIMD演算はSoA構造がよい(場合がほとんど)

12どうやって最適化コードを生成するか• 粒子間相互作用計算の一般化• DSLの設計• DSLの読込(一般的なものを使っているので今回は割愛)• オプションから適切なループ構造を選択• アーキテクチャ(命令セット)に合わせてカーネルを生成

13

2粒子間相互作用計算の一般化

1. 前処理が必要であればここで行う
2. EPI変数及び一時変数でカーネル計算で使われる変数をロード
3. 相互作用のアキュムレートに必要なFORCE変数のロード
4. EPJ変数のロード．(2)と同様
5. 相互作用の計算を行い，FORCE変数にアキュムレート
6. アキュムレートしたFORCE変数をもとの配列にストア
14一般的なSIMD化作業

1. i/jループ形状をSIMD幅に合わせて変更
2. ロードでSoAからAoSにする
3. SIMD化されたカーネル計算を書く
4. ストアでAoSからSoAにする
15

DSL設計の要件• SoAを入力をもらったら相互作用を計算してSoAで出力する• 入出力の構造体の情報が必要• 2重ループをユーザに書かせない• 最適化で大事な要因で，アーキテクチャごとに変わるループ形状はPIKG側で決める• 計算中の精度の管理はユーザに任せる• といっても入出力のデータの型を指定するだけ• 混合精度にも対応(まだできてない)• 入出力以外の一時変数は型推論される• 最適化で重要となる関数(sqrt/rsqrt/invとか)は予約語にしてPIKGが最適なものを選んで使う

16

実際のDSL(再掲)mjai =rij
3(ei + ej + |rij |)
uimj
=
mi
✏i + ✏j + |rij |EPIrx, ry, rz, εEPJrx, ry, rz, m, ε
FORCEax, ay, az, uVariable deﬁnition:[class_type] type varname [: member_name]
Function deﬁnition:function name(variables…)[statements…]
return val
end
Kernel deﬁnition:
variable (=|+=|-=) expression
17

AoS <-->SoA変換
AoS -> SoA
x0y0z0w0x1y1z1

w1x2y2z2

w2x3y3z3

w3

x0

x1x2

x3y0

y1

y2y3z0z1z2z3w0

w1

w2

w3

Gather load / Structure load / Load+Transpose
SoA -> AoS
x0

x1

x2

x3

y0

y1

y2

y3z0

z1z2z3

w0

w1

w2

w3

x0y0z0

w0

x1

y1z1

w1x2y2z2

w2x3y3z3

w3

Scatter store / Structure store / Transpose+Store18

カーネルを翻訳• PIKG内ではDSLの式をそれぞれ入れ子構造で保持• C = A + B だと(=, C, (+, A B)) のようなイメージ• 与えられたオプションによってAVX2/AVX-512/ARM
SVEなどのイントリンジックスに変換• (例)C = A + B• reference: C = (A + B);• AVX2 : C = ̲mm256̲add̲ps(A,B);• AVX-512: C = ̲mm512̲add̲ps(A,B);• ARM SVE: C = svadd̲f32̲z(pg,A,B);
19

生成コード(最内ループ)
リファレンスARM SVE

20

AVX2/AVX-512でのベンチマーク
120

100

Phantom-GRAPEと比較(https://bitbucket.org/kohji/phantom-grape/src/master/)80• 10万粒子，θ=0.5• Skylake Xeon 1コア計測• パフォーマンスはカーネル計

60

40

20064
350

128256512

1024

2048

PG, AVX2PIKG, AVX2
PG, AVX-512
PIKG, AVX-512

300200用リストを使うかという指標．大きいほど余分な計算が多くなる

150
100500算部分のみ，計算時間は相互作用計算全体• 横軸は何粒子まで同じ相互作

250
Time / sec

Performance / Gfops• FDPSのN体サンプルでPG, AVX2PIKG, AVX2
PG, AVX-512
PIKG, AVX-512

64

128

256

512

Ng

1024

2048

21

富岳概要
https://news.livedoor.com/article/detail/18423524/

https://www.r-ccs.riken.jp/jp/post-k/overview.html
22富岳の最適化で知るべき数字• A64FX̲Microarchitecture̲Manual参照• SIMD演算に使えるレジスタは32本• SVEの演算レイテンシは基本9• コアあたりの演算器は2つなので(雑に考えると)演算ばかりでも18個依存関係のない命令が連続で発行されな
いと演算レイテンシが見えてその分性能が落ちる• L1$レイテンシはSVEは11，その他は8．読込は2ライン，書き込みは1ライン• ほかにもあるが相互作用カーネル作るときはこのくらい23

富岳向け最適化
元コード• ループ分割• 長い(というわけでもないものも)
ループを分割してレジスタ負荷を減らす

for(int i=0;i<N;i++){ A(i); B(i);}for(int i=0;i<N;i++){ A(i);}for(int i=0;i<N;i++){ B(i);}for(int i=0;i<N/n;i++){for(int j=0;j<n;j++) A(n*i+j);for(int j=0;j<n;j++) B(n*i+j);}• ストリップマイニング• L1$にのるようループ回転数を分割
for(int i=0;i<N/n;i++){• 最内ループを定数(n)回転にできる for(int j=0;j<n/m;j+=m){A(n*i+j+0); … A(n*i+j+m-1);• ループアンロール}• 長い命令レイテンシを隠蔽
for(int j=0;j<n/m;j+=m){• アンロール段数(m)は分割したルー B(n*i+j+0); … B(n*i+j+m-1);
プごとに最適値がある}}24

最適化の効果A64FX 2.0 GHz (1コア, 単精度ピーク128Gﬂops)
コンパイラ: FCC 4.1.0 20200415
オプション: -Kfast -Nclang
条件: 512 x 512のN2ループ
40

Monopole
Quadrupole

35

Performance / Gfops

30
25
20
15
10
50

RefSVE+LoopFission+LoopUnroll※富岳共用開始後の性能を保証するものではありません 25

区間多項式近似• 近似したい関数の定義域をm個の区間に分割• それぞれの区間についてn次のminimax多項式で近似• SVE(AVX-512も)のテーブル参照命令でレジスタをテーブルに使える➡例えば単精度でm = 16にすると, n+1個のレジスタを使う• 近似精度は分割数と定義域の分割の仕方，次数で変わる• 除算やルートを減らせる / FMAリッチn⇣⌘lXf

PPA(x; k) =(k)x0ak,l xl=0

<latexit sha1_base64="Ug9fRAqjAJUGHGnww7j/ANXbwmo=">AAACs3ichVHPTxNBFH6sIFhQil5MvGxsIG2izStKJBoSjBdOpIAFki7dzC7TMtnZH5mdNsXN/gP+Ax48aWKM8ep/wMV/gANc8WI8QsLFg6/bTQgQ9U1m5ptv3vfmmxknkiLWiEcjxo3RsZvjE7cKk1O370wXZ+5uxmFXubzhhjJU2w6LuRQBb2ihJd+OFGe+I/mW470a7G/1uIpFGLzW+xHf8VknEG3hMk2UXVxttxIrjLhiOlQB83lSr79M03L/hVdZMq2467cCO5FLmJrMTrxHMrUkb+uy2Tcfm30bW0nZq6SWEp09XWlJu1jCKmZhXge1HJQgj3pY/AwW7EIILnTBBw4BaMISGMTUmlADhIi4HUiIU4REts8hhQJpu5TFKYMR69HYoVUzZwNaD2rGmdqlUyR1RUoTZvEQv+Apfsev+BN//7VWktUYeNmn2RlqeWRPv72/cf5flU+zhr0L1T89a2jDYuZVkPcoYwa3cIf63pt3pxvP12eTOfyIv8j/BzzCA7pB0DtzP63x9fdQoA+oXX3u62Bzvlp7Up1fe1paXsy/YgIewEMo03s/g2VYgTo06NxvcAwn8MNYMJqGY+wOU42RXHMPLoXh/wE5C6e6</latexit>

ak,l : k番目の領域のl次の係数, x0(k): k番目の領域の始点-7(0) -6(1) -5(2) -4(3) -3(4) -2(5) -1(6) 0(7) 1(8) 2(9) 3(10) 4(11) 5(12) 6(13) 7(14) 8(15) 9222220 x20 2x0 x0 x0 x0 2x0 2x0 2x0 2x0 2x0 2x0 xx0 2x0 2x0 2x0 2 26
<latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit><latexit sha1_base64="(null)">(null)</latexit>コードで見る区間多項式近似(SPHカーネルを例に)リファレンスコード

シリアル区間多項式近似コード

SIMD化区間多項式近似コード• sqrtが消える(sqrtのレイテンシは98)• 上記は4次多項式で近似しているが，ここは3次でも精度に影響がない(必要な精度にあわせて演算を削れる)• テーブル命令のレイテンシは6(加減乗算の2/3)

27

まとめ
• 粒子間相互作用計算カーネルジェネレータPIKGの紹介• DSLで粒子間相互作用計算を記述してジェネレータにかけると，オプションによってさまざまなアーキテクチャに最適化されたカーネルが生成される• 現在ARM SVE/AVX-512/AVX2向けの最適化が利用可能• 富岳向けに最適化オプションを追加し，およそ20から30% の効率のカーネルを生成できている• GPU(CUDA)やPEZY-SC2などにも対応予定• 近日githubで公開予定
謝辞
本研究は，文部科学省「富岳」成果創出加速プログラム「宇宙の構造形成と進化から惑星表層環境変動までの統一的描像の構築」の一環として実施したものです。
28

その他富岳を使って• 使い勝手としては基本的に京と変わらない(らしい)• SIMD化は手でやったほうが性能は出る(のでPIKG作ったと言える)• できたてのシステムではよくあることでしょうが，実装の過程でコンパイラのバグを何個か踏んだ．対処法はあるが直ってなかったりするのでポータルのレポートやマニュアル(既知の問題とか)はしっかり読みましょう• (頑張って最適化したい人向け)A64FXの気持ちになるためには理研シミュレータ＋gem5可視化ソフト(Konata)が便利• ロボ太先生(@kaityo256)の記事を読もう• https://qiita.com/kaityo256/items/
00fc50221d86ce3ﬀ2ea
29