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

はじめに

前回の記事では4倍アンロール+SIMD化をしたにも関わらず高速化できなかった。

前回の反省点

やはり粒子番号を逐次でストアしていくのでは条件分岐も多くなって効率が悪い。これをやめてパックドデータとしてストアするやり方を考える。

方針

まず粒子番号のストア箇所を再度見てみる。

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_++;
}

大きく分けて4つの処理がある。
1. i粒子とj粒子の番号を昇順に並び替える。
2. 小さい方(以降key粒子と表記)をkey_particles_に、大きい方(以降partner粒子と表記)をpartner_particles_にストアする。
3. あるkey粒子のまわりで$r_{s}$以内にいる粒子数をカウントして、number_of_partnersに格納する。
4. pairが何個あるのかを数える。

まず3についてだが、間接参照がはいっているのでSIMD化が難しい。scatter命令があればできるかも知れないが、いまHaswellでの実行を考えているので使えない。ひとまず3をSIMD化することは諦めて、後処理に回すことにする。

1はminやmax関数を使えばできそう。2も簡単そうだが、$r_{s}$以内にいる粒子についてのみストアする処理が必要なので少し複雑になる。4については前回の記事で256bit maskを理をmm256_movemask_pdを使って22bitデータに直した。得られた32bitデータ(flagと呼ぶ)を用いて$r_{s}$以内にいる粒子数は計算できそう。

SIMD化に向けた準備

最初にAVX命令を用いて256bitとしてまとめてストアできるように、データ構造を少し変更する。今のようにkeyとpartner別々にストアする限り32bit x 4 = 128bit単位のストアになってしまう。ここを最初に変更することにしよう。

これは簡単で

int32_t key_particles[MAX_PAIRS];
int32_t partner_particles[MAX_PAIRS];

となっているのを

int32_t key_partner_particles[MAX_PAIRS][2];

とすれば済む。ストアするときはパックドデータとして

// vkey_part_id = {ik_a, ip_a, ik_b, ip_b, ik_c, ip_c, ik_d, ip_d}
// ik means id of key particle
// ip means id of partner particle
_mm256_storeu_si256(reinterpret_cast<__m256i*>(key_partner_particles_[number_of_pairs_]), vkey_part_id);

のようにする。ここでik_a(bcd) ip_a(bcd)となっているのはそれぞれkey粒子、partner粒子の番号である。key粒子とpartner粒子の4つのペアをまとめた256bitデータを作りまとめてストアする。

i粒子とj粒子の番号を昇順に並び替える

これは簡単でmaxとmin関数のそのまま適用すればよい。

v4di vi_id = _mm256_set_epi64x(i, i, i, i); //i粒子番号
v4di vj_id = _mm256_set_epi64x(ja, jb, jc, jd); //j粒子番号
v8si vkey_id = _mm256_min_epi32(vi_id, vj_id); //iとj粒子番号のうち小さい方
v8si vpart_id = _mm256_max_epi32(vi_id, vj_id); //iとj粒子番号のうち大きい方

しかし、この後keyとpartnerのペア4つを256bitデータとしてストアするのでvkey_idとvpart_idをペアにして256bitデータに変更しよう。やりたいことはvpart_idをvkey_idを互い違いになるように混ぜ合わせること。まず、vpart_idを左へ4Byteシフトする。

vpart_id = _mm256_slli_si256(vpart_id, 4);

そしてvkey_idとのorをとる。

v8si vpart_key_id = _mm256_or_si256(vkey_id, vpart_id);

これで256bitデータとしてkey粒子とpartner粒子の番号が用意できる。
一連の流れを図にするとこんな感じ

図0 vpart_key_idができるまで (iとかjaとか書いていないところは0で埋められている。)
image

key粒子とpartner粒子のストア

256bitデータkey粒子とpartner粒子の番号が用意できたので、$r_{s}$以内のものだけを取り出してストアする処理を実装する。やり方は色々ありそうだが、今回はvkey_part_idを適当に並び替えて、$r_{s}$以内のペアだけをレジスタの右側にくるようにする。
何をやりたいかというと以下のコードのコメントに書いたように、レジスタの中身をshuffleで並び替える。

v4df dr2_flag = _mm256_cmp_pd(dr2_abc, vsl2, _CMP_LE_OS);
const int32_t flag = _mm256_movemask_pd(dr2_flag);

