Dotnet developer, C++ learner, Unreal Engine lover
One popular random number generator that you can find out in the wild is the xoroshiro128 — or the xor-rotate-shift-rotate — by David Blackman and Sebastiano Vigna (source). It uses a state of four 32-bit integers (hence the 128-part in the name), and a combination of xor and bit-shifts to advance the state. (Pseudorandom) generators based on xor and shifts seem to be fast and have relatively good properties. I am by no means an expert on this topic, but Xorshift RNGs is one paper that describes some of the properties of these generators.
The xoroshiro128++ mentioned in the previous paragraph is of course not the only one - in fact, there are many similar algorithms as the reader can see for him/herself: A PRNG Shootout, but in my experiments I have picked this one to test with (but may play around with others later).
xoroshiro128++ has some nice properties for relevant for multithreading, as we can use the jump and long_jump functions to advance the state and create a large number of non-overlapping (up to a point) starting points for processing in different threads. In my experiment, I use this to produce non-overlapping starting points for vectorized number generation.
I am assuming that the reader is on approximately the same level (a.k.a. “still so much to learn”) as I am. This should be a reasonable assumption as expert programmers probably wouldn’t be reading this anyway :)
Since I am still learning about SIMD, this article may contain errors or mistakes of varying degree. I hope the reader will take the time to comment if they notice something that is totally wrong!
Implementation
Before we look at the vectorized generator, let’s take a look at the original implementation. It is in fact quite short; it starts with a structure holding the state (this is an adaption by me) and a helper function (source):
namespace xoroshiro128
{
struct state_t
{
uint32_t v[4];
};
inline uint32_t rotl(const uint32_t x, int k)
{
return (x << k) | (x >> (32 - k));
}
}
After this we then have the function that generates the actual values:
namespace xoroshiro128
{
inline uint32_t next(state_t* state)
{
const uint32_t result = rotl(state->v[0] + state->v[3], 7) + state->v[0];
const uint32_t t = state->v[1] << 9;
state->v[2] ^= state->v[0];
state->v[3] ^= state->v[1];
state->v[1] ^= state->v[2];
state->v[0] ^= state->v[3];
state->v[2] ^= t;
state->v[3] = rotl(state->v[3], 11);
return result;
}
}
The parameter should perhaps be a reference, but it was a raw pointer in the original so I left it like that. I’m not a mathematician, so I won’t claim to understand why it works, but all we need to know is really that if we shift and xor in this way, the resulting numbers will have good properties. What we as programmers care more about is how to use it, which couldn’t be simpler:
xoroshiro128::state_t rng_state{ 1,2,3,4 };
uint32_t v = xoroshiro128::next(&rng_state);
There are also the jump and jump_long functions, but they are omitted here for brevity.
Vectorizing it
If you stare at the next function for a minute, you realize that it should be possible to vectorize this with relatively little effort, as the instructions in use are only simple bit-shifts and xors. Since I’m still learning about SIMD this is what I set out to do.
My approach requires a slight change of the structure of the generator state; so instead of using an array of individual state_t structures, I will break this up and use a structure of arrays:
namespace xoroshiro128::avx512
{
struct states_t
{
static constexpr size_t size = 16;
std::array<uint32_t, size> _0{};
std::array<uint32_t, size> _1{};
std::array<uint32_t, size> _2{};
std::array<uint32_t, size> _3{};
};
}
This may feel a bit awkward to some — using something like std::array<state_t, size> states{} may be more intuitive to some depending on where you come from (not geographically …), but I assure you that it’s actually quite neat!
The structure is such that a state_t in the previous sense “cuts across” all arrays, e.g., _0[0], _1[0], _2[0], _3[0] represent one generator state when combined. The array member have simply been named after the indices of the single 4-element array in state_ t.
In effect, states_t hold 16 state_t structures! This allows us to simplify the load and store operations to and from the SIMD registers. If we assume that we have a function that takes in a states_t like so:
inline void next(states_t& in_states, std::array<uint32_t, 16>& out_values)
We can then load the states into registers in four lines of code:
__m512i s_0 = _mm512_loadu_epi32(in_states._0.data());
__m512i s_1 = _mm512_loadu_epi32(in_states._1.data());
__m512i s_2 = _mm512_loadu_epi32(in_states._2.data());
__m512i s_3 = _mm512_loadu_epi32(in_states._3.data());
I hope that the reader now sees why it is so powerful to use a structure of arrays instead of an array of structures - the data is already laid out properly in memory, so we don’t need to jump around to find all that goes into one register.
How do we know which intrinsic to use? We can use the Intel Intrinsics Guide to filter, and in this case I knew I wanted to use the 512-bit registers (availability of which may depend on your computer), I’m working with 32-bit unsigned integers and we want to load them into a register (represented by the type __m512i). Just like assembly programming, we have to manually load the data to a register before we can use it. When we are done we store it back using … you guessed it! store instead of save. After looking at a few of these you begin to be able to guess what the name should be, so I don’t think there is a need to memorize them all.
After loading the data into our register, we must perform the left rotate (as I assume rotl stands for …) and generate the results:
__m512i intermediate; rotl<7>(_mm512_add_epi32(s_0, s_3), intermediate);
__m512i const result = _mm512_add_epi32(intermediate, s_0);
_mm512_storeu_epi32(out_values.data(), result);
// Corresponds to:
// uint32_t const result = xoroshiro128::rotl(state->v[0] + state->v[3], 7) + state->v[0];
If we assume for the moment that rotl is given, then all that the original code did here was to add some numbers. This we can do using the _mm512_add_epi32 intrinsic (again searching the intrinsics guide for add and filter on instruction set). At the end we store all 16 numbers to the output variable - again this is simple as the array provides us with a way to reach the underlying storage.
And with this we are half done actually, we have generated 16 unsigned integers! What remains to be done is to advance the state of all generators simultaneously. This we do, as above, by mapping the existing lines to SIMD instructions. After computing the result in the original code, the next section contains a number of bitwise xor operations. This we can do using _mm512_xor_epi32, and since all lines look “the same” it is a simple copy-paste operation:
s_2 = _mm512_xor_epi32(s_2, s_0); // state->v[2] ^= state->v[0];
s_3 = _mm512_xor_epi32(s_3, s_1); // state->v[3] ^= state->v[1];
s_1 = _mm512_xor_epi32(s_1, s_2); // state->v[1] ^= state->v[2];
s_0 = _mm512_xor_epi32(s_0, s_3); // state->v[0] ^= state->v[3];
The next xor is preceded by the calculation of an intermediate variable using a left-shift. Again, a direct mapping to SIMD:
__m512i const t = _mm512_rol_epi32(s_1, 9); //const uint32_t t = state->v[1] << 9;
s_2 = _mm512_xor_epi32(s_2, t); // state->v[2] ^= t;
Finally another left-rotation:
rotl<11>(s_3, s_3); // state->v[3] = xoroshiro128::rotl(state->v[3], 11);
Finally we need to store the state back to the original arrays, which is, as we have seen before, very easy with the memory layout we have chosen:
_mm512_storeu_epi32(in_states._0.data(), s_0);
_mm512_storeu_epi32(in_states._2.data(), s_1);
_mm512_storeu_epi32(in_states._2.data(), s_2);
_mm512_storeu_epi32(in_states._3.data(), s_3);
Even if the reader was not familiar with SIMD when opening this article, these instructions and how to find them should be somewhat familiar by now. The full code looks like this:
/**
* Given 16 input states, generates 16 unsigned 32-bit integers, and advances all states.
* Assumes AVX512 instructions to be available on the target cpu.
*/
inline void next(states_t& in_states, std::array<uint32_t, 16>& out_values)
{
__m512i s_0 = _mm512_loadu_epi32(in_states._0.data());
__m512i s_1 = _mm512_loadu_epi32(in_states._1.data());
__m512i s_2 = _mm512_loadu_epi32(in_states._2.data());
__m512i s_3 = _mm512_loadu_epi32(in_states._3.data());
// First generate the numbers
// uint32_t const result = xoroshiro128::rotl(state->v[0] + state->v[3], 7) + state->v[0];
__m512i intermediate; rotl<7>(_mm512_add_epi32(s_0, s_3), intermediate);
__m512i const result = _mm512_add_epi32(intermediate, s_0);
_mm512_storeu_epi32(out_values.data(), result);
// Advance generator states
s_2 = _mm512_xor_epi32(s_2, s_0); // state->v[2] ^= state->v[0];
s_3 = _mm512_xor_epi32(s_3, s_1); // state->v[3] ^= state->v[1];
s_1 = _mm512_xor_epi32(s_1, s_2); // state->v[1] ^= state->v[2];
s_0 = _mm512_xor_epi32(s_0, s_3); //state->v[0] ^= state->v[3];
__m512i const t = _mm512_rol_epi32(s_1, 9); //const uint32_t t = state->v[1] << 9;
s_2 = _mm512_xor_epi32(s_2, t); // state->v[2] ^= t;
rotl<11>(s_3, s_3); // state->v[3] = xoroshiro128::rotl(state->v[3], 11);
_mm512_storeu_epi32(in_states._0.data(), s_0);
_mm512_storeu_epi32(in_states._2.data(), s_1);
_mm512_storeu_epi32(in_states._2.data(), s_2);
_mm512_storeu_epi32(in_states._3.data(), s_3);
}
What remains now is the implementation of the rotl function. Recall that the code was very simple, consisting of two shifts and one bitwise or. By searching the intrinsics guide we can find direct correspondences to these operations:
template <uint8_t K>
inline void rotl(__m512i const in_val, __m512i& out_val)
{
auto const lhs = _mm512_rol_epi32(in_val, K); // (x << k)
auto const rhs = _mm512_ror_epi32(in_val, 32-K); // (x >> (32 - k))
out_val = _mm512_or_epi32(lhs, rhs);
}
Here the parameter k that appeared in the original version has been converted into a template parameter. This is because I got an error stating that the second argument to _mm512_rol_epi32 must be an immediate value - so I made it a compile-time constant.
Using this is very simple, although the initialization takes a few more lines. Here is an example of both initialization and usage:
xoroshiro128::avx512::states_t generators{};
{
xoroshiro128::state_t initialization_state{1, 2, 3, 4};
for (uint32_t const i: std::ranges::views::iota(0u, xoroshiro128::avx512::states_t::size))
{
xoroshiro128::long_jump(&initialization_state);
generators._0[i] = initialization_state.v[0];
generators._1[i] = initialization_state.v[1];
generators._2[i] = initialization_state.v[2];
generators._3[i] = initialization_state.v[3];
}
}
std::array<uint32_t, xoroshiro128::avx512::states_t::size> result{};
xoroshiro128::avx512::next(generators, result);
long_jump could be vectorized as well, but you probably won’t need to call that as much as you do next.
That’s All Folks!
We have seen how to manually vectorize a pseudo-random generator. The original implementation only uses simple operations that map directly to SIMD intrinsics, which turn the vectorization int more of a mapping operation. This allows us to perform this transformation without needing to understand what actually makes the generator tick, although in reality you should probably add some unit- or acceptance tests to be sure no mistakes were made.
There could of course be mistakes in the above code, so maybe I should take my own advice and write some tests. In the future I would like to try other methods to generate numbers, such as moving the computations to the GPU. In fact, the goal of this experiment was to see at what rate I could generate pseudo/random numbers. To that end, here are the results of a google benchmark of the two implementations ran on my computer:
Run on (16 X 3793 MHz CPU s)
CPU Caches:
L1 Data 32 KiB (x8)
L1 Instruction 32 KiB (x8)
L2 Unified 1024 KiB (x8)
L3 Unified 16384 KiB (x1)
--------------------------------------------------------------------------------------------------------
Benchmark Time CPU Iterations UserCounters...
--------------------------------------------------------------------------------------------------------
BM_SingleThread/repeats:25/manual_time_mean 122 ms 121 ms 25 Num generated=108.687M/s
BM_SingleThread/repeats:25/manual_time_median 143 ms 141 ms 25 Num generated=87.6607M/s
BM_SingleThread/repeats:25/manual_time_stddev 27.8 ms 27.5 ms 25 Num generated=27.236M/s
BM_SingleThread/repeats:25/manual_time_cv 22.85 % 22.68 % 25 Num generated=25.06%
BM_Avx512/repeats:25/manual_time_mean 943 ms 939 ms 25 Num generated=1.07835G/s
BM_Avx512/repeats:25/manual_time_median 896 ms 891 ms 25 Num generated=1.11629G/s
BM_Avx512/repeats:25/manual_time_stddev 147 ms 144 ms 25 Num generated=120.893M/s
BM_Avx512/repeats:25/manual_time_cv 15.58 % 15.32 % 25 Num generated=11.21%
To summarize, the original implementation generated numbers at about 88 M/s (stddev 30 M/s), whereas the vectorized version generated about 1.1 G/s! Of course, we should be careful with benchmarks like these, but still it’s fun to do I think - it tickles something deep inside.
The code is really messy but can be found in github.
Thanks for reading this far!
