AVX命令を用いたVerlet list構築のSIMD化 その1

はじめに

Verlet list構築のSIMD化を行って一応完成したのでメモ。ソースコードはここにある。いくつか技術的困難があったので何回かに分けて記事として投稿することにした。

(はじめに言っておくと、この記事の通りの順番にSIMD化したわけではないので、ソースコードを見ていると若干違和感を感じられるかもしれない。)

SIMD化前のコード

まずSIMD化するまえのコードを示す。

for (int32_t icell = 0; icell < all_cell_; icell++) { // すべてのicellについて
  const auto icell_beg = cell_pointer_[icell];    // icellに所属している粒子データの先頭アドレス
  const auto icell_size = number_in_cell_[icell]; // icellに所属している粒子データの終端アドレス
  const int32_t* pid_of_neigh_cell_loc = &ptcl_id_of_neigh_cell_[icell][0]; // 隣接セルに含まれる粒子データ
  const int32_t num_of_neigh_cell = ptcl_id_of_neigh_cell_[icell].size(); // 隣接セルに含まれる粒子データ数
  for (int32_t l = 0; l < icell_size; l++) { // icellにいるすべての粒子について
    const auto i = l + icell_beg;
    const auto qi = q[i];
    for (int32_t k = l + 1; k < num_of_neigh_cell; k++) { //icellの隣接セルに含まれるすべての粒子について
      const auto j = pid_of_neigh_cell_loc[k];
      RegistInteractPair(qi, q[j], i, j); // search_length_以内であればペアをVerlet listに登録
    }
  }
}

ここでRegistInteractPairの関数内では
1. 二つの粒子間の距離を計算
2. search_length_未満であれば粒子iと粒子jの番号を昇順に並び替える
3. 粒子番号が小さい方をkey_particles_、粒子番号が大きい方をpartner_particles_の配列にそれぞれストア
4. 粒子ペアの数と、key_particles_に相互作用している粒子数を計算
の四つの処理を行う。実装は次の通りになる。

void RegistInteractPair(const Vec& qi,
                        const Vec& qj,
                        const int32_t index1,
                        const int32_t index2) {
  const auto dx = qj.x - qi.x;
  const auto dy = qj.y - qi.y;
  const auto dz = qj.z - qi.z;
  const auto r2 = dx * dx + dy * dy + dz * dz; // 二粒子間の距離を計算
  if (r2 > search_length2_) return; // search_length以上であれば以降の処理をスキップ
  int i, j;
  if (index1 < index2) { // indexを昇順に並び替える。
    i = index1;
    j = index2;
  } else {
    i = index2;
    j = index1;
  }
  key_particles_[number_of_pairs_] = i;     // <-SIMD化できる?
  partner_particles_[number_of_pairs_] = j; // <-SIMD化できる?
  number_of_partners_[i]++;                 // <-SIMD化できる?
  number_of_pairs_++;                       // <-SIMD化できる?
}

まず考えること

ひとまず最内ループを4倍展開してSIMD化する方針でいくことにする。この部分はLJ力計算のSIMD化と同じやり方。上のコードで「SIMD化できる?」と示した箇所はSIMD化が難しそうなのでひとまず32bitずつストアすることにする。

二粒子間距離の計算

LJ力計算のSIMD化と全く同じように座標データを_mm256_load_pdで一発でロードしたのち、相対ベクトルを計算する。

const auto ja = pid_of_neigh_cell_loc[k + l + 1];
const v4df vqja = _mm256_load_pd(reinterpret_cast<const double*>(q + ja));
v4df dvqa = vqja - vqi;

const auto jb = pid_of_neigh_cell_loc[k + l + 2];
const v4df vqjb = _mm256_load_pd(reinterpret_cast<const double*>(q + jb));
v4df dvqb = vqjb - vqi;

const auto jc = pid_of_neigh_cell_loc[k + l + 3];
const v4df vqjc = _mm256_load_pd(reinterpret_cast<const double*>(q + jc));
v4df dvqc = vqjc - vqi;