// flagの値が5であるとする。2進数だと0101
// vkey_part_id = {0, 4, 1, 5, 2, 6, 3, 7}であるとする。

// hashに応じてvkey_part_idの内容を並び替えて、ストアするデータをレジスタの右側に寄せる
shuffle(vkey_part_id, hash);

// 寄せた後の値はvkey_part_id = {0, 4, 2, 6, 0, 0, 0, 0}になる。
// リトルエンディアンであることに注意。

問題はshffuleをどのように実装するかだ。
128bit境界の壁を越えてデータを並び替える必要があるので、permute系の命令を使うことになりそうだ。即値を引数にとらないものを探すと

__m256i _mm256_permutevar8x32_epi32 (__m256i a, __m256i idx)

があった。Intelのguideによると処理の詳細は

FOR j := 0 to 7
   i  := j*32
   id := idx[i+2:i]*32
   dst[i+31:i] := a[id+31:id] 
ENDFOR
dst[MAX:256] := 0

となっている。gatherのような処理ができる。

idxとflagの値の対応関係をlookup tableとして用意しておき、$r_{s}$以内の粒子ペアデータがレジスタの右側に固まるようにしよう。

const int32_t shfl_table_[16][8] = {
  {0, 0, 0, 0, 0, 0, 0, 0}, // flag: 0000
  {6, 7, 0, 0, 0, 0, 0, 0}, // flag: 0001
  {4, 5, 0, 0, 0, 0, 0, 0}, // flag: 0010
  {4, 5, 6, 7, 0, 0, 0, 0}, // flag: 0011
  {2, 3, 0, 0, 0, 0, 0, 0}, // flag: 0100
  {2, 3, 6, 7, 0, 0, 0, 0}, // flag: 0101
  {2, 3, 4, 5, 0, 0, 0, 0}, // flag: 0110
  {2, 3, 4, 5, 6, 7, 0, 0}, // flag: 0111
  {0, 1, 0, 0, 0, 0, 0, 0}, // flag: 1000
  {0, 1, 6, 7, 0, 0, 0, 0}, // flag: 1001
  {0, 1, 4, 5, 0, 0, 0, 0}, // flag: 1010
  {0, 1, 4, 5, 6, 7, 0, 0}, // flag: 1011
  {0, 1, 2, 3, 0, 0, 0, 0}, // flag: 1100
  {0, 1, 2, 3, 6, 7, 0, 0}, // flag: 1101
  {0, 1, 2, 3, 4, 5, 0, 0}, // flag: 1110
  {0, 1, 2, 3, 4, 5, 6, 7}  // flag: 1111
};

あとで自分で見直しても分かるように具体例を二つ図として示しておく。

image

コードでは次のようになる。

// shuffleする際のidxをlook up tableから取得
v8si idx = _mm256_load_si256(reinterpret_cast<const __m256i*>(shfl_table_[flag]));
// レジスタ内を並び替えてr_{s}以内にある粒子ペアをレジスタの右側に寄せる。
vpart_key_id = _mm256_permutevar8x32_epi32(vpart_key_id, idx);

このlookupテーブルに基づいてvkey_part_idを並び替えると、レジスタの右側にストア$r_{s}$以内にある粒子ペアの番号だけが固まった状態になる。この後はストアするだけ。

pairが何個あるのかを数える

flagの情報に基づいて$r_{s}$以内にある粒子ペアの数を数える。
これもlookupテーブルを用いて実装しようと思ったのだが、popcnt32という命令でできる。popcnt32は32bit整数のたっているbitの数を数える命令である。例えばflagの値が2進数で0101であったとき、 popcnt32(flag)は2になる。コードで示すと、

number_of_pairs_ += _popcnt32(flag);

のようになる。

reference版との速度比較

上記のアイデアを元に実装したのがこれのMakePairListFusedLoopSIMDである。「最内ループ4倍アンロール+距離計算のSIMD化+粒子番号ストアのSIMD化」を行うと結果は

実行時間
reference 22723[ms]
距離計算をSIMD化して逐次ストア 24222[ms]
距離計算とストアをSIMD化 17821[ms]

と無事に高速化することができた。

まとめ

ストアの部分のSIMD化も合わせて行うことで大きく高速化することができた。しかし、4粒子同時に処理している割には加速の度合いが小さい気がする。次回はアンロールするループを一つ外側にすることでさらに高速化する。