分子動力学シミュレーションにおける力計算のSIMD化 最内ループ展開と最外ループ展開のときの性能比較

動機

短距離分子動力学シミュレーションの力計算のSIMD化は大体において最内ループ(j粒子ループ)を展開して行う。このやり方は最も自然なSIMD化の方法だと思うが、wide SIMDのときに(512bitや1024bit)この最内ループ展開のやり方でちゃんと性能がでるのか少し疑問に思える。というのも短距離分子動力学シミュレーションの最内ループ長はそこまで長くないからである。(40 ~ 100くらいである。)

この記事によるとAVX命令を用い最内ループ4倍アンロールのSIMD化により、約2倍程度の高速化ができたみたいだが、AVX512ではこれがちゃんと4倍程度になるのだろうか。

AVX512を使える環境が手元にないため、SIMD化されたLJ力計算の倍精度演算コードを単精度演算コードに直し、最内ループのアンロールを倍にして、最内ループ長の長さに依存して実行時間がどのように変化するのかを検証してみた。

ソースコードはここにある。

SIMD化する前のスカラー版のコード

こんな感じで20行ちょっとの非常に短いコード。

void reference(void) {
  const auto c24 = 24.0 * dt;
  const auto c48 = 48.0 * dt;
  for (int i = 0; i < N; i++) { // i粒子ループ
    for (int k = 0; k < M; k++) { // j粒子ループ
      const auto j = list[k]; // すべてのi粒子についてひとまず共通のリストを与える。
      const auto dx = q[j].x - q[i].x;
      const auto dy = q[j].y - q[i].y;
      const auto dz = q[j].z - q[i].z;
      const auto r2 = (eps2 + dx * dx + dy * dy + dz * dz); // 発散を防ぐためeps2を足す。
      const auto r6 = r2 * r2 * r2;
      const auto df = (c24 * r6 - c48) / (r6 * r6 * r2);
      p[i].x += df * dx;
      p[i].y += df * dy;
      p[i].z += df * dz;
      p[j].x -= df * dx;
      p[j].y -= df * dy;
      p[j].z -= df * dz;
    }
  }
}

i粒子とj粒子の二重ループになっていて、通常j粒子ループを展開してSIMD化する。
今回はi粒子ループ展開したコードとj粒子ループ展開したコードの実行時間を比較する。AVX命令を使ってSIMD化する。

結果

実行環境
CPU: Intel(R) Xeon(R) CPU E5-2680 v3

実行時間とj粒子ループ長の関係 (ついでにdouble版の4倍展開した場合の結果も載せた。)
bench.png

グラフの凡例はそれぞれ

凡例 詳細 精度
double 1x4 j粒子ループを4倍展開してSIMD化 倍精度
double 4x1 i粒子ループを4倍展開してSIMD化 倍精度
double ref SIMD化していないもの。      倍精度
float 1x8 j粒子ループを8倍展開してSIMD化 単精度
float 8x1 i粒子ループを8倍展開してSIMD化 単精度
float ref SIMD化していないもの。      単精度

実行時間比とj粒子ループ長の関係 (同じようにdouble版の4倍展開した場合の結果も載せた。)
ratio.png

グラフの凡例はそれぞれ

凡例 詳細
double ref/1x4 j粒子ループを4倍展開してSIMD化版とreference版との実行時間比
double ref/4x1 i粒子ループを4倍展開してSIMD化版とreference版との実行時間比
float ref/1x8 j粒子ループを8倍展開してSIMD化版とreference版との実行時間比
float ref/8x1 i粒子ループを8倍展開してSIMD化版とreference版との実行時間比

当たり前のことだが、i粒子ループを展開する方が性能はいい。今回すべての粒子について共通の近接粒子リストを与えるようにしたため、i粒子ループ展開により、計算/ロードの回数の比率が向上し、j粒子ループ展開のときよりも高速化されている。i粒子ループとj粒子ループを展開する場合の性能差は1.6 ~ 1.8倍くらいあることがわかる。ループ長が短くなってくると下手したら2倍くらい性能差が出てくる。

まとめ

4、8倍展開それぞれについて、j粒子ループ、i粒子ループを展開する場合で実行時間がどの程度変化するのかを調べた。i粒子ループで展開すると、(i粒子が共通の近接粒子リストを持っているのであれば、)j粒子ループで展開する場合とくらべ1.6倍から1.8倍くらい高速化されることがわかった。i粒子を展開した方が広いSIMD幅を使って高い性能を出せることがわかった。

これを見ているとi粒子ループで展開するのがベストな気がしてくるが、欠点もあって
1. 展開する倍数個のi粒子について共通のVerlet listを構築しないといけない。
2. 使用するレジスタの本数が多くなる。

1は、4または8個のi粒子について共通のVerlet listを構築するしないといけないことを意味していて、それをやるのはかなり厄介。