はじめに
分子動力学法でよく使われる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化するのが簡単だろう1。このSIMD化では、同じi粒子と相互作用するj粒子には重なりがないため、アトミック処理が要らないという利点もある。
さて、分子動力学法において、あるi粒子と相互作用するj粒子は、メモリ上でバラバラに並んでいる(つまり関節参照アクセスが要求される)。なので、4倍展開したj粒子の、例えばx座標、y座標をそれぞれ4つまとめてロードしようとすると、ロードが12回発生してしまう2。
このメモリアクセスを減らすため、データの持ち方を工夫する。具体的には、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倍になりましたよ、ということ。
それぞれ環境が違うので正当な比較になっているか自信がないのだけれど3、傾向としては
- 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のコードを整理したもの。