この記事では全国一億三千万のビット演算マニアのために、AVX512命令セットからビット演算に使えそうなものを紹介します。

AVX512の基本

AVX512とは

Advanced Vector Extensions(AVX)というx86/x64のベクトル演算拡張命令を512ビット幅に拡張したものです。普通のCPUではSkylake-SP/X/Wから使えるようになりました。 512bitのデータを一括で処理できます。ビット演算なら512並列ですよ512並列。

AVX512になって変わった主なこと

レジスタ幅が最大512bit(ZMMレジスタ)に（それはそう）
論理レジスタ数が16から32に倍増
マスクレジスタ(k0-k7)の追加
各種ビット操作・バイト操作命令の追加

AVX512の命令セット

AVX512は複数の要素から構成されており、プロセッサによってサポートの状況が異なります。

以下に書くサブセットとその説明を書きます。

AVX512-F (Foundation) 基本命令セット。AVX512をサポートするプロセッサは必ずこれをサポートする。
AVX512-CD (Conflict Detection) データの衝突検出。Hist処理とかがベクトル化できる。 vpconflictd - Qiita
AVX512-DQ (Doubleword and Quadword) 32bit/64bit単位での処理。
AVX512-BW (Byte and Word) 8bit/16bit単位での処理。
AVX512-VL (Vector Length) 128bit(XMM)/256bit(YMM)レジスタを操作する。

最新のCPU(Skylake-SP/X/W)でもまだ使えない拡張もたくさんあります。

AVX512-PF (Pre Fetch) プリフェッチ命令。
AVX512-ER (Exponential and Reciprocal) 指数関数や逆数関数。
AVX512-IFMA52 (Integer Fused Multiply Add) 整数の積和算を倍精度の積和算演算器を借りて52ビット精度で実行する。
AVX512-VBMI (Vector Byte Manipulation Instructions) BWに入っていないbyte操作の命令。
AVX512-VNNI (Vector Neural Network Instructions) ニューラルネットワーク処理用の整数演算命令。16bit整数同士の積を32bit整数と足す。
AVX512-4VNNIW (Vector Neural Network Instructions Word variable precision) ニューラルネットワーク処理用の整数演算命令。16bit整数同士の積を32bit整数と足すのを4回繰り返す。
AVX512-4FMAPS (Fused Multiply Accumulation Packed Single precision) 単精度の内積を4回繰り返す。これもニューラルネットワーク処理用。
AVX512-VPOPCNTDQ (Vector POP CouNT) ベクトル化されたpopcount。
AVX512-VBMI2 (Vector Byte Manipulation Instructions2) BWにもVBMIにも入っていないbyte操作の命令。
AVX512-BITARG (不明) ビット操作関連の命令。

マスク演算

AVX512では比較演算の結果などはマスクレジスタと呼ばれる特殊なレジスタに格納されます。

AVX512では命令にマスクとしてマスクレジスタを指定することができます。

VPANDQ (and演算, 64bit単位のマスク)を例にして疑似コードで表すと、

for j in range(0, 8):
  i = j*64
  if mask[j]:
    dest[i+63:i] = src[i+63:i]
  else:
    dest[i+63:i] = a[i+63:i] AND b[i+63:i]

このようにマスクでビットが立っている場所を演算せずに値を残すことが可能です。また、値を残す代わりに0にすることもできます。

intrinsicでは

dest = _mm512_and_epi64(a, b); // マスク無し
dest = _mm512_mask_and_epi64(src, mask, a, b); // マスクあり（マスクされたらsrcを残す）
dest = _mm512_maskz_and_epi64(mask, a, b); // マスクあり（マスクされたら0）

のように使います。

マスク同士の演算

このマスクレジスタはSIMD レジスタ(XMM/YMM/ZMMレジスタ)とも通常のレジスタとも違うので、通常のレジスタとマスクレジスタとの間で値をやりとりするのにもコストがかかります。そこで、いくつかの演算(and, or, addなど)はマスクレジスタ同士で行えるようになっています。

命令の解説

AVX512-F

VPTERNLOG{D,Q}

ビット単位の任意の3項演算ができます。演算内容は即値で与えます。 A & (B | ~C)なら

A	B	C	A&(B\|~C)
1	1	1	1
1	1	0	1
1	0	1	0
1	0	0	1
0	1	1	0
0	1	0	0
0	0	1	0
0	0	0	0

