疑似乱数xoshiro/xoroshiroについて(2) 〜SIMDによる並列処理 by @STDNUL

前回のおさらい

前回の記事 - 疑似乱数xoshiro/xoroshiroについて(1) では、疑似乱数のxoshiro/xoroshiroシリーズは、ジャンプ関数を使うことで、系列の重複しない状態空間を低コストで生成することができ、これによってマルチスレッドやSIMDで並列動作させても安全に使うことができることについて言及している。

なお、この記事はApple M1の性能検証を兼ねているので Apple M1のサポートする命令などなんかも目を通していただけると幸いである。

AoSとSoA

どっちがどっちというと割といろんな解釈のしようもあるのだが、SIMDプログラミングの世界では、SIMDのベクトル幅に合わせて同じ要素をもったデータをひとまとめにして演算を行いやすくすることを「SoA化」とよく言う。

xoshiropシリーズにおいては、s[0], s[1], s[2] s[3]で演算の種類が異なっているが、そこで、s[0]同士、s[1]同士で纏めてSIMDで処理しやすくする。
xoshiro256の要素は64ビット単位なので、128ビットSIMDであれば2つ、256ビットSIMDであれば4つ束ねることになるであろう。

そんなわけで、128-bit Advanced SIMD(NEON)向けのデータ構造は以下の通りである。
ちょうど2つのxoshiro256の内部状態配列をインターリーブした構造になる。

union XOSHIRO256_X2_STATE {
    struct{
        uint64x2_t vs0, vs1, vs2, vs3;
    };
    uint64_t state64[8];
};
typedef union XOSHIRO256_X2_STATE xoshiro256_x2_state_t;

また、レイテンシの大きい演算を多用する場合はSIMDレジスタ数の整数倍の配列にして、演算をインターリーブすることもある。

余談: SFMTなども「再並列化」できる

2つとか4つ、或いはそれ以上のSFMTの状態ベクトルをインターリーブすることでAVX2やAVX-512でSIMD並列動作させるというアプローチも当然にありうる。
SFMTを並列化する場合は、MT配列の128ビットごとを2つとか4つごとにパックして256ビットSIMDや512ビットSIMDに適用するような使い方になるであろう。

SIMD化の留意点

IntelのSIMD実装では64ビットの乗算やビットローテート演算をサポートするのはAVX-512からである。
乗算に関しては、5倍は2ビットの左シフトと加算、9倍は3ビットの左シフトと加算で代替できる。
これも結局命令数が増えてしまうので、極力使いたくはない。
ビットローテートは、サポートされていない場合、左シフトして右シフトして論理和という素朴な方法をとるほかない。

64ビットARMアーキテクチャ向けのSIMD(NEON)実装

構造体は64ビットを2並列に取る SoA 構成を取る。
Apple以外のARMでは場合によってはSIMDを使わない方が速いことも十分ありうるが、ユニットの構成次第である。
SIMDユニットが強化されるCortex-X1などでは十分選択肢に入ってくる。

ローテートが無いので左シフト＋右シフト＋論理和で行うのだが、右シフト＋加算を1命令でこなす命令があるので、これと左シフトを組み合わせた2命令でローテートを行う（ローテート演算ではビットが重なる部分がないので論理和を加算で代用できる）。

もちろん乗算は左シフト＋加算で表現する。右シフト＋加算があるのなら左シフト＋加算を一括で行う命令が切実に欲しいと思うが、なぜか無い。
ビット論理演算や加算にシフトやローテートを組み合わせて1命令で処理できるARMスカラ命令と比べると、ARMの強みが活かせてなくて弱い感じがする。

#define rotl64_neon(x, n) vsraq_n_u64(vshlq_n_u64(x, n), x, (64-n));

__inline void next_neon_common(xoshiro256_x2_state_t* state) {
    uint64x2_t t = vshlq_n_u64(state->vs1, 17);
    uint64x2_t vs0 = state->vs0;
    uint64x2_t vs1 = state->vs1;
    uint64x2_t vs2 = state->vs2;
    uint64x2_t vs3 = state->vs3;

    vs2 = veorq_u64(vs2, vs0);
    vs3 = veorq_u64(vs3, vs1);
    vs1 = veorq_u64(vs1, vs2);
    vs0 = veorq_u64(vs0, vs3);
    vs2 = veorq_u64(vs2, t);
    vs3 = rotl64_neon(vs3, 45);

    state->vs0 = vs0;
    state->vs1 = vs1;
    state->vs2 = vs2;
    state->vs3 = vs3;
}

__inline uint64x2_t next_neon_plus(xoshiro256_x2_state_t* state) {
    uint64x2_t result = veorq_u64(state->vs0, state->vs3);
    next_neon_common(state);
    return result;
}

__inline uint64x2_t next_neon_plus_plus(xoshiro256_x2_state_t* state) {
    uint64x2_t result = vaddq_u64(state->vs0, state->vs3);
    result = rotl64_neon(result, 23); //rol 23;
    result = vaddq_u64(result, state->vs0);
    next_neon_common(state);
    return result;
}

