はじめに
AVX512Fにはvpternlogdという、任意の三項論理演算を実現する命令がある。ちょうど手元にこれを使ったら高速化できそうなプログラムがあったのだが、そこで扱ってるデータ型が__int128、つまりGCC拡張の128ビット整数だった。これにvpternlogdを使おうとしたらなかなか素直にいかなかったので覚書。
やりたいこと (スカラ版)
3つのビットフィールドがあった時、同じ桁に一つだけビットがたっていたら1、それ以外は0、という演算がしたい。つまり、テーブルにするとこんな感じ。
| src1 | src2 | src3 | result |
|---|---|---|---|
| 0 | 0 | 0 | 0 |
| 0 | 0 | 1 | 1 |
| 0 | 1 | 0 | 1 |
| 0 | 1 | 1 | 0 |
| 1 | 0 | 0 | 1 |
| 1 | 0 | 1 | 0 |
| 1 | 1 | 0 | 0 |
| 1 | 1 | 1 | 0 |
これをビット演算で実現すると、たとえばこんな感じになるだろう。
result = (src1 & src2 & src3) ^ (src1 ^ src2 ^ src3);
128ビット整数にランダムなビット列を与えて、上記のビット演算を実装するとこんな感じになるだろう。
__attribute__((aligned(64))) __int128 t1 = rand_bit();
__attribute__((aligned(64))) __int128 t2 = rand_bit();
__attribute__((aligned(64))) __int128 t3 = rand_bit();
std::cout << "Input" << std::endl;
std::cout << i2b(t1) << std::endl;
std::cout << i2b(t2) << std::endl;
std::cout << i2b(t3) << std::endl;
std::cout << "Output" << std::endl;
// Scalar
__int128 t4 = (t1 & t2 & t3) ^ (t1 ^ t2 ^ t3);
std::cout << i2b(t4) << std::endl;
ここでrand_bitは適当に作ったランダムなビット列を生成する関数で、i2bは整数をビット列の文字列に変換する関数。実行結果はこんな感じ。
Input
10011010010110111010001011001000100100000100011111000111000101011010110000101111110011111010010011100100110011101000000010001110
01110101001000000011010111000101100001111100011100110110011001111000110001100101011011001001111101111001101001100011110111010110
10001111111001000001000100100001110110010110011110001010111100001101011001000001010000110110101010000001011100111001100101110000
Output
01100000100111111000011000101100010011101010000001111001100000100111001000001010101000000101000100011100000110010010010000101000
確かに、３つの入力うち1つだけビットが立っている桁のみ、結果のビットが立っている。
vpternlogd版
さて、vpternlogdは、このresultの8ビットフィールドを即値で与えることで、任意の三項論理演算を一度に行うことができる。上記のテーブルの例なら、与えるべき即値は2+4+16 = 22となる。
_mm_ternarylogic_epi32
まず、__int128は、32ビット整数4個分のデータを持っているのだが、ざっと探した感じ、それがメモリ上で連続しているという保証を見つけることができなかった。でもまぁ連続で取られているだろうと思って、
__int128 t1 = rand_bit();
__m128i x1 = _mm_load_si128((__m128i*)(&t1));
みたいなコードを書いた。これで128ビットのデータがxmmレジスタに乗ったことにしよう。
さて、vpternlogdがxmmレジスタを引数に取ってくれればそのまま呼ぶだけ。これに対応する組み込み関数は_mm_ternarylogic_epi32なので、これが動けば
__m128i x1 = _mm_load_si128((__m128i*)(&t1));
__m128i x2 = _mm_load_si128((__m128i*)(&t2));
__m128i x3 = _mm_load_si128((__m128i*)(&t3));
__m128i x4 = _mm_ternarylogic_epi32(x1,x2,x3,22);
みたいに呼ぶだけで良い。しかし、上記を実行するとSIGILLで落ちてしまう。調べてみると、vpternlogdがxmmやymmレジスタを呼ぶにはAVX512VL拡張が必要。手元にあったKNLにあるフラグはAVX512F, PF, ER, CDで、VLには対応していなかった。というわけでそこをなんとかしないといけない。
インラインアセンブラ版
どうせxmmとzmmは同じデータを共有しているんだから、xmmをzmmにして、そのままvpternlogdを呼んでしまえば良い。インラインアセンブラで書くとこんな感じ。
__attribute__((noinline))
__m128i ternlog(__m128i x1, __m128i x2, __m128i x3){
__asm__ (
"vpternlogd $22,%zmm2,%zmm1,%zmm0\n\t"
);
}
呼び出し規約として、引数の__m128iはxmm0, xmm1に順番に入っているはずで、返り値が__m128iの時は単にxmm0に値を突っ込めばよかったはず。どうせ上位ビットはどうなっていようが関係ないので、それをそのままzmmレジスタにしてしまっている。
なお、インライン展開されると結果がおかしくなるので、インライン展開を抑制している。この関数を使うとこんな感じになる。
__m128i x4 = ternlog(x1,x2,x3);
std::cout << i2b(x4) << std::endl;
結果も正しく計算できている。
_mm512_ternarylogic_epi32
インラインアセンブラ版は毎回callが入ってしまうので遅い。もしかしたらインラインアセンブラ込みで正しくインライン展開させる方法があるのかもしれないが、少なくとも僕は知らない。というわけで組み込み関数を使ってみる。最初はinsert系命令でxmmをymmに、ymmをzmmにinsertしていくのかと思ったけど、対応する命令vinserti32x4とかvinserti32x8がそれぞれAVX512VL, AVX512DQ拡張なので使えない。そんなわけで__int128を64バイトアラインにして、無理やりzmmにロード、結果も無理やりzmmから__int128にストアしてみる。こんな感じ？
__attribute__((aligned(64))) __int128 t1 = rand_bit();
__attribute__((aligned(64))) __int128 t2 = rand_bit();
__attribute__((aligned(64))) __int128 t3 = rand_bit();
__m512i z1 = _mm512_load_si512((__m512i*)(&t1));
__m512i z2 = _mm512_load_si512((__m512i*)(&t2));
__m512i z3 = _mm512_load_si512((__m512i*)(&t3));
__m512i z4 = _mm512_ternarylogic_epi32(z1,z2,z3,22);
__attribute__((aligned(64))) __int128 t5[4];
_mm512_store_si512((__m512i*)(t5),z4);
std::cout << i2b(t5[0]) << std::endl;
一応__int128 t5[4]として64バイト分の領域を確保して、そこに値を突っ込んでいるが、これで良いのがあまり自信が無い。とりあえず結果は正しそう。
組み込み関数をインライン展開＋最適化させる
一度__int128をzmmにロードして、vpternlogdをかける方法でうまくいきそうだったので、それを関数にしてコンパイラによる最適化に期待してみる。
__int128
ternlog512(__int128 &t1, __int128 &t2, __int128 &t3){
__m512i z1 = _mm512_load_si512((__m512i*)(&t1));
__m512i z2 = _mm512_load_si512((__m512i*)(&t2));
__m512i z3 = _mm512_load_si512((__m512i*)(&t3));
__m512i z4 = _mm512_ternarylogic_epi32(z1,z2,z3,22);
__attribute__((aligned(64))) __int128 t[4];
_mm512_store_si512((__m512i*)(t),z4);
return t[0];
}
入力される__int128は64バイトアラインされていることが前提。これはこう呼べる。
__attribute__((aligned(64))) __int128 t1 = rand_bit();
__attribute__((aligned(64))) __int128 t2 = rand_bit();
__attribute__((aligned(64))) __int128 t3 = rand_bit();
std::cout << i2b(ternlog512(t1,t2,t3)) << std::endl;
希望としては、ternlog512がインライン展開され、不要なロードストアが最適化されてほしい。上記のアセンブリはこうなる。
vmovups 64(%rsp), %zmm1
vmovups 128(%rsp), %zmm0
movq %r11, 192(%rsp)
movq %r8, 200(%rsp)
vpternlogd $22, 192(%rsp), %zmm0, %zmm1
あ、なんかできてるっぽい。rand_bitの結果が64(%rsp)、128(%rsp)、192(%rsp)にはいっていて、それをそのままzmmにvmovupsして、vpternlogdが呼ばれている。
まとめ
GCC拡張の128ビット整数型__int128に対して、AVX-512の三項論理演算命令vpternlogdをかけてみた。vpternlogdにxmmを渡せなかったり、組み込み関数でxmmからzmmに簡単に変換できなかったりといろいろあったが、最終的には理想的なアセンブリが吐かれるコードを作ることができた。
しかし、AVX-512は便利そうな命令多いんだけど、いろいろ拡張があって、いざ使おうとすると「対応してません (Illegal instruction)」ってすげなく言われるのつらい。
上記のソースコードは以下においてある。