LJの力計算のSIMD化(たぶん完結編)

はじめに

これまでのあらすじ。

とりあえずO(N^2)ループでインテルコンパイラより早そうなコードが書けた。しかし、

実用的なMDコードはペアリストによる間接参照、およびカットオフがあるが、それが反映されていない
上記のコードはAoSでやっているが、実はSoAでやったらインテルコンパイラはもっと良いコードを吐くかもしれない

という疑念があった。というわけで、

相互作用ペアリストを作成し、間接参照とカットオフがある状態で速度を計測
AoSとSoAの速度差を調べる

ということを目的としてベンチマークコードを書いた。ただしペアリストは手抜きでO(N^2)で作っている。

コード

コードは長くなったのでhttps://github.com/kaityo256/lj_simd においておく。それぞれ

force_aos.cpp: AoSで、力の計算を何種類かと、シャッフル命令を使って手でSIMD化したコードを含む。
force_soa.cpp: 上記のSoA版。シャッフル命令を使わない馬鹿SIMD化を含む。

力の計算ルーチンは以下の通り。

force_pair: ペアごとにループを回す。一重ループになる。
force_sorted: 相互作用相手でソートして、二重ループにしたもの。
force_next: 手で二段のソフトウェアパイプライニングを施したもの。手持ちの、intrinsicsを使わない条件ではx86系で一番早いコード。
force_intrin: 手でSIMD化したもの。

計測条件、パラメータは以下の通り。

Intel(R) Xeon(R) CPU E5-2680 v3 @ 2.50GHz
icpc (ICC) 16.0.3
コンパイルオプション: -xHOST -O3 -std=c++11
カットオフ3.0、ペアリストに登録する距離3.3
数密度: 1.0
粒子配置: FCC+ランダムな摂動

SIMD化の方針

相互作用相手でソートした状態で、外側をi粒子、内側をj粒子と呼ぶと、j粒子のループを4倍にアンロールして馬鹿SIMD化する。ただし、データをAoSで持っておき、xyz座標を_mm256_load_pdで一度にYMMにロードする(1要素無駄になる)。その後シャッフル命令とか使ってなんとなく4ペアの力を同時に計算する。カットオフがあるため、カットオフ距離以上にいる奴は力をゼロクリアしないといけない。そのため、前回のコードに

(カットオフ距離 - 粒子間距離)でマスクを作る
距離が負の時、カットオフ距離よりも粒子間距離の方が長いため、マスクを使って_mm256_blendv_pdで、力をゼロクリアする

という処理が追加される。後は同じで、データがAoSになってるから、YMMの転置っぽいことをして、それぞれを_mm256_store_pdを使って運動量に書き戻せばOK。

結果

それぞれの関数を100回実行するのにかかった時間。粒子数が119164(〜5.5MB)とかなので、完全にキャッシュにのっていると思われる。

データ配置	計算手法	実行時間	備考
AoS	force_pair	8.194431	ペアごとにループをまわしたもの
AoS	force_sorted	7.010013	相互作用相手でソートして、二重ループにしたもの
AoS	force_next	5.818571	手で二段のソフトウェアパイプライニングを施したもの
AoS	force_intrin	3.888375	手でSIMD化したもの (シャッフル命令あり)
SoA	force_pair	9.554387	ペアごとにループをまわしたもの(SoA版、以下同様)
SoA	force_sorted	7.105266	相互作用相手でソートして、二重ループにしたもの
SoA	force_next	5.634103	手で二段のソフトウェアパイプライニングを施したもの
SoA	force_intrin	5.846895	手でSIMD化したもの (シャッフル命令なし)

まとめのようなもの

AoSで持つよりも、SoAで持つ方が若干速度が向上したが、その差はほんの少しだけだった。O(N^2)ループのような、きれいなデータアクセスだともっと差がつくのだが、間接参照だらけだとあまり差がつかないっぽい。

手でSIMD化したバージョン、結構がんばって書いた割に速度向上は33%程度。これも比較的高密度だからで、もう少し低密度、数密度0.5だと13%程度。うーん、もう少し早くなって欲しかったなぁ。手で書いた二段ソフトウェアパイプライニングがそれなりに効果的なので、手で四段とかパイプライン組んで、なおかつそれをSIMD化とかしたら早くなるのかなぁ。もはや人内が書けるコードではない気がするけど。

倍以上になったら詳細な説明でも書こうとおもったけど、思ったほどでもなかったので、ソースだけ置いておきます。

追記(SoA版のSIMD化)

もともとの動機は普通に馬鹿SIMD化すると、データのパック/アンパックの処理が重いため高速化しないのでは？というところだったので、データのパック/アンパックを伴う、シャッフル命令を全く使わないバージョンのSIMD化コードを追記。SIMD化する前のコード(7.105266秒)よりは高速化された(5.846895)けれど、手でソフトウェアパイプライニングしたもの(5.634103秒)よりは遅いので、やっぱりシャッフル命令を使ってload/storeをケチった方が早いですな。