const auto jd = pid_of_neigh_cell_loc[k + l + 4];
const v4df vqjd = _mm256_load_pd(reinterpret_cast<const double*>(q + jd));
v4df dvqd = vqjd - vqi;

4つの相対ベクトルから4つの距離の二乗のデータを作りたい。これを行うためにshuffle命令を用いて相対ベクトルのデータを並び替えることにする。

dvqa = {dxa, dya, dza, 0}  4x4transpose {dxa, dxb, dxc, dxd}
dvqb = {dxb, dyb, dzb, 0}        ->     {dya, dyb, dyc, dyd}
dvqc = {dxc, dyc, dzc, 0}               {dza, dzb, dzc, dzd}
dvqd = {dxd, dyd, dzd, 0}               {  0,   0,   0,   0}

このあと、dvqa*dvqa + dvqb*dvqb + dvqc*dvqc で4つの距離の二乗のデータを作ることができる。

カットオフ処理

このあと、どの番号の粒子がsearch_length_以内にあるのかを判定する必要がある。これには比較関数というものを用いる。

__m256d _mm256_cmp_pd (__m256d a, __m256d b, const int imm8)

Intel intrinsics guideを初めて見たときはちょっとギョッとしてしまったが、4つの倍精度データそれぞれについて第三引数で指定した比較演算を実行して、trueであれば0xFFFFFFFFFFFFFFFFを、そうでなければ0を代入した値を返す。比較演算は32種類用意されている。

今行いたいのは

dr^{2} \le r^{2}_{s}

であり、_CMP_LE_OSというのがそれに相当するようなので、

// dr2_abcdには距離の二乗の情報が粒子4つ分格納されている。
// vsl2はsearch_lengthの二乗のデータが4つ分格納されている。
v4df dr2_flag = _mm256_cmp_pd(dr2_abcd, vsl2, _CMP_LE_OS);

で$r_{s}$以内にある粒子がどれかがわかる。
力の計算のときには、これをmaskとして力の情報を0クリアするのだが、いまは$r_{s}$以内にいるのかそうでないのかだけが知りたい。256bitもいらないので、これを32bitデータに直すために次の関数を用いる。

int _mm256_movemask_pd (__m256d a)

先ほどのmaskをこの関数に渡すと例えば、

256 bit mask                                                        32 bit
ffffffffffffffff0000000000000000ffffffffffffffff0000000000000000 -> 00001010

となる。
この後で各ビットが0か1であるかを見てやれば各粒子が$r_{s}$以内にいるのかどうかが分かる。右へbitシフトしてやれば次のようにして計算できる。

int flag = _mm256_movemask_pd(mask);
const bool ja_in_range = flag & 1;
if (ja_in_range) RegistInteractPair(qi, q[ja], i, ja);

flag >>= 1;
const bool jb_in_range = flag & 1;
if (jb_in_range) RegistInteractPair(qi, q[jb], i, jb);

flag >>= 1;
const bool jc_in_range = flag & 1;
if (jc_in_range) RegistInteractPair(qi, q[jc], i, jc);

flag >>= 1;
const bool jd_in_range = flag & 1;
if (jd_in_range) RegistInteractPair(qi, q[jd], i, jd);

bit shiftとmaskを繰り返すことで各ビットが1か0であるかを得ることができる。得られたビットに応じて粒子idの情報をストアするかしないかを逐一判断していけばよい。

結果

上記のことを踏まえて実装したのがこれのMakePairListFusedLoopSIMDSeqStoreである。残念ながら遅くなってしまった...

実行時間
reference 22674[ms]
simd_seq_store 24222[ms]

まとめ

最初のステップとしてLJ力計算のSIMD化で使われている方法を参考にしながら「最内ループ4倍展開+距離計算のSIMD化+粒子番号は逐次ストア」を行ってみた。次回は逐次ストアをしないやり方を紹介し、referenceより早くする。