From 27565a007cd9ceafe4fc58f5569d0bee8d1a286c Mon Sep 17 00:00:00 2001 From: notoraptor Date: Tue, 3 Dec 2019 11:35:44 -0500 Subject: [PATCH] Add implementation for RNG-MRG, using MRG31k3p random number generator. --- myia/rng_mrg.py | 290 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_rng.py | 89 ++++++++++++++ 2 files changed, 379 insertions(+) create mode 100644 myia/rng_mrg.py create mode 100644 tests/test_rng.py diff --git a/myia/rng_mrg.py b/myia/rng_mrg.py new file mode 100644 index 000000000..45dc65340 --- /dev/null +++ b/myia/rng_mrg.py @@ -0,0 +1,290 @@ +""" +Implementation of MRG31k3p random number generator for Myia. + +Based on Theano implementation (2019/11/05): +https://github.com/Theano/Theano/blob/master/theano/sandbox/rng_mrg.py + +While Theano implementation uses many parallel streams in a matrix to +generate numbers, current myia implementation below uses only +1 stream vector for simplicity. + +Generator code in SSJ package (L'Ecuyer & Simard): +http://www.iro.umontreal.ca/~simardr/ssj/indexe.html +Page up-to-date (2019/11/05): +http://simul.iro.umontreal.ca/ssj/indexe.html +Original Github project page (2019/11/05): +https://github.com/umontreal-simul/ssj +Original JAVA implementation on Github (2019/11/05): +https://github.com/umontreal-simul/ssj/blob/master/src/main/java/umontreal/ssj/rng/MRG31k3p.java + +The MRG31k3p algorithm was published in: + +P. L'Ecuyer and R. Touzin, Fast Combined Multiple Recursive Generators with +Multipliers of the form a = +/- 2^d +/- 2^e, Proceedings of the 2000 Winter +Simulation Conference, Dec. 2000, 683-689. + +The conception of the multi-stream from MRG31k3p was published in: + +P. L'Ecuyer and R. Simard and E. Jack Chen and W. David Kelton, +An Object-Oriented Random-Number Package with Many Long Streams +and Substreams, Operations Research, volume 50, number 6, 2002, 1073-1075. +""" + +from dataclasses import dataclass + +import numpy as np + +# ---------- +# Constants. +# ---------- + +M1 = np.int32(2147483647) # 2^31 - 1 +M2 = np.int32(2147462579) # 2^31 - 21069 +MASK12 = np.int32(511) # 2^9 - 1 +MASK13 = np.int32(16777215) # 2^24 - 1 +MASK2 = np.int32(65535) # 2^16 - 1 +MULT2 = np.int32(21069) +A1p72 = np.asarray([[1516919229, 758510237, 499121365], + [1884998244, 1516919229, 335398200], + [601897748, 1884998244, 358115744]], + dtype='int64') +A2p72 = np.asarray([[1228857673, 1496414766, 954677935], + [1133297478, 1407477216, 1496414766], + [2002613992, 1639496704, 1407477216]], + dtype='int64') +A1p134 = np.asarray([[1702500920, 1849582496, 1656874625], + [828554832, 1702500920, 1512419905], + [1143731069, 828554832, 102237247]], + dtype='int64') +A2p134 = np.asarray([[796789021, 1464208080, 607337906], + [1241679051, 1431130166, 1464208080], + [1401213391, 1178684362, 1431130166]], + dtype='int64') + +i0 = np.int32(0) +i7 = np.int32(7) +i9 = np.int32(9) +i15 = np.int32(15) +i16 = np.int32(16) +i22 = np.int32(22) +i24 = np.int32(24) + +# Hashable versions of useful nd-arrays above +# (to avoid hash error during myia compilation). +a_1_p_134 = tuple( + tuple(np.int64(A1p134[i, j]) for j in range(3)) for i in range(3)) +a_2_p_134 = tuple( + tuple(np.int64(A2p134[i, j]) for j in range(3)) for i in range(3)) + + +# Tuple of d-types supported by RNG, and index for each type in tuple. +SUPP_TYPES = (np.float16, np.float32, np.float64) +FLOAT16, FLOAT32, FLOAT64 = 0, 1, 2 + + +# ------------------ +# Private functions. +# ------------------ + +def _mat_vec_mod_m(a, i, j, k, m): + """ + Specific matrix-vector operation used to increment random state. + + Given 3x3 matrix `a` and vector v = (i, j, k), compute: + output = numpy.sum((a * v) % m, axis=1) % m + :param a: a 2-d tensor (matrix) with shape (3, 3) + :param i: a scalar + :param j: a scalar + :param k: a scalar + :param m: a scalar to use for modulo + :return: a vector with 3 values + """ + i = np.int64(i) + j = np.int64(j) + k = np.int64(k) + m = np.int64(m) + o0 = (((a[0][0] * i) % m) + ((a[0][1] * j) % m) + ((a[0][2] * k) % m)) % m + o1 = (((a[1][0] * i) % m) + ((a[1][1] * j) % m) + ((a[1][2] * k) % m)) % m + o2 = (((a[2][0] * i) % m) + ((a[2][1] * j) % m) + ((a[2][2] * k) % m)) % m + return np.int32(o0), np.int32(o1), np.int32(o2) + + +def _mrg_new_stream(seed=12345): + """ + Create and return a new random state (vector of 6 32-bits integers). + + Seed can be an integer or a vector of 6 integers. + :type seed: int | tuple + :return: a tuple with 6 32-bits integers. + """ + if isinstance(seed, int): + if seed == 0: + raise Exception('seed should not be 0') + elif seed >= M2: + raise Exception('seed should be less than M2') + rstate = (np.int32(seed), np.int32(seed), np.int32(seed), + np.int32(seed), np.int32(seed), np.int32(seed)) + elif len(seed) == 6: + if seed[0] == 0 and seed[1] == 0 and seed[2] == 0: + raise Exception('First 3 values of seed should not be all 0') + if seed[3] == 0 and seed[4] == 0 and seed[5] == 0: + raise Exception('Last 3 values of seed should not be all 0') + if seed[0] >= M1 or seed[1] >= M1 or seed[2] >= M1: + raise Exception('First 3 values of seed should be less than M1') + if seed[3] >= M2 or seed[4] >= M2 or seed[5] >= M2: + raise Exception('Last 3 values of seed should be less than M2') + rstate = (np.int32(seed[0]), np.int32(seed[1]), np.int32(seed[2]), + np.int32(seed[3]), np.int32(seed[4]), np.int32(seed[5])) + else: + raise Exception("seed should be 1 integer or 6 integers") + return rstate + + +def _mrg_next_value(rstate, norm, mask, offset, dtype): + """ + Generate next value from given MRG stream, norm, mask, offset and dtype. + + Return new value and updated MRG stream. + """ + x11 = rstate[0] + x12 = rstate[1] + x13 = rstate[2] + x21 = rstate[3] + x22 = rstate[4] + x23 = rstate[5] + + y1 = (((x12 & MASK12) << i22) + (x12 >> i9) + + ((x13 & MASK13) << i7) + (x13 >> i24)) + + if y1 < 0 or y1 >= M1: + y1 = y1 - M1 + y1 = y1 + x13 + if y1 < 0 or y1 >= M1: + y1 = y1 - M1 + + x13 = x12 + x12 = x11 + x11 = y1 + + y1 = ((x21 & MASK2) << i15) + (MULT2 * (x21 >> i16)) + if y1 < 0 or y1 >= M2: + y1 = y1 - M2 + y2 = ((x23 & MASK2) << i15) + (MULT2 * (x23 >> i16)) + + if y2 < 0 or y2 >= M2: + y2 = y2 - M2 + y2 = y2 + x23 + if y2 < 0 or y2 >= M2: + y2 = y2 - M2 + y2 = y2 + y1 + if y2 < 0 or y2 >= M2: + y2 = y2 - M2 + + x23 = x22 + x22 = x21 + x21 = y2 + + new_rstate = x11, x12, x13, x21, x22, x23 + + if x11 <= x21: + value = SUPP_TYPES[dtype](((x11 - x21 + M1) & mask) + offset) * norm + else: + value = SUPP_TYPES[dtype](((x11 - x21) & mask) + offset) * norm + return value, new_rstate + + +# --------------- +# RNG public API. +# --------------- + + +@dataclass +class MyiaRandomState: + """ + Data class to represent a random state specialized for a d-type. + + Supported dtypes: float16, float32, float64 + """ + + rstate: tuple + mask: object + offset: object + norm: object + dtype: object + + +def myia_random_state(dtype, seed=12345): + """ + Generate a new random state for given dtype and seed. + + :param dtype: a string: either 'float16', ´float32' or 'float64'. + :param seed: an integer or a tuple of 6 integers. + :return: a MyiaRandomState object. + """ + rstate = _mrg_new_stream(seed) + if dtype == 'float16': + mask = 0x7fff + offset = 1 + norm = np.float16(3.0458e-05) + dtype = FLOAT16 + elif dtype == 'float32': + mask = 0xffffffff + offset = 0 + norm = np.float32(4.6566126e-10) + dtype = FLOAT32 + elif dtype == 'float64': + mask = 0xffffffff + offset = 0 + norm = np.float64(4.656612873077392578125e-10) # 1./2^31 + dtype = FLOAT64 + else: + raise Exception('Unsupported type for random state') + return MyiaRandomState(rstate, mask, offset, norm, dtype) + + +def myia_increment_state(state: MyiaRandomState): + """ + Increment a random state and return new updated random state. + + Useful to create a new random state from a previous existing one. + Not useful in the random generation process. + :param state: an existing MyiaRandomState object. + :return: a MyiaRandomState object. + """ + rs = state.rstate + new_rstate = (_mat_vec_mod_m(a_1_p_134, rs[0], rs[1], rs[2], M1) + + _mat_vec_mod_m(a_2_p_134, rs[3], rs[4], rs[5], M2)) + return MyiaRandomState( + new_rstate, state.mask, state.offset, state.norm, state.dtype) + + +def myia_next_value(state: MyiaRandomState): + """ + Generate a scalar using given state. + + :param state: a MyiaRandomState object. + :return: a tuple: (new scalar value, updated state). + """ + sample, new_rstate = _mrg_next_value( + state.rstate, state.norm, state.mask, state.offset, state.dtype) + return sample, MyiaRandomState( + new_rstate, state.mask, state.offset, state.norm, state.dtype) + + +# -------------------------- +# MRG public high-level API. +# -------------------------- + + +def myia_uniform(state: MyiaRandomState, low=0.0, high=1.0): + """ + Generate a uniform scalar between low and high. + + :param state: a MyiaRandomState object. + :param low: lower bound for uniform value. + :param high: upper bound for uniform value. + :return: a tuple: (new uniform scalar value, updated state). + """ + u, next_state = myia_next_value(state) + r = u * (high - low) + low + return r, next_state diff --git a/tests/test_rng.py b/tests/test_rng.py new file mode 100644 index 000000000..526157be9 --- /dev/null +++ b/tests/test_rng.py @@ -0,0 +1,89 @@ +""" Test random number generator. +Notes: + - Compilation is slow (about 1 minute). + - Got several other compilation issues related to dtype handling, as dtype + is currently a string. We may need to set a better support for strings + in Myia. + - Exceptions are not correctly handled with @myia decorator. + Many "illegal" errors are raised, including illegal primitives, + illegal registered types, and other things. I can fix that, + but ultimately we will still need to correctly support strings + (so that exceptions could be correctly printed in backend). +""" + +import numpy as np + +from myia import myia +from myia.rng_mrg import ( + MyiaRandomState, + myia_increment_state, + myia_random_state, + myia_uniform, +) + + +def pyth_random(): + state = myia_random_state('float64') + v1, state = myia_uniform(state) + v2, state = myia_uniform(state) + v3, state = myia_uniform(state) + v4, state = myia_uniform(state) + v5, state = myia_uniform(state) + v6, state = myia_uniform(state) + v7, state = myia_uniform(state) + v8, state = myia_uniform(state) + v9, state = myia_uniform(state) + v10, _ = myia_uniform(state) + return v1, v2, v3, v4, v5, v6, v7, v8, v9, v10 + + +@myia +def myia_random(): + return pyth_random() + + +@myia +def run_myia_random_state(): + return myia_random_state('float32') + + +@myia +def run_myia_increment_state(s): + return myia_increment_state(s) + + +def test_myia_random_state(): + state = run_myia_random_state() + assert isinstance(state, MyiaRandomState) + assert isinstance(state.rstate, tuple) + assert state.rstate == (np.int32(12345),) * 6 + + +def test_myia_increment_state(): + state_1 = run_myia_random_state() + state_2 = run_myia_increment_state(state_1) + state_3 = run_myia_increment_state(state_2) + state_4 = run_myia_increment_state(state_3) + assert state_2.rstate == tuple(np.int32(val) for val in ( + 336690377, 597094797, 1245771585, 85196284, 523477687, 2094976052)) + assert state_3.rstate == tuple(np.int32(val) for val in ( + 502033783, 1322587635, 1964121530, 1949818481, 1607232546, 1462898381)) + assert state_4.rstate == tuple(np.int32(val) for val in ( + 739421137, 1475938232, 730262207, 1630192198, 324551134, 795289868)) + + +def test_myia_random_generation(): + # Hide runtime warnings about overflow in integer operations. + err_orig = np.seterr(all='ignore') + + pyth_res = pyth_random() + myia_res = myia_random() + assert pyth_res == myia_res + assert pyth_res == tuple(np.float64(val) for val in ( + 0.7353244530968368, 0.6142074400559068, 0.11007806099951267, + 0.6487741703167558, 0.36619443260133266, 0.10882294131442904, + 0.5330547927878797, 0.9783797566778958, 0.9151237849146128, + 0.8509745532646775)) + + # Get numpy back to its original warning config. + np.seterr(**err_orig)