C++
AVX
AVX2

LJの力計算のSIMD化ステップ・バイ・ステップ その0

More than 1 year has passed since last update.

はじめに

SIMD化、してますか?>挨拶

「京」が128bit、現在のx86系の石が256bit、そしてポスト「京」がSIMD幅可変、ハード512bitという話もあって、もはやSIMD化しないといろいろつらいというか、ポスト「京」が走りだす時に、えらいひとたちがSIMD化率を気にしだしたり、アプリケーションごとにSIMD化率を提出させられたりするような未来が見えなくもないので、私もSIMD化してみようと思います。

で、「ステップ・バイ・ステップ」とあるけれど、これは教材という意味ではなくて、単に私がどういう形でSIMD化をしていったかの作業記録を保存しようという試みです。もっとすごい人はもっとスパッとSIMD化しているんだと思います。

背景

これまでのあらすじ。

で書いたように、古典分子動力学法でよく使われるカットオフ付きLennard-Jones (LJ)ポテンシャルの力計算をAVX/AVX2の命令を使ってSIMD化してみた。その結果、SIMD化していない最速コードに比べて、SIMD化によって33%程度実行時間が短縮したんだけれど、4倍幅使ってるのに、ちょっと加速率が物足りない。特に問題なのは、粒子系が低密度の場合の加速率が悪い、ということだった。

SIMD化は、ループを4倍展開して自明に独立な演算を4つ作ってSIMD化する、馬鹿SIMD化を採用しているのだが、高密度でループ長が60前後、低密度だと30前後になるため、4倍展開のゴミが見えている可能性がある。

そこで、低密度粒子径でも、もっとループ長が取れる形の計算を馬鹿SIMD化することで、もっと高速化できないか考えた。これから書くのはその作業記録。

ソースコード

SIMD化の元になるコードは https://github.com/kaityo256/lj_simdstep/tree/master/step0 においてある。
ペアリストをO(N^2)で作成していて時間がかかるので、ファイルにキャッシュするようにした。力の計算が4種類定義されているが、それぞれ以下の通り

  • force_pair ペア毎に力を計算する
  • force_sorted i粒子をキーにソートし、j粒子に対してループを回す
  • force_sorted_swp force_sortedを手でソフトウェアパイプライニングしたもの
  • force_sorted_intrin force_sortedを手でSIMD化したもの

ペアリストを作成すると、こんな感じで粒子のペアができる。

pair.png

このペア番号でループを回すのがforce_pair。

で、この方法だと、i粒子とj粒子の情報を毎回ロード/ストアしなければならない。LJの力計算はかなり軽いので、このメモリアクセスのペナルティはわりと大きい。そこで、i粒子をキーにしてソートする。

sorted.png

すると、i粒子の情報がレジスタに載るため、メモリアクセスが減る。これをやったのがforce_sorted。

i粒子の情報がレジスタに載っている前提では、力の計算は

  1. j粒子の座標を取ってくる
  2. 相対ベクトルを計算する
  3. 距離を計算する
  4. 距離から力を計算する
  5. 力から力積ベクトルを計算する
  6. 力積ベクトルを運動量に書き戻す

というステップをとる。このうち、1-2と3-6をループごとにずらして、独立な計算にする。1-2の処理をA、3-6の処理をBとすると、ソフトウェアパイプライニングというのは以下のようなイメージ。

swp.png

上が普通の処理、下がソフトウエアパイプライニングしたもの。全体の処理内容は変わっていないが、BnはAnが終わらないと実行できないが、BnとAn+1は独立に実行できるから、それを明示的に重ねてIPCを増やしましょう、という方法論。x86ではわりと効果的な最適化手法。もちろんコンパイラは勝手にはやってくれない。

SIMD化方針

SIMD化については、とりあえずデータを

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

と4次元ベクトルにしておいて、vmovupdで4成分を一度にロード、あとは適当にシャッフルして4つの力を一度に計算する。

で、低密度で性能が出なかったのは、ループ長が短いせいもあるだろうと考え、i粒子でソートする前のもの、つまりペアでループを回して力を計算する、force_pairをSIMD化することを考える。

ただし、そのままではメモリアクセスのペナルティが大きく、性能が出ないことが予想されるため、

  • ペアごとの力の計算ルーチン force_pairを ソフトウェアパイプライニングする
  • それを4倍展開する
  • それをSIMD化する

というステップを踏むことにしよう。

性能

現状での性能は以下の通り。

計算手法 実行時間[sec] 備考
pair 8.166132 ペアごとにループをまわしたもの
sorted 7.029136 相互作用相手でソートして、二重ループにしたもの
sorted_swp 5.79193 ソートしたものを手でソフトウェアパイプライニングしたもの
sorted_intrin 3.890581 ソートしたものを手でSIMD化したもの

目標は、最後の「sorted_intrin」の速度を超えること。

次回へ続く