LJの力計算のSIMD化ステップ・バイ・ステップ その3.5

LJの力計算のSIMD化ステップ・バイ・ステップ その3.5

はじめに

LJのSIMD化の試み。4倍展開したループをSIMD化するが、その途中経過。

• その0
• その1
• その2
• その3
• その3.5 ←イマココ
• その4
• その5
• その6

コードは https://github.com/kaityo256/lj_simdstep においてある。

今回、ようやくAVX2命令が出てきた。

デバッグ

ループが4倍展開されているので、SIMD化するのはわりと「そのまま」できる。だが、ymmの上位と下位がごっちゃになったり、依存関係を間違えてロードしてしまったりとバグを入れることがすごく多い(経験談)。というわけで、SIMD化の「途中」のコードを晒しておく。誰かの参考になれば。

実例1

以下は相対座標を計算するコードである。

  double dx_b1 = q[j_a1][X] - q[i_a1][X];
  double dy_b1 = q[j_a1][Y] - q[i_a1][Y];
  double dz_b1 = q[j_a1][Z] - q[i_a1][Z];

先に述べたように、データ構造が[x1,y1,z1,0,x2,y2,z2,0, …]と並んでいるので、xyz座標をload_pd一発で持ってこられる。そこで、左辺と右辺をそれぞれ取得し、引けば上記のSIMD化は完成である。

  v4df vqi_a1 = _mm256_load_pd((double*)(q + i_a1));
  v4df vqj_a1 = _mm256_load_pd((double*)(q + j_a1));
  v4df vdq_b1 = vqj_a1 - vqi_a1;

デバッグ用にprint256という関数を用意してある。これでv4dfベクトル(ymmレジスタ)を表示し、たしかにこれが(dx_b1, dy_b1, dz_b1, 0.0)というベクトルになっているかを確認する。

実例2

以下は４倍展開の1番目のペアの相対距離の計算ルーチン。

const double r2_1 = (dx_a1 * dx_a1 + dy_a1 * dy_a1 + dz_a1 * dz_a1);

先程、ymmレジスタには(dx,dy,dz,0)という形で相対ベクトルが入っている。これを自乗すると(dx^2, dy^2, dz^2)になる。 これを_mm256_permute4x64_pd(vr2_1x, 201)で (dy^2, dz^2, dx^2, 0)に、_mm256_permute4x64_pd(vr2_1x, 210)で (dz^2, dx^2, dy^2,0)にして、全部足すと (r2_1, r2_1, r2_1, 0)になるはずである。

そこを確認しているコードが

    v4df vr2_1x = vdq_a1*vdq_a1;
    v4df vr2_1y = _mm256_permute4x64_pd(vr2_1x, 201);
    v4df vr2_1z = _mm256_permute4x64_pd(vr2_1x, 210);
    v4df vr2_1 =  vr2_1x + vr2_1y + vr2_1z;
    /*
    printf("%.10f\n",r2_1);
    print256(vr2_1);
    exit(1);//ここまでOK
    */

いちいち表示してexit(1)しては「ここまでOK」とチェックをしている。地道な作業である。まぁ、世の中print文デバッグは基本なのである。こうして全ての計算をスカラ計算とベクトル計算で同時に行い、全ての結果が正しいことを確認したところまでが

https://github.com/kaityo256/lj_simdstep/tree/master/step3.5

にある。

こういう「SIMD化作業途中」のコードが晒されることってあまりなさそうなので晒してみた。

次回に続く。