なので即値として0b11010000 = 0xD0を与えてやれば良いです。 intrinsicは_mm512_ternarylogic_epi64(A, B, C, 0xD0)のように使います。これを使えば他のビット論理演算命令いらないじゃんとなりそうですが、怪しいページによるとスループットが2命令/サイクルとandやorなどの3命令/サイクルに比べると劣るので適材適所で使うことになりそうです。

VP{AND,OR,XOR,ANDNOT}{D,Q}

and, orなどの演算です。マスクが32bit/64bit単位で指定できます。

VPS{L,R}{L,A}{,V}{W,D,Q}

16bit/32bit/64bit単位のビット{論理,算術}シフトです。V付きはシフト量もベクトルで指定することで、各データごとに異なるシフト量を指定できます。

VPRO{L,R}{,V}{D,Q}

待望のローテート命令です。でも32bit/64bitだけ。

AVX512-CD

VPLZCNT{D,Q}

leading zero countです。上位ビットから見て最初に1が立っているビットの位置を調べる。整数→浮動小数の変換で前からほぼ同じことができましたが、キャストとか0の特別扱いとか諸々せずに済むようになりました。ところでなんでConflict Detectionに入っているんでしょうか。

AVX512-VBMI

VPERMI2B

各バイトの下位7bitをインデックスにして、２つのZMMレジスタをくっつけた128bytesの中から表引きします。これまでの128bit境界に捕われていたり、即値で指定しないといけなかったりしたシャッフル命令とはおさらばできます。とはいえスループットやレイテンシが大きかったら従来のシャッフル命令も併用しないといけなさそうですね。

VPMULTISHIFTQB

各64bit要素に対して各バイトの分だけ右シフトしたデータを返します。

for i in range(0, 8):
  q = i*64
  for j in range(0, 8):
    tmp8 = 0
    ctrl = src1[q+j*8+7:q+j*8] & 0x3F
    for l in range(0, 8):
      tmp8[l] = src2[q+((ctrl+l)&0x3F)]
    dst[q+j*8+7:q*j*8] = tmp[7:0]

AVX512-VBMI2

VPSH{L,R}D{,V}{W,D,Q}

２つの値をつなげてシフトします。多倍長シフトとかに使えますね。V付きの命令はシフト量もベクトルで指定することで、各データごとに異なるシフト量を指定できます。

AVX512-POPCNT

VPOPCNT{D,Q}

32bit/64bit単位でのpopulation count（1になっているビットの数を数える）です。

AVX512-BITARG

VPOPCNT{B,W}

8bit/16bit単位でのpopcountです。なぜかVPOPCNT{D,Q}とは別の拡張になっています。

VPSHUFBITQMB

~~謎命令。~~各64bit要素に対して、各バイトをインデックスとして64bitデータの指定したビットを取ってきてマスクにします。疑似コードを見るのが早いと思います。

for i in range(0, 8): # qword
  for j in range(0, 8): # byte
    m = src2[i*64+j*8+7:i*64+j*8] & 0x3F
    res[i*8+j] = src1[i*64+m]
return res

AVX512以外の拡張

VPCLMULQDQ

繰り上がりなしの64bit乗算を並列で行います。ビット演算的にはCRCの計算とかpclmulqdqを用いたビット単位のunpack【ビット演算テクニック Advent Calendar 2016 11日目】 - prime's diaryとかに使えそうですね。

便利Intrinsic

_mm512_reduce_{or,and}_epi{32,64}

32bit/64bitデータ16/8個のOR/ANDを取ります。

_mm512_reduce_and_epi64なら、

reduced[63:0] = 0xFFFFFFFFFFFFFFFF
for j in range(0, 8):
    i = j*64
    reduced[63:0] = reduced[63:0] AND a[i+63:i]
return reduced[63:0]

まとめ

AVX512には~~どう使うかよくわからない~~便利そうな命令がたくさん追加されています。

使えるCPUを手に入れた暁にはうまく使ってビット演算ライフを満喫しましょう！

明日のKMCアドベントカレンダーはwass88さんによる「役に立たないこと」の予定です。お楽しみに！

adventar.org

prime's diary

そすうの日々を垂れ流しちゃうやつだよ

ビット演算マニアのためのAVX512入門【KMCアドベントカレンダー 22日目】