AVX で遊んでみる

Programming

今年の１月に発表されたIntelの新CPU「Sandy Bridge」から搭載された新機能に、「AVX」というものがあります。AVXを用いると、256ビットのレジスタを用いて、例えばdouble型の計算を4つ同時に行うなど高速に演算をすることができます。
AVXを使うためには、専用のコードを書かなければいけませんが、C++では単体で（アセンブラを用いずに）AVX対応のコードを書くことができるので、いろいろ遊んでみました。

コードの作法

・#include "immintrin.h" を記述
　AVXを使うための関数がここに記述されているので、includeします。
・変数のalignment
　変数の前に「_declspec(align(32))」というおまじないを書いておけば、自動で境界が32バイト境界になるように調整されます。このおまじないがあると、わりと何も考えなくてもAVX命令を使えるようです。
・データ型
　普通int, doubleなどの型を使いますが、AVXでは__m256i, __m256dなどの型を使います。これは直接int, doubleなどとの変換ができませんが、配列の場合ポインタはそのままキャストできるので、かなり楽になります。
逆アセンブリを見たところ、配列をキャストしてAVX型にしたものは、アライメントされているものに対してのみ有効なload, store命令を用いているので、先述した「アライメントのおまじない」を書かないとだめです（エラーになります）。

実際の例

double in[SZ], out[SZ] があって、out[i] = sqrt(in[i]) とする例を示します。

//in, outをアライメントして確保
_declspec(align(32)) double in[SZ], out[SZ];

//in にデータをセットする処理

//AVXを用いて計算
__m256d *in2 = (__m256d *) in, *out2 = (__m256d *) out;
for(int i=0;i<SZ/4;i++) out2[i] = _mm256_sqrt_pd(in2[i]);

これだけです。intrinsic 命令のおかげで、かなり簡潔に記述することができます。
ちなみに、これを普通にsqrtを用いて書いた場合でも、アセンブリを見るとSSEの平方根命令を使っています。これが例えばceilになると、ceil関数を使って書いたものは実際に関数を呼び出しているため、非常に遅くなりました。ちなみにAVXにはceil命令もあり、_mm256_ceil_pd で実行することができます。実験の結果、sqrtはAVXが通常の2倍、ceilに至っては通常の35倍で計算することができました。

AVX を使ってできること

AVXで新たに追加された命令でできることを分かる範囲でまとめてみました。

  • 256bit相当の浮動小数点数(float*8, double*4)について同時に、加減乗除、最大値最小値、平方根、切り捨て、切り上げ、四捨五入
  • (a, b, c, d) と (p, q, r, s)について、(a-p, b+q, c-r, d+s), (a+b, p+q, c+d, r+s), (a-b, p-q, c-d, r-s)を求める関数
  • 256bit同時の論理演算（AND, OR, XOR, NAND?)

参考資料

Intelの公式ページから、「Intel Intrinsic Guide」がダウンロードできます。これを見ると、C++で使える恐らくすべてのintrinsic 命令を調べることができます。但し、物によっては説明を読んでも全く意味不明というものもあります（特に比較などの制御引数をとる命令）。それ以前に、includeすべきヘッダファイルの名前がどこにも書いてありません。同じページからPDF版の資料をダウンロードできますが、そこにすらヘッダファイルの名前はありませんでした。
http://software.intel.com/en-us/avx/