はじめに

分子動力学における近接粒子リスト(Verlet listとも)の構築をするコードを論文[1]に記述されているGrid search methodに基づいて実装した。C++で実装したのち、CUDAとthrustを用いてGPUでも動くようにした。ソースコードはここにある。

この記事ではCUDAとthrustでどのように実装したかについて説明する。

分子動力学法のアルゴリズムについてはkaityo256先生が分子動力学法ステップ・バイ・ステップ その1から分子動力学法ステップ・バイ・ステップ その5にまとめているのでそちらを見てからのほうがわかりやすいと思う。

Grid search method

手順を簡単におさらいすると
1. i番目のセル(iセル)の周りにあるセル(jセル)の番号を格納した配列を作成。(これは一度だけ計算すればよい。)
2. システムを格子で区切ってやり、セルに分割する。(格子のメッシュのサイズ$l$は$l \ge r_{s}$となるように選ぶ。)この時に以下の三つを行う。
  (a) 粒子がどの番号のセルに所属しているのかを計算
  (b) 各セルに粒子が何個入っているかを計算
  (c) iセルに所属している粒子情報が配列の何番目から始まるのかを示すインデックスを作成する。
  (d) セル番号をキーとして粒子情報のsortを行う(これは力を計算する際のキャッシュヒット率の向上を見越している。)
3. 1と2で得られた情報をもとにすべてのセルについて以下の処理を行う
  (a) iセルに所属している粒子とjセルに所属している粒子の組み合わせすべてを取り出す。
  (b) 上記で得た組み合わせすべてについて距離を計算し、$r_{s}$以下のペアを取り出しVerlet listとして登録する。

CUDAとthrustを用いた実装

上記の処理をCUDAとthrustを用いて実装する。

1はの処理は並列性が十分にないので、CPUで計算したデータをGPUへ転送する。
2.(a)の処理については並列性が自明であり、次のようなkernelで計算できる。

neighlist_gpu.hpp
// セルの番号を計算する関数
__host__ __device__ __forceinline__
int32_t gen_hash(const int32_t* idx,
                 const int32_t* cell_numb) {
  return idx[0] + (idx[1] + idx[2] * cell_numb[1]) * cell_numb[0];
}

// 各粒子がどのセル番号に所属しているかを計算する関数
template <typename Vec>
__global__ void gen_cell_id(const Vec* q,
                            int32_t* cell_id_of_ptcl,
                            const int32_t particle_number) {
  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (tid < particle_number) {
    const auto qi = q[tid]; // 粒子の座標データ
    int32_t idx[] = { static_cast<int32_t>(qi.x * params_d::inv_cell_leng[0]), // 所属するセルの番号を計算
                      static_cast<int32_t>(qi.y * params_d::inv_cell_leng[1]),
                      static_cast<int32_t>(qi.z * params_d::inv_cell_leng[2]) }; 
    if (idx[0] == params_d::cell_numb[0]) idx[0]--; // 粒子がシミュレーション領域の境界にいる時の処理
    if (idx[1] == params_d::cell_numb[1]) idx[1]--;
    if (idx[2] == params_d::cell_numb[2]) idx[2]--;
    cell_id_of_ptcl[tid] = gen_hash(idx, params_d::cell_numb); // セルの番号を計算
  }
}

2.(b)だが、これはhistogramを作ることに等しい。セルの番号をキーとして各キーに所属している粒子数を計算する処理と考えることができる。これについてはThrustを用いたhistogram作成でまとめたのでこれを参照して欲しい。

2.(c)だが、これはhistogramを作成したあとで、各キーの開始アドレスを求めることに等しい。これもThrustを用いたhistogram作成 各キーの開始アドレスを計算を参考にしてほしい。

2.(d)について。これはcellの番号をkeyとしてthrust::sort_by_keyを呼べばそれで済む。しかし、粒子の位置と運動量を並べ替える際に逐一sort_by_keyを呼び出すのは冗長に思える。
そこで、thrust::sequenceとthrust::scatterの二つを用いることでsort処理自体は一回で済むようにした。

まずthrust::sequenceで0からparticle_number - 1の連続値を作成し、ptcl_id_in_cell_という配列に格納する。

neighlist_gpu.hpp
    thrust::sequence(ptcl_id_in_cell_.thrust_ptr,
                     ptcl_id_in_cell_.thrust_ptr + particle_number);

そのあとでsort_by_keyによりセルの番号をキーとしてptcl_id_in_cell_をsortする。

neighlist_gpu.hpp
    thrust::sort_by_key(cell_id_of_ptcl_.thrust_ptr, // cell_id_of_ptcl_は各粒子のセル番号
                        cell_id_of_ptcl_.thrust_ptr + particle_number,
                        ptcl_id_in_cell_.thrust_ptr);

ptcl_id_in_cell_を用いれば「粒子情報をどこから持って来れば良いか」がわかる。すなわち次のようにすることでsortが完了する。

// copy to buffer
for (int i = 0; i < particle_number; i++) {
  q_buf[i].x = q[i].x;
  q_buf[i].y = q[i].y;
  q_buf[i].z = q[i].z;
}