__inline uint64x2_t next_neon_star_star(xoshiro256_x2_state_t* state) {
    //result = rotl(state->s[1] * 5, 7) * 9;
    uint64x2_t temp   = vshlq_n_u64(state->vs1, 2); // x4
    uint64x2_t result = vaddq_u64(state->vs1, temp); // x5
    result = rotl64_neon(result, 7); //roｔl 7;
    temp   = vshlq_n_u64(result, 3); // x8
    result = vaddq_u64(result, temp); // x9
    next_neon_common(state);
    return result;
}

命令数削減に威力を発揮するSHA3命令セット

威力は発揮するにはするのだが、 xoshiro256** は使いたくないな...というのが率直な感想。x86-64スカラ演算では xoshro** よりも速いというのはわかるのだが。

XAR(vxarq_u64)

64ビット2並列の RotateR((A XOR B), N) (Nは即値)を一度に可能にする命令である。ソースレジスタオペランドの一方に0を指定することで、右ローテート命令として使うこともできる。
もちろんこれがある以上は左シフトと右シフト＋加算命令の組み合わせを使う理由は微塵もない。
64ビット要素のSIMDでビットローテート操作が必要ならとりあえずXAR命令を使っておけという程度にはXARは便利である。なお、32ビット, 16ビット, 8ビット版はない。従ってxoshiro128の方では使えない。

EOR3(veor3q_u64)

MT/SFMTでも使える優等生。なぜ便利なのに組み込み関数を用意しないのか。一見重複してXORをとっているように見えて無駄に見えるが、どうやらEOR(veorq_u64)命令と同じレイテンシで2回分を実行できるのでこれで十分である。

コード

__inline void next_sha3_common(xoshiro256_x2_state_t* state) {
    uint64x2_t t = vshlq_n_u64(state->vs1, 17);
    uint64x2_t vs0 = veor3q_u64(state->vs0, state->vs1, state->vs3);
    uint64x2_t vs1 = veor3q_u64(state->vs1, state->vs0, state->vs2); 
    uint64x2_t vs2 = veor3q_u64(state->vs2, state->vs0, t);
    uint64x2_t vs3 = vxarq_u64(state->vs3, state->vs1, 19);

    state->vs0 = vs0;
    state->vs1 = vs1;
    state->vs2 = vs2;
    state->vs3 = vs3;
}

__inline uint64x2_t next_sha3_plus(xoshiro256_x2_state_t* state) {
    uint64x2_t result = veorq_u64(state->vs0, state->vs3);
    next_sha3_common(state);
    return result;
}

__inline uint64x2_t next_sha3_plus_plus(xoshiro256_x2_state_t* state) {
    uint64x2_t result = vaddq_u64(state->vs0, state->vs3);
    result = vxarq_u64(result, vdupq_n_u64(0), 41); //roｔl 23;
    result = vaddq_u64(result, state->vs0);
    next_sha3_common(state);
    return result;
}

__inline uint64x2_t next_sha3_star_star(xoshiro256_x2_state_t* state) {
    //result = rotl(state->s[1] * 5, 7) * 9;
    uint64x2_t temp   = vshlq_n_u64(state->vs1, 2); //x4
    uint64x2_t result = vaddq_u64(state->vs1, temp); //x5
    result            = vxarq_u64(result, vdupq_n_u64(0), 57); //roｔl 7
    temp              = vshlq_n_u64(result, 3); //x8
    result            = vaddq_u64(result, temp); //x9
    next_sha3_common(state);
    return result;
}

ベンチマーク

Apple M1上でxoshiro256++の生成速度を計測してみたのが以下の通りであるが、50,000,000回の乱数生成に対し生成時間以下の様になった。
通常のxoshiro256がSERIAL。さらにそれを4並列で呼び出し、演算のレイテンシの穴埋めをはかったものがx4 PARALLELであるが、SIMD版(NEON/SHA3)も同様に2並列(SIMD版単体)と8並列（2x4並列）を用意している。
50,000,000回にしたのは処理時間をSFMTと比較するためだが、標準のSFMTはインスタンスの並列化していないのでフェアじゃないかもしれない。

xoshiro256++ benchmark for 64-bit x 50000000 randoms generation
64-bit SERIAL:60.27ms
64-bit x4 PARALLEL:35.36ms
NEON 64-bit x2 VECTOR:62.81ms
NEON 64-bit x2 VECTOR x4 PARALLEL:33.52ms
NEON(SHA3) 64-bit x2 VECTOR:30.50ms
NEON(SHA3) 64-bit x2 VECTOR x4 PARALLEL:23.02ms

ローテート＋加算とかシフト＋ビット論理演算みたいな演算は、ARMの得意とするところなので、それが使えないどころかローテートすら2命令かかるNEON(ASIMD)はパフォーマンス的に苦しい。
言い換えるなら64ビットのARMはSIMD化しなくても十分速いのがxoshiro256++である。
そして8並列でも2048ビット（256バイト）しか消費しないので、MTより遥かに経済的と言える。

SIMD/マルチスレッドへの更なる応用

当然ジャンプ関数もSoA化された内部状態に対してSIMDで並列処理できる。xoshiro256には周期の違う2つのジャンプ関数があるので、例えばlong_jumpで周期の長いジャンプを行った状態空間配列をSoA化し、SIMD化したjump関数によって更なる並列実行インスタンスを作るなどの方法をとることができる。
その辺りは各自工夫してみて欲しい。