はじめに

分子動力学法でよく使われるLennard-Jones (LJ)ポテンシャルの力計算のSIMD化について調べていたのだけれど、Haswellで試したときよりも、Skylakeで試した方がSIMD化の恩恵が大きかった。詳細は詰めきれていないが、とりあえず覚書として事実だけ残しておく。

基本アルゴリズム(SIMD化前)

短距離古典分子動力学法では、力の及ぶ範囲にカットオフがある。そこで、そのカットオフ距離内にいる粒子を探し、ペアリスト(Verletリスト)を作る。この時、Bookkeeping法という最適化手法を用いるために、カットオフ距離よりもやや長い距離だけ離れたペアもリストに登録しておく。力の計算はそのペアリストを参照して行うため、ループを回す時に「相互作用範囲内であるかどうか」のチェックが必要になる。

ナイーブな実装(pair)

とりあえずこんなペアリストができたとする。

粒子番号の若い方を「i粒子」、その相方を「j粒子」と呼ぶことが多い。ナイーブな力の計算では、この粒子対ごとに力を計算すれば良い。これを参照実装として「pair」と呼ぶことにする。なお、i粒子でソートする時間はどうせBookkeeping法で消せるので以下の実行時間にはカウントしない。

i粒子でソート(sorted)

先程の実装ではペアごとに座標の読み込み3成分2粒子、運動量の読み込み&書き込み3成分2粒子が発生する。LJの力計算はとても軽いので、これではメモリアクセスが多すぎて性能が出ない。そこで、i粒子でソートする。こんな感じ。

これによりi粒子の情報が座標、運動量ともにレジスタに乗るので、メモリアクセスは半分になる。ここまでした実装を「sorted」と呼ぶ。

ソフトウェアパイプライニング(sorted_swp)

i粒子の情報がレジスタに乗った状態で、j粒子との力を計算するためには

j粒子の座標を読み込む
i粒子との距離を計算する
距離から力(力積)を計算する
j粒子の運動量を読み込み、力積を加えてからメモリに書き戻す

という手続きを実行する必要がある。ここで、1〜4は順番に実行しなければならない。図解するとこんな感じ。

このうち、1と2の処理だけ先にやって、あとは「n番目の3と4の処理と、n+1番目の1と2の処理を一つのループ内で実行する」という、ソフトウェアパイプライニング(swp)をやる。こんな感じ。

行うべき処理は変わっていない(実際にはちょっとだけ増える)が、ループ内で独立な命令を増やしてIPCを向上させる意図がある。特に、スーパースカラにおいて「メモリの読み書き」と「演算」を重ねる意図があるが、実際にどういう理屈で性能が向上しているかはきちんとは理解していない。これを「sorted_swp」と呼ぶ。

AVX2によるSIMD化 (sorted_swp_intrin)

さて、「i粒子でソート」し、さらに「ソフトウェアパイプライニング」したものをAVX2によりSIMD化する。外側がi粒子、内側がj粒子に関する二重ループになるので、もっとも単純には内側、すなわちj粒子に関するループを4倍展開してSIMD化するのが簡単だろう¹。このSIMD化では、同じi粒子と相互作用するj粒子には重なりがないため、アトミック処理が要らないという利点もある。

さて、分子動力学法において、あるi粒子と相互作用するj粒子は、メモリ上でバラバラに並んでいる(つまり関節参照アクセスが要求される)。なので、4倍展開したj粒子の、例えばx座標、y座標をそれぞれ4つまとめてロードしようとすると、ロードが12回発生してしまう²。

このメモリアクセスを減らすため、データの持ち方を工夫する。具体的には、3次元の計算をするのに4次元のベクトルをAoS形式で保持する。こんな感じ。

double q[N][4], p[N][4];

ここでは座標をq、運動量をpで表現している。そうすると、運動量と座標が1粒子ごとにメモリ上で256bitでアラインされるので、ymmレジスタに一発で読み込むことができる(この時、1成分だけ無駄になる)。

なので、4つのj粒子の座標を、4つのロード命令で持ってこられる。

i粒子も同様にymmレジスタにロードしておくと、その差を取ることで相対座標ベクトルになる。これを4粒子分計算する。

ここから、4ペアそれぞれについての距離を計算するため、ymmレジスタを転置する。

あとはその自乗和をとれば、4ペアそれぞれの相対距離の平方が得られるので、力積を4ペア分同時に計算することができる。

運動量のへ書き戻しも、ちょっと工夫すればできる。まず、4つ計算できた力のうち、1成分をブロードキャストする。これをスカラー量だと思って、あとはベクトルの計算をそのまますれば良い。

