AVX2/AVX-512を用いたLennard-Jones系 ポテンシャルの力計算のSIMD化

2018年2月17日 x86/x64最適化勉強会8@サイボウズ・ラボ（東京）
発表資料

Transcript

1. 1/26 AVX2/AVX-512を用いたLennard-Jones系 ポテンシャルの力計算のSIMD化 @kaityo256 2018年2月17日 x86/x64最適化勉強会8@サイボウズ・ラボ（東京） 1. MDの説明 2. SIMD化前の最適化 3. AVX2を使った最適化 4. AVX-512を使った最適化 5. まとめ Outline https://github.com/kaityo256/lj_simd 今日発表するコードはここに置いてあります 東京大学物性研究所 渡辺宙志 物性犬
2. 2/26 分子動力学法とは 分子動力学法とは Molecular Dynamics Simulation (MD) 分子間に働く力を定義し、ニュートンの運動方程式を数値的に解く手法 科学技術分野で広く使われている 例：創薬(タンパク質のシミュレーション)、エネルギー関連(燃料電池のシ ミュレーション)、材料開発(高分子の破壊シミュレーション) 、etc ただし、力の到達距離によりコードの性格は大きく変わる ・ 遠距離力 ・ 重力多体系 ・ クーロン相互作用のある系 ・ 短距離力 ・ 相互作用距離に上限(カットオフ)がある 計算手法として 分子間に働く力を計算し、運動量と位置を更新するステップを繰り返す ホットスポットは力の計算 本発表では短距離古典分子動力学法の高速化手法について説明します Movie
3. 3/26 ハイブリッド並列化 (1/2) Process Process Process Process Process Process Process Process Flat-MPI 計算領域を分割し、CPUコアそれぞれに割り当てる 各CPUコア間で直接通信を行う 利点：プログラムが楽 欠点：プロセス数の増加でメモリ使用量が飛躍的に増加 → 「京」フルノードflat-MPIは不可能 → ハイブリッド不可避
4. 4/26 ハイブリッド並列化 (2/2) Thread Thread Thread Thread Process Process 通信 Thread Thread Thread Thread プロセスもスレッドも領域分割してしまう 利点： ・ループ分割に比べればプログラムが楽 ・スレッド並列とSIMD化を同時に考えなくて良い 欠点： ・通信がやや面倒くさい cf. HW, M. Suzuki, and N. Ito: Comput. Phys. Commun. 184 2775-2784 (2013) 普通のハイブリッド： プロセス→領域分割 スレッド→ループ分割 スレッド数が増えると、ホットスポット以外もスレッド並列化する必要があり面倒 疑似flat-MPI法
5. 5/26 ペアリストとBookkeeping法 (1/2) ペアリストとは？ 相互作用距離(カットオフの距離)以内にある粒子対のリスト 全粒子対についてチェックすると 高速に粒子対を作成する方法 → グリッド探索 ) ( 2 N O グリッド探索 ・空間をグリッドに切り、その範囲に存在する粒子を登録する→ () ・ 相互作用する粒子対がブロックを二つまたが ないように空間をブロックに分ける ・ 住所録を作る ・ 住所録から逆引きして相互作用相手を探す
6. 6/26 ペアリストとBookkeeping法 (2/2) Bookkeepingによるペアリスト再利用 相互作用ペア探索(グリッド探索を用いる)は重い処理 → 探索距離にマージンを設けて何ステップか再利用する 相互作用距離 探索距離 マージン ・ マージンを「ペアリストの寿命」とする ・系の最速粒子が進んだ距離の二倍だけマージンを減らす ・ マージンが無くなったらペアリスト再構築 相互作用距離 マージン 最悪のケース 最速と次最速の原子が、お互いに近づく時 → 系の最速の原子がマージンの半分を 横切るまではペアリストが有効 有効期限判定 ペアリストに、相互作用していないペアが含まれる → 最内ループに条件分岐が入る
7. 7/26 Lennard-Jonesポテンシャルの力計算 (1/2) r dx dy ポテンシャル 力 力積 ここをまとめる 計算したい量 相対座標 相対距離 ※ ポテンシャルが距離の偶数ベキの場合、逆数平方根が不要 c.f. 重力多体系
8. 8/26 Lennard-Jonesポテンシャルの力計算 (2/2) 減算*3 乗算*3 + 加算*2 乗算*2 乗算*2 乗算*2 + 減算*1 + 除算1 乗算*3 + 加算*3 乗算*3 + 加算*3 加減乗算 * 27 + 除算*1 LJの力計算は軽い(メモリアクセスがボトルネック)
9. 9/26 i粒子によるソート 相互作用相手でソート 相互作用ペア 1 9 0 2 0 1 2 7 3 9 3 5 1 2 1 4 2 8 3 4 0 5 2 6 i粒子でソート 1 2 5 2 4 9 6 7 8 4 5 9 Sorted List 0 1 2 5 1 2 4 9 2 6 7 8 3 4 5 9 配列表現 0 1 2 3 i粒子の情報がレジスタにのる → 読み込み、書き込みがj粒子のみ → メモリアクセスが半分に i粒子 j粒子 i粒子 j粒子
10. 10/26 ソフトウェアパイプライニング (SWP) ループ内でメモリアクセスと計算に依存関係があるため逐次処理となる 最適化前 最適化後 相互作用相手を先読みし、独立なメモリアクセスと演算を増やす (IPC向上が狙い) 赤：メモリアクセス 青：演算
11. 11/26 SIMD化無しの最適化の効果 ナイーブな実装 i粒子でソート さらにSWP Haswell Skylake ×1.14 ×1.24 ×1.44 ×1.07 ナイーブな実装 i粒子でソート さらにSWP 原子数: 12万, 数密度: 1.0, カットオフ: 3.0 100回の力の計算にかかった時間 [ms] この状態からSIMD化をがんばっていく
12. 12/26 AVX2を用いたSIMD化 (1/5) ナイーブな実装 ループを4倍展開し、4つの座標データをymmレジスタにロードする 4つの座標データをymmレジスタにロードする ※ y, zも同様 ※ y, zも同様 実装してみると非常に遅い ナイーブなpackは多数のメモリアクセスを伴う gatherも遅いし、scatterは無い。
13. 13/26 AVX2を用いたSIMD化 (2/5) パディング付きAoS データを粒子数*4成分の二次元配列で宣言 double q[N][4], p[N][4]; Z Y X Z Y X Z Y X Z Y X うれしいこと (x,y,z)の3成分のデータを256bitレジスタに一発ロード ただし、1成分(64bit) は無駄になる Z Y X Z Y X Z Y X Z Y X movupd/movapd ymm Z Y X
14. 14/26 AVX2を用いたSIMD化 (3/5) データの転置 Z Y X i粒子 Z Y X - j粒子 dz dy dx 相対座標 dz1 dy1 dx1 dz2 dy2 dx2 dz3 dy3 dx3 dz4 dy4 dx4 4ペア分計算 dx1 dx2 dx3 dx4 dy1 dy2 dy3 dy4 dz1 dz2 dz3 dz4 転置 r4 r3 r2 r1 自乗和 4ペア分の距離 v4df vdq_1 = vq_j1 - vq_i; v4df vdq_2 = vq_j2 - vq_i; v4df vdq_3 = vq_j3 - vq_i; v4df vdq_4 = vq_j4 - vq_i; v4df tmp0 = _mm256_unpacklo_pd(vdq_1, vdq_2); v4df tmp1 = _mm256_unpackhi_pd(vdq_1, vdq_2); v4df tmp2 = _mm256_unpacklo_pd(vdq_3, vdq_4); v4df tmp3 = _mm256_unpackhi_pd(vdq_3, vdq_4); v4df vdx = _mm256_permute2f128_pd(tmp0, tmp2, 0x20); v4df vdy = _mm256_permute2f128_pd(tmp1, tmp3, 0x20); v4df vdz = _mm256_permute2f128_pd(tmp0, tmp2, 0x31); v4df vr2 = vdx * vdx + vdy * vdy + vdz * vdz; 相対ベクトル 転置 自乗和
15. 15/26 AVX2を用いたSIMD化 (4/5) D2 C2 B2 A2 D1 C1 B1 A1 負 正 正 負 src1 src2 mask D2 C1 B1 A2 相互作用距離とカットオフ距離でマスクを作成し、力積をゼロクリア 0 0 0 0 src1 src2 0 0 ペアAとDが相互作用範囲外 Bookkeeping法により、相互作用範囲外のペアもいる → vblendvpdを使ったマスク処理 マスク処理
16. 16/26 AVX2を用いたSIMD化 (5/5) ナイーブな実装 i粒子でソート さらにSWP AVX2実装 Haswell Skylake 原子数: 12万, 数密度: 1.0, カットオフ: 3.0 100回の力の計算にかかった時間 [ms] ×1.14 ×1.24 ×1.63 ×1.44 ×1.07 ×2.31 ・ SIMD化無しの最速実装に比べてHwで1.6倍、SLで2.3倍高速化 ・ AVX2を用いた加速度合いはSkylakeの方が高い ナイーブな実装 i粒子でソート さらにSWP AVX2実装
17. 17/26 AVX-512を用いたSIMD化 (1/4) gather/scatter D C B A メモリ レジスタ D C B A AVX2まではバラバラに持ってくると遅かった H G F E D C メモリ レジスタ H G F E AVX-512では、メモリから8要素同時にバラバラにロードできる (gather) D C B A B A バラバラにストアできるscatterもある
18. 18/26 AVX-512を用いたSIMD化 (2/4) ナイーブな実装 i粒子でソート さらにSWP AVX2実装 KNL Skylakeとの比較 KNLではAVX2を使うより、コンパイラに任せた方が早い Core i5とくらべて、gather/scatterが効かないと6〜7倍遅いが、効くとまぁまぁ KNLでの実測 とりあえずそのままKNLで測定してみる ナイーブな実装 i粒子でソート さらにSWP AVX2実装
19. 19/26 AVX-512を用いたSIMD化 (3/4) for (int k = 0; k < np; k++) { const int j = sorted_list2[kp + k]; double dx = q[j][X] - qix; double dy = q[j][Y] - qiy; double dz = q[j][Z] - qiz; double r2 = (dx * dx + dy * dy + dz * dz); double r6 = r2 * r2 * r2; double df = ((24.0 * r6 - 48.0) / (r6 * r6 * r2)) * dt; if (r2 > CL2) df = 0.0; pfx += df * dx; pfy += df * dy; pfz += df * dz; p[j][X] -= df * dx; p[j][Y] -= df * dy; p[j][Z] -= df * dz; } コンパイラは何をしたか？ 1. ループを8倍展開 2. 相互作用インデックスをvmovdqu32で8個持ってくる 3. x,y,z成分にgatherをかけて力を計算 4. vpconflictdでインデックスの衝突検出 5. 衝突がなければそのままscatter、あれば逐次書き戻し ソースの最内ループ i粒子 j粒子 コンパイラはj粒子のインデックスに衝突がないことを知らない
20. 20/26 AVX-512を用いたSIMD化 (4/4) 手でgather/scatterを書く コンパイラは衝突検出をしている 我々は衝突がないことを知っているので、vpconflictdの分早くなるはず ×1.93 ×1.09 若干早くはなったが、もう少し頑張りたい ナイーブな実装 i粒子でソート 手でgather/scatterを書く コンパイラがgather/scatter+vpconflictd を使ってSIMD化 手でgather/scatterだけを 使ってSIMD化
21. 21/26 データ構造の工夫 (1/4) データ構造をSoAに AoS Z Y X Z Y X Z Y X Z Y X SoA Z Z Z Z Z Y Y Y Y Y X X X X X ループの端数処理最適化 ・ループ回転数が8の倍数でない余りの処理も一度にまわしてしまう ・「余り」はvpcmpqでマスクを作ってゼロクリア (kand) ・運動量の書き戻しをマスク付きscatterに (ループ端数をまたぐとインデックス衝突の可能性がある) ソフトウェアパイプライニング(SWP) ループを「半個ずらし」して、異なる処理を重ね、IPC向上を狙う ・データが連続になるため、j粒子のインデックスのビットシフトが不要になる
22. 22/26 データ構造の工夫 (2/4) ×1.93 ×1.09 ×1.21 ×1.09 ×1.32 ナイーブな実装 i粒子でソート 手でgather/scatterを書く SoA+ループ端数処理最適化 SoA+ ループ端数処理最適化+SWP 最適化の効果 ナイーブにgather/scatterを書くより30%以上高速化
23. 23/26 データ構造の工夫 (3/4) もう少し頑張る Z Z Z Z Z Y Y Y Y Y X X X X X SoA VZ VZ VZ VZ VZ VY VY VY VY VY VX VX VX VX VX 座標 運動量 Cache Line Size (512 Byte) AoS 8-packed 0 VZ VY VX 0 Z Y X 0 VZ VY VX 0 Z Y X 一つの粒子の情報 Cache Line Size (512 Byte) Cache Line Size (512 Byte) データ構造をSoAからAoSの8要素パックに 一つの粒子の情報がちょうどキャッシュラインサイズに一致する キャッシュライン完食率 75%
24. 24/26 データ構造の工夫 (4/4) ナイーブな実装 i粒子でソート 手でgather/scatterを書く SoA+ループ端数処理最適化 SoA+ ループ端数処理最適化+ SWP AoS 8要素+ループ端数処理最適化+SWP ×1.93 ×1.09 ×1.21 ×1.09 ×1.05 ×1.39 コンパイラの自動ベクトル化より51%高速化 ナイーブにgather/scatterを書くより40%高速化 AoS 8要素の効果 ×1.51
25. 25/26 Skylake (Xeon)でのAVX2とAVX-512 ×1.22 ×1.28 ×1.39 ×1.03 ナイーブな実装 i粒子でソート AVX-512: 手でgather/scatterを書く AVX-512: AoS 8要素 +ループ端数+SWP AVX2: AoS 8要素+ループ端数+SWP Skylake (Xeon)での測定 Skylake (Xeon)では、AoS 8要素 + AVX2の方がAVX-512最速実装より若干早い コンパイラは gather/scatterを 吐いてない
26. 26/26 まとめ Skylake (Core-i5) ではSIMD化無しの最速実装から、AVX2を使うことで 3倍以上高速に KNLでの最速実装はキャッシュラインに最適化したデータ構造+手動 gather/scatter＋ループ端数最適化＋SWP マスクロード/ストアは便利 gather/scatterがそこそこ早いとレジスタのシャッフルがいらなくて楽 結果のまとめ 感想 インテルコンパイラわりと賢い (KNLでは最速実装とコンパイラに任せた場合の差が30%程度) Skylake (Xeon)での最速実装はキャッシュラインに最適化したデータ構 造＋AVX2 最適化もAVX2より楽 (コード量も短くなる)