dSFMTをAVX-512に対応させる

What is dSFMT？

dSFMTは数値シミュレーションに特化した倍精度浮動小数点数型の疑似乱数をダイレクトに生成できるMersenne Twisterの派生版です。
スパコン富岳の前の世代「京」のKMATH_RANDOMなんかにも採用実績があります。

ただこのDSFMTは128ビットを超える並列処理を通常想定していないため、そのままでは素直に256ビットや512ビット単位の並列処理をすることができません。
2倍とか4倍のメモリ空間を食うけど複数の状態空間を並列に回すと言う解決策はあるにはあるけどあんまり美しくありません。

これをどうにかしたと言う話をまとめます。

そのままだとうまく並列化できない

もともと128ビット単位の逐次処理を想定した乱数アルゴリズムなので当然といえば当然なのですが
lung変数の更新の部分でシャッフル＋XORの依存関係が連鎖していて、そのままじゃ並列化できないんです。

0	1	2	3
WordReverse(lung)
XOR (lung, Z[0])
	WordReverse(lung)
	XOR(lung, Z[1])
		WordReverse(lung)
		XOR(lung, Z[2])
			WordReverse(lung)
			XOR(lung, Z[3])

ここでいうWordReverseというのは、SSE2のコードでいうこれです。

    y = _mm_shuffle_epi32(u->si, SSE2_SHUFF);

不幸中の幸い、このシャッフル操作は可逆で、もう一度シャッフルするともとにもどすことができます。ここ重要なところ。

lungと排他的論理和を取る値は先出しで計算できる

XORは演算順序の交換法則及びビットシフトやシャッフル演算他その他のビット演算に対する分配法則が成り立ちます。また $WordReverse(WordReverse(n)) = n$ になります。すなわち

lung_{n+1} = WordReverse(lung_n) \oplus X_{n}

lung_{n+2} 
= WordReverse(lung_{n+1}) \oplus X_{n+1}
= lung_n  \oplus WordReverse(X_{n}) \oplus X_{n+1}

あとは帰納的に計算するだけですね。lungが不明の状態でも、lungまたはその反転値とXORをとれば求める値がすぐに求められるので、依存関係を断ち切って並列化することができます。
つまり(a << SL1) ^ bの方にWordReverseをかけて階段状に処理するように書き換えてしまいましょう。

Operation	Lane:0	Lane:1	Lane:2	Lane:3
r := WordReverse( ShiftLeft128(z) )	0	WordReverse(z[0])	WordReverse(z[1])	WordReverse(z[2])
z := XOR(z, r))	z[0]	z[1] ^ r[1]	z[2] ^ r[2]	z[3] ^ r[3]
r := ShiftLeft256(z)	0	0	z[0]	z[1]
y = Expand(lung)	WordReverse(lung)	lung	WordReverse(Lung)	lung
z := XOR3(z, y, r)	z[0] ^ y[0]	z[1] ^ y[1]	z[2] ^ y[2] ^ r[2]	z[3] ^ y[3] ^ r[3]

いまこの部分の説明をしています。