これを4ペア分それぞれに実行すれば良い。

なお、相互作用距離範囲外にあるペアについては、マスク処理で力積をゼロクリアする。

このSIMD化を、「i粒子でソート」し、さらに「ソフトウェアパイプライニング」したループについて適用したものを「sorted_swp_intrin」と呼ぶ。

計算結果

以上のコードを、HaswellとSkylakeで試す。

まずHaswellの実行環境から。

Intel(R) Xeon(R) CPU E5-2680 v3 @ 2.50GHz
L2 2MB (0.25MB/core)
L3 30MB
g++ (GCC) 5.1.0

Skylakeは手元のiMac (Retina 5K, 27-inch, Late 2015)を使った。

3.3 GHz Intel Core i5
L2 Cache (per Core): 256 KB
L3 Cache: 6 MB
Memory: 8 GB
g++ (Homebrew GCC 6.3.0_1) 6.3.0

GCCのバージョンが違うが、この2つの違いは今回に関してはあまり効いていない気がする。

計算条件は以下の通り。

システムサイズ: 50×50×50
粒子数: 119164
数密度: 0.95
カットオフ: 3.0
探索距離: 3.3
力を100回計算するのにかかった時間(ペアリスト構築の時間は含まない)

粒子数が12万なので、1粒子あたりの情報が48バイトとして6MB弱、つまりSkylakeでぎりぎりL3に入るくらい。Haswellの方は余裕で入ってる。コアあたりのL2のサイズは同じくらい。

コンパイルオプションはどちらも-O3 -mtune=native -mavx2 -std=c++11とした。

結果はこんな感じ。

Haswell

手法	実行時間[s]	一段階前からの加速率
pair	12.811105	1
sorted	8.316984	1.540354653
sorted_swp	6.207216	1.339889574
sorted_swp_intrin	3.300018	1.880964286

Skylake

手法	実行時間[s]	一段階前からの加速率
pair	6.983079	1
sorted	5.362384	1.302234044
sorted_swp	4.516331	1.187331929
sorted_swp_intrin	2.007175	2.25009329

「一段階前からの加速率」をグラフ化したもの。

ここで、「一段階前からの加速率」というのは、順番に最適化を施していった時の、一つ前を基準とした加速率のこと。例えば「i粒子ソート+ソフトウェアパイプライニング」までやった状態で、さらに「SIMD化」したら、Haswellなら1.9倍に、Skylakeなら2.25倍になりましたよ、ということ。

それぞれ環境が違うので正当な比較になっているか自信がないのだけれど³、傾向としては

SIMD化以外の最適化による加速率は、SkylakeよりHaswellの方が恩恵が大きい
SIMD化による恩恵は、HaswellよりSkylakeの方が大きい

ということは言える気がする。

Haswell + icpc

ついでに、Haswellのマシンにはインテルコンパイラも入っていたので、そっちでも試した。コンパイルオプションは icpc -O3 -xHOST -std=c++11。

手法	実行時間[s]	一段階前からの加速率
pair	8.19292	1
sorted	7.045052	1.162932509
sorted_swp	5.814481	1.21163901
sorted_swp_intrin	3.043919	1.910195705

最適化の余地が大きい、ナイーブなコードを食わせた時の最適化は、インテルコンパイラの方ががんばっていますね。AVX2の組み込み関数を使った場合の加速率はGCCと同程度かな。

まとめ

とりあえず、AVX2を生書きした場合の加速に関しては、HaswellよりSkylakeの方が恩恵が大きかった。「まっとうなgatherやscatterがあればこんなややこしいことしなくていいんじゃね？」とか言わないのが大人です。

今回使ったコードは

においておく。基本的にはLJの力計算のSIMD化ステップ・バイ・ステップその6のコードを整理したもの。

個人的にはこういうSIMD化を「バカSIMD化」と呼んでいる。自明並列計算を「バカパラ」と呼ぶことにならっている。 ↩
またAVX2にはgatherはあるがかなり遅く、かつscatterが無いのでストアが多数発生してしまう。 ↩
まずL2キャッシュは同じで、どちらもL3に粒子情報が全て乗っていて、かつアセンブリざっと見た感じではコンパイラのバージョンの違いはあまり出てないように見える。そして同じアルゴリズムを基準に規格化しているので、この加速率はHaswellとSkylakeの命令処理の違いを見ているつもりなのだが・・・、本当にそうなってるかは自信がない。 ↩