// gather
for (int i = 0; i < particle_number; i++) {
  q[i].x = q_buf[ptcl_id_in_cell_[i]].x;
  q[i].y = q_buf[ptcl_id_in_cell_[i]].y;
  q[i].z = q_buf[ptcl_id_in_cell_[i]].z;
}

copyもgatherもthrustに実装されている。二つを用いてptcl_id_in_cell_を用いて粒子情報をsortする関数を作ると次のようになる。

neighlist_gpu.hpp
  template <typename T>
  void CopyGather(thrust::device_ptr<T>& __restrict src,
                  thrust::device_ptr<T>& __restrict buf,
                  const thrust::device_ptr<int32_t>& __restrict key,
                  const int size) {
    thrust::copy(src, src + size, buf);
    thrust::gather(key, key + size, buf, src);
  }

...

    //粒子の座標データをsort
    CopyGather(q.thrust_ptr,
               buffer_,
               ptcl_id_in_cell_.thrust_ptr,
               particle_number);
    //粒子の運動量データをsort
    CopyGather(p.thrust_ptr,
               buffer_,
               ptcl_id_in_cell_.thrust_ptr,
               particle_number);

これを用いてsort処理自体は一回だけ行い、他のデータの並び替えについてはcopyとgatherを用いて行う。

3.(a)と3.(b)についてはCUDAのkernel関数を書くことにする。
各セルについて周囲のセルの粒子情報を抜き出し、距離を計算し$r_{s}$以下であればリストとして登録する処理を書くと次のようになる。各処理の意味するところはコメントを参照。

kernel_impl.cuh
__global__ void make_neighlist_naive(const Vec* q, // 粒子座標データ
                                     const int32_t* cell_id_of_ptcl, // 各粒子が所属しているセルの番号
                                     const int32_t* neigh_cell_id, // iセルの周囲に存在するjセルの番号
                                     const int32_t* cell_pointer, // 各セルの情報が開始アドレスを指定するindex
                                     int32_t* neigh_list, // Verlt listの結果を格納する配列
                                     int32_t* number_of_partners, // 各i粒子の周りr_{s}内に存在する粒子数
                                     const Dtype search_length2,
                                     const int32_t particle_number) {
  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (tid < particle_number) {
    const auto qi = q[tid];
    const auto i_cell_id = cell_id_of_ptcl[tid]; // 粒子の所属しているセルの番号
    int32_t n_neigh = 0;
    for (int32_t cid = 0; cid < 27; cid++) { // jセルについて(iセル + 周りの26個のセル)
      const auto j_cell_id = neigh_cell_id[27 * i_cell_id + cid];
      const auto beg_id = cell_pointer[j_cell_id    ]; // jセルに所属する粒子情報の開始アドレス
      const auto end_id = cell_pointer[j_cell_id + 1]; // jセルに所属する粒子情報の終了アドレス
      for (int32_t j = beg_id; j < end_id; j++) { // jセルに含まれる粒子すべてについて
        const auto drx = qi.x - q[j].x;
        const auto dry = qi.y - q[j].y; 
        const auto drz = qi.z - q[j].z;
        const auto dr2 = drx * drx + dry * dry + drz * drz; // 距離を計算し
        if (dr2 > search_length2 || j == tid) continue;
        neigh_list[particle_number * n_neigh + tid] = j; // r_{s}以内であればVerlt listに登録する。
        n_neigh++;
      }
    }
    number_of_partners[tid] = n_neigh;
  }
}

このあとshared memory, read only cache, warp vote functionsを使うなどいろいろチューニングをしたのだが、それは省略する。

CPU版との処理時間の比較

CPU版はマルチコア実行ではないし、SIMD化もしていない(というか試みたがかなり難しく断念)ので実際はそこまで差がないはずだが、10から20倍程度GPU版の方が早いという結果が出た。

image

Tesla K40tとGTX TITANで2倍の性能差があるのが気になる。これについてはストライド幅を変えた時のメモリバンド幅比較 (Tesla K40tとGeforce GTX TITAN)で書いた通り、ユニットストライドでないメモリアクセスをした時のメモリバンド幅がTeslaとTitanで2倍違っていることに起因する。

まとめ

Grid search methodの実装をthrustとCUDAを用いて行ってみた。GPUで実行して10倍から20倍の性能向上が見られることがわかった。この結果はGPUで近傍粒子リストの作成を行っても性能上大きな問題はないことを意味する。力計算の部分は多くのMDソフトウェアで実装されているようにGPUで非常に効率よく実行することができる。時間発展の部分についても並列性は自明であり、性能はそこまで悪くないはずだ。すなわちMDの処理のほとんどをGPUで行うことができる。CPUとGPUで完全に独立に計算できることでガチヘテロな計算の実装も少しは軽減されるのではないだろうか。

参考文献

[1] H. Watanabe, M. Suzuki, and N. Ito, Progress of Theoretical Physics, 126(2), 203-235 (2011).