inline static __m512i do_recursion_x4_avx512( __m512i a, __m512i b, __m512i* u) {
    const __m512i shift128_reverse_idx = 
        _mm512_setr_epi32(0, 0, 0, 0, 19,18,17,16, 23,22,21,20, 27,26,25,24);
    const __m512i expand_lung_idx = 
        _mm512_setr_epi32(15,14,13,12, 12,13,14,15, 15,14,13,12, 12,13,14,15);
        
    __m512i v, w, x, y, z;

    z = _mm512_slli_epi64(a, DSFMT_SL1);
    z = _mm512_xor_si512(z, b);

    /* ↓↓↓↓↓ここから↓↓↓↓↓ */
    x = _mm512_permutex2var_epi32(_mm512_setzero_si512(),shift128_reverse_idx , z);
    z = _mm512_xor_si512(z, x);
    x = _mm512_inserti64x4(_mm512_setzero_si512(), _mm512_castsi512_si256(z), 1);
    y = _mm512_permutexvar_epi32(expand_lung_idx, _mm512_load_epi64(u));  /* lung */
    y = _mm512_ternarylogic_epi32(y, x, z, 0x96);
    _mm512_store_epi64(u, y);
    /* ↑↑↑↑↑ここまで↑↑↑↑↑*/

    v = _mm512_srli_epi64(y, DSFMT_SR);
    w = _mm512_and_si512(y, _mm512_broadcast_i32x4(sse2_param_mask.i128));
    v = _mm512_ternarylogic_epi32(v, a, w, 0x96);
    return v;

あっさり並列化できてしまいました。ビット論理演算の交換法則って偉大ですね。困った時はカルノーマップでも書いてみるといいかもしれません。

これでlungの依存関係を気にすることなく並列処理できますね。
ただし、いくらでも並列化できるかというとそうでもなく、依存関係があるのでDSFMT_N - DSFMT_POS個先までしか先読みできません。

性能評価

御宅はいいからまず性能をって話ですよね。はい。

• CPU: Core i7-1068NG7
• OS: macOS 13.4
• コンパイラ: Apple clang version 13.1.6 (clang-1316.0.21.2.5)

現行版(SSE2)

まず現行最新のdSFMTのSSE2版のベンチマークが以下のような感じです。

% ./test-sse2-M2203 -s
consumed time for generating 100000000 randoms.
ST BLOCK [0, 1) AVE:  52ms.
ST BLOCK (0, 1] AVE:  59ms.
ST BLOCK (0, 1) AVE:  58ms.
ST BLOCK [1, 2) AVE:  50ms.
ST SEQ [0, 1) 1 AVE: 126ms.
ST SEQ [0, 1) 2 AVE: 123ms.
total = 500017062.435534
ST SEQ (0, 1] 1 AVE: 129ms.
ST SEQ (0, 1] 2 AVE: 122ms.
total = 500032937.564466
ST SEQ (0, 1) 1 AVE: 133ms.
ST SEQ (0, 1) 2 AVE: 137ms.
total = 500017062.435534
ST SEQ [1, 2) 1 AVE: 117ms.
ST SEQ [1, 2) 2 AVE: 112ms.
total = 1500067062.434869

コンパイルオプションのみでのAVX-512最適化

そして上記のSSE2のコードをAVX-512を使うようにしたビルドでテストしてみます。

% ./test-avx512-M2203 -s
consumed time for generating 100000000 randoms.
ST BLOCK [0, 1) AVE:  45ms.
ST BLOCK (0, 1] AVE:  45ms.
ST BLOCK (0, 1) AVE:  45ms.
ST BLOCK [1, 2) AVE:  43ms.
ST SEQ [0, 1) 1 AVE: 118ms.
ST SEQ [0, 1) 2 AVE: 113ms.
total = 500017062.435534
ST SEQ (0, 1] 1 AVE: 118ms.
ST SEQ (0, 1] 2 AVE: 110ms.
total = 500032937.564466
ST SEQ (0, 1) 1 AVE: 120ms.
ST SEQ (0, 1) 2 AVE: 125ms.
total = 500017062.435534
ST SEQ [1, 2) 1 AVE: 117ms.
ST SEQ [1, 2) 2 AVE: 106ms.
total = 1500067062.434869

これも最新のclangコンパイラはビット論理演算(pand/pandn/por/pxor)が２つ並んだ場合に１つのvpternlogqにまとめる最適化程度はやってくれるようで、ブロック生成の場合で1〜2割くらいきいています。非SIMD処理で自動ベクトル化されている部分もあったりします。
割と優秀ですね。

更新関数4並列化実装

そして本題のこれ。クロック可変なので時間にはばらつきはありますが、ブロック生成でおおよそSSE2の半分程度の時間（2倍の生成速度）、シーケンシャル生成でも微妙に性能改善しているっぽいのがいい感じです。
これでdSFMT2203と全く同じ乱数列を生成しますから質は十分担保されています。

consumed time for generating 100000000 randoms.
ST BLOCK [0, 1) AVE:  27ms.
ST BLOCK (0, 1] AVE:  26ms.
ST BLOCK (0, 1) AVE:  27ms.
ST BLOCK [1, 2) AVE:  24ms.
ST SEQ [0, 1) 1 AVE: 116ms.
ST SEQ [0, 1) 2 AVE:  95ms.
total = 500017062.435534
ST SEQ (0, 1] 1 AVE: 116ms.
ST SEQ (0, 1] 2 AVE:  95ms.
total = 500032937.564466
ST SEQ (0, 1) 1 AVE: 121ms.
ST SEQ (0, 1) 2 AVE: 102ms.
total = 500017062.435534
ST SEQ [1, 2) 1 AVE: 123ms.
ST SEQ [1, 2) 2 AVE:  88ms.
total = 1500067062.434869

512ビット（更新関数４並列実行）化可能なdSFMTのバージョン

今回2203だけを512ビット化したのは、ZMMレジスタ5本（AVX2でもYMMレジスタ10本）の状態空間に収まるのでオンレジスタで処理できて高速化しやすいという点です。
4253は1024ビットSIMDまで見据えるならありかも、と言う感じです。
オンレジスタ処理を諦め、端数が出て良いのであれば11213以上も十分検討範囲です。

端数はマスクレジスタを使ったり128ビットSIMDで処理すればよいのでそれほど気にすることはありません。コードが複雑になる点だけは頭が痛いところですが。

	DSFMT_N	DSFMT_POS1	512-bit SIMD Capable
521	4	3	×
1279	12	9	×
2203	20	7	◎
4253	40	19	◎
11213	107	37	○
19937	191	117	○
44997	427	304	○
86243	829	231	○
132049	1269	371	○
216091	2077	1890	○

小言

ARMスパコン（SVE）なんかへの移植とかはご自由に。
場合分けが必要になってベクタ長可変のメリットなくなるんじゃないかなと思うんですけど。

今回のコードブランチは以下に挙げております
https://github.com/magurosan/dSFMT/tree/AVX512