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
   今日発表するコードはここに置いてあります
   東京大学物性研究所 渡辺宙志
   物性犬
   

   View Slide

2. 2/26
   分子動力学法とは
   分子動力学法とは
   Molecular Dynamics Simulation (MD)
   分子間に働く力を定義し、ニュートンの運動方程式を数値的に解く手法
   科学技術分野で広く使われている
   例：創薬(タンパク質のシミュレーション)、エネルギー関連(燃料電池のシ
   ミュレーション)、材料開発(高分子の破壊シミュレーション) 、etc
   ただし、力の到達距離によりコードの性格は大きく変わる
   ・ 遠距離力
   ・ 重力多体系
   ・ クーロン相互作用のある系
   ・ 短距離力
   ・ 相互作用距離に上限(カットオフ)がある
   計算手法として
   分子間に働く力を計算し、運動量と位置を更新するステップを繰り返す
   ホットスポットは力の計算
   本発表では短距離古典分子動力学法の高速化手法について説明します
   Movie
   

   View Slide

3. 3/26
   ハイブリッド並列化 (1/2)
   Process Process
   Process Process
   Process Process
   Process Process
   Flat-MPI
   計算領域を分割し、CPUコアそれぞれに割り当てる
   各CPUコア間で直接通信を行う
   利点：プログラムが楽
   欠点：プロセス数の増加でメモリ使用量が飛躍的に増加
   → 「京」フルノードflat-MPIは不可能
   → ハイブリッド不可避
   

   View Slide

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法
   

   View Slide

5. 5/26
   ペアリストとBookkeeping法 (1/2)
   ペアリストとは？
   相互作用距離(カットオフの距離)以内にある粒子対のリスト
   全粒子対についてチェックすると
   高速に粒子対を作成する方法 → グリッド探索
   )
   ( 2
   N
   O
   グリッド探索
   ・空間をグリッドに切り、その範囲に存在する粒子を登録する→ ()
   ・ 相互作用する粒子対がブロックを二つまたが
   ないように空間をブロックに分ける
   ・ 住所録を作る
   ・ 住所録から逆引きして相互作用相手を探す
   

   View Slide

6. 6/26
   ペアリストとBookkeeping法 (2/2)
   Bookkeepingによるペアリスト再利用
   相互作用ペア探索(グリッド探索を用いる)は重い処理
   → 探索距離にマージンを設けて何ステップか再利用する
   相互作用距離
   探索距離
   マージン
   ・ マージンを「ペアリストの寿命」とする
   ・系の最速粒子が進んだ距離の二倍だけマージンを減らす
   ・ マージンが無くなったらペアリスト再構築
   相互作用距離 マージン
   最悪のケース
   最速と次最速の原子が、お互いに近づく時
   → 系の最速の原子がマージンの半分を
   横切るまではペアリストが有効
   有効期限判定
   ペアリストに、相互作用していないペアが含まれる
   → 最内ループに条件分岐が入る
   

   View Slide

7. 7/26
   Lennard-Jonesポテンシャルの力計算 (1/2)
   r
   dx
   dy
   ポテンシャル
   力
   力積
   ここをまとめる
   計算したい量
   相対座標
   相対距離
   ※ ポテンシャルが距離の偶数ベキの場合、逆数平方根が不要 c.f. 重力多体系
   

   View Slide

8. 8/26
   Lennard-Jonesポテンシャルの力計算 (2/2)
   減算*3
   乗算*3 + 加算*2
   乗算*2
   乗算*2
   乗算*2 + 減算*1 + 除算1
   乗算*3 + 加算*3
   乗算*3 + 加算*3
   加減乗算 * 27 + 除算*1 LJの力計算は軽い(メモリアクセスがボトルネック)
   

   View Slide

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粒子
   

   View Slide

10. 10/26
    ソフトウェアパイプライニング (SWP)
    ループ内でメモリアクセスと計算に依存関係があるため逐次処理となる
    最適化前
    最適化後
    相互作用相手を先読みし、独立なメモリアクセスと演算を増やす
    (IPC向上が狙い)
    赤：メモリアクセス
    青：演算
    

    View Slide

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化をがんばっていく
    

    View Slide

12. 12/26
    AVX2を用いたSIMD化 (1/5)
    ナイーブな実装
    ループを4倍展開し、4つの座標データをymmレジスタにロードする
    4つの座標データをymmレジスタにロードする
    ※ y, zも同様
    ※ y, zも同様
    実装してみると非常に遅い
    ナイーブなpackは多数のメモリアクセスを伴う
    gatherも遅いし、scatterは無い。
    

    View Slide

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
    

    View Slide

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;
    相対ベクトル
    転置
    自乗和
    

    View Slide

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を使ったマスク処理
    マスク処理
    

    View Slide

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実装
    

    View Slide

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もある
    

    View Slide

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実装
    

    View Slide

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粒子のインデックスに衝突がないことを知らない
    

    View Slide

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化
    

    View Slide

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粒子のインデックスのビットシフトが不要になる
    

    View Slide

22. 22/26
    データ構造の工夫 (2/4)
    ×1.93
    ×1.09
    ×1.21
    ×1.09
    ×1.32
    ナイーブな実装
    i粒子でソート
    手でgather/scatterを書く
    SoA+ループ端数処理最適化
    SoA+ ループ端数処理最適化+SWP
    最適化の効果
    ナイーブにgather/scatterを書くより30%以上高速化
    

    View Slide

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%
    

    View Slide

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
    

    View Slide

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を
    吐いてない
    

    View Slide

26. 26/26
    まとめ
    Skylake (Core-i5) ではSIMD化無しの最速実装から、AVX2を使うことで
    3倍以上高速に
    KNLでの最速実装はキャッシュラインに最適化したデータ構造+手動
    gather/scatter＋ループ端数最適化＋SWP
    マスクロード/ストアは便利
    gather/scatterがそこそこ早いとレジスタのシャッフルがいらなくて楽
    結果のまとめ
    感想
    インテルコンパイラわりと賢い
    (KNLでは最速実装とコンパイラに任せた場合の差が30%程度)
    Skylake (Xeon)での最速実装はキャッシュラインに最適化したデータ構
    造＋AVX2
    最適化もAVX2より楽 (コード量も短くなる)
    

    View Slide