Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
283 lines (217 sloc) 9.01 KB
import math
from numba import (cuda, float32, float64, uint32, int64, uint64, from_dtype,
jit)
import numpy as np
# This implementation is based upon the xoroshiro128+ and splitmix64 algorithms
# described at:
#
# http://xoroshiro.di.unimi.it/
#
# and originally implemented by David Blackman and Sebastiano Vigna.
#
# The implementations below are based on the C source code:
#
# * http://xoroshiro.di.unimi.it/xoroshiro128plus.c
# * http://xoroshiro.di.unimi.it/splitmix64.c
#
# Splitmix64 is used to generate the initial state of the xoroshiro128+
# generator to ensure that small seeds don't result in predictable output.
# **WARNING**: There is a lot of verbose casting in this file to ensure that
# NumPy casting conventions (which cast uint64 [op] int32 to float64) don't
# turn integers into floats when using these functions in the CUDA simulator.
#
# There are also no function type signatures to ensure that compilation is
# deferred so that import is quick, and Sphinx autodoc works. We are also
# using the CPU @jit decorator everywhere to create functions that work as
# both CPU and CUDA device functions.
xoroshiro128p_dtype = np.dtype([('s0', np.uint64), ('s1', np.uint64)],
align=True)
xoroshiro128p_type = from_dtype(xoroshiro128p_dtype)
@jit
def init_xoroshiro128p_state(states, index, seed):
'''Use SplitMix64 to generate an xoroshiro128p state from 64-bit seed.
This ensures that manually set small seeds don't result in a predictable
initial sequence from the random number generator.
:type states: 1D array, dtype=xoroshiro128p_dtype
:param states: array of RNG states
:type index: uint64
:param index: offset in states to update
:type seed: int64
:param seed: seed value to use when initializing state
'''
index = int64(index)
seed = uint64(seed)
z = seed + uint64(0x9E3779B97F4A7C15)
z = (z ^ (z >> uint32(30))) * uint64(0xBF58476D1CE4E5B9)
z = (z ^ (z >> uint32(27))) * uint64(0x94D049BB133111EB)
z = z ^ (z >> uint32(31))
states[index]['s0'] = z
states[index]['s1'] = z
@jit
def rotl(x, k):
'''Left rotate x by k bits.'''
x = uint64(x)
k = uint32(k)
return (x << k) | (x >> uint32(64 - k))
@jit
def xoroshiro128p_next(states, index):
'''Return the next random uint64 and advance the RNG in states[index].
:type states: 1D array, dtype=xoroshiro128p_dtype
:param states: array of RNG states
:type index: int64
:param index: offset in states to update
:rtype: uint64
'''
index = int64(index)
s0 = states[index]['s0']
s1 = states[index]['s1']
result = s0 + s1
s1 ^= s0
states[index]['s0'] = uint64(rotl(s0, uint32(55))) ^ s1 ^ (s1 << uint32(14))
states[index]['s1'] = uint64(rotl(s1, uint32(36)))
return result
@jit
def xoroshiro128p_jump(states, index):
'''Advance the RNG in ``states[index]`` by 2**64 steps.
:type states: 1D array, dtype=xoroshiro128p_dtype
:param states: array of RNG states
:type index: int64
:param index: offset in states to update
'''
index = int64(index)
jump = (uint64(0xbeac0467eba5facb), uint64(0xd86b048b86aa9922))
s0 = uint64(0)
s1 = uint64(0)
for i in range(2):
for b in range(64):
if jump[i] & (uint64(1) << uint32(b)):
s0 ^= states[index]['s0']
s1 ^= states[index]['s1']
xoroshiro128p_next(states, index)
states[index]['s0'] = s0
states[index]['s1'] = s1
@jit
def uint64_to_unit_float64(x):
'''Convert uint64 to float64 value in the range [0.0, 1.0)'''
x = uint64(x)
return (x >> uint32(11)) * (float64(1) / (uint64(1) << uint32(53)))
@jit
def uint64_to_unit_float32(x):
'''Convert uint64 to float32 value in the range [0.0, 1.0)'''
x = uint64(x)
return float32(uint64_to_unit_float64(x))
@jit
def xoroshiro128p_uniform_float32(states, index):
'''Return a float32 in range [0.0, 1.0) and advance ``states[index]``.
:type states: 1D array, dtype=xoroshiro128p_dtype
:param states: array of RNG states
:type index: int64
:param index: offset in states to update
:rtype: float32
'''
index = int64(index)
return uint64_to_unit_float32(xoroshiro128p_next(states, index))
@jit
def xoroshiro128p_uniform_float64(states, index):
'''Return a float64 in range [0.0, 1.0) and advance ``states[index]``.
:type states: 1D array, dtype=xoroshiro128p_dtype
:param states: array of RNG states
:type index: int64
:param index: offset in states to update
:rtype: float64
'''
index = int64(index)
return uint64_to_unit_float64(xoroshiro128p_next(states, index))
TWO_PI_FLOAT32 = np.float32(2 * math.pi)
TWO_PI_FLOAT64 = np.float64(2 * math.pi)
@jit
def xoroshiro128p_normal_float32(states, index):
'''Return a normally distributed float32 and advance ``states[index]``.
The return value is drawn from a Gaussian of mean=0 and sigma=1 using the
Box-Muller transform. This advances the RNG sequence by two steps.
:type states: 1D array, dtype=xoroshiro128p_dtype
:param states: array of RNG states
:type index: int64
:param index: offset in states to update
:rtype: float32
'''
index = int64(index)
u1 = xoroshiro128p_uniform_float32(states, index)
u2 = xoroshiro128p_uniform_float32(states, index)
z0 = math.sqrt(-float32(2.0) * math.log(u1)) * math.cos(TWO_PI_FLOAT32 * u2)
# discarding second normal value
# z1 = math.sqrt(-float32(2.0) * math.log(u1))
# * math.sin(TWO_PI_FLOAT32 * u2)
return z0
@jit
def xoroshiro128p_normal_float64(states, index):
'''Return a normally distributed float32 and advance ``states[index]``.
The return value is drawn from a Gaussian of mean=0 and sigma=1 using the
Box-Muller transform. This advances the RNG sequence by two steps.
:type states: 1D array, dtype=xoroshiro128p_dtype
:param states: array of RNG states
:type index: int64
:param index: offset in states to update
:rtype: float64
'''
index = int64(index)
u1 = xoroshiro128p_uniform_float32(states, index)
u2 = xoroshiro128p_uniform_float32(states, index)
z0 = math.sqrt(-float64(2.0) * math.log(u1)) * math.cos(TWO_PI_FLOAT64 * u2)
# discarding second normal value
# z1 = math.sqrt(-float64(2.0) * math.log(u1))
# * math.sin(TWO_PI_FLOAT64 * u2)
return z0
@jit
def init_xoroshiro128p_states_cpu(states, seed, subsequence_start):
n = states.shape[0]
seed = uint64(seed)
subsequence_start = uint64(subsequence_start)
if n >= 1:
init_xoroshiro128p_state(states, 0, seed)
# advance to starting subsequence number
for _ in range(subsequence_start):
xoroshiro128p_jump(states, 0)
# populate the rest of the array
for i in range(1, n):
states[i] = states[i - 1] # take state of previous generator
xoroshiro128p_jump(states, i) # and jump forward 2**64 steps
def init_xoroshiro128p_states(states, seed, subsequence_start=0, stream=0):
'''Initialize RNG states on the GPU for parallel generators.
This initializes the RNG states so that each state in the array corresponds
subsequences in the separated by 2**64 steps from each other in the main
sequence. Therefore, as long no CUDA thread requests more than 2**64
random numbers, all of the RNG states produced by this function are
guaranteed to be independent.
The subsequence_start parameter can be used to advance the first RNG state
by a multiple of 2**64 steps.
:type states: 1D DeviceNDArray, dtype=xoroshiro128p_dtype
:param states: array of RNG states
:type seed: uint64
:param seed: starting seed for list of generators
'''
# Initialization on CPU is much faster than the GPU
states_cpu = np.empty(shape=states.shape, dtype=xoroshiro128p_dtype)
init_xoroshiro128p_states_cpu(states_cpu, seed, subsequence_start)
states.copy_to_device(states_cpu, stream=stream)
def create_xoroshiro128p_states(n, seed, subsequence_start=0, stream=0):
'''Returns a new device array initialized for n random number generators.
This initializes the RNG states so that each state in the array corresponds
subsequences in the separated by 2**64 steps from each other in the main
sequence. Therefore, as long no CUDA thread requests more than 2**64
random numbers, all of the RNG states produced by this function are
guaranteed to be independent.
The subsequence_start parameter can be used to advance the first RNG state
by a multiple of 2**64 steps.
:type n: int
:param n: number of RNG states to create
:type seed: uint64
:param seed: starting seed for list of generators
:type subsequence_start: uint64
:param subsequence_start:
:type stream: CUDA stream
:param stream: stream to run initialization kernel on
'''
states = cuda.device_array(n, dtype=xoroshiro128p_dtype, stream=stream)
init_xoroshiro128p_states(states, seed, subsequence_start, stream)
return states
You can’t perform that action at this time.