Merge pull request #32 from bashtage/mrg32k3a-perf

Mrg32k3a perf

Author: bashtage

randomstate/interface.pyx

@@ -499,7 +499,7 @@ cdef class RandomState:
 
         if state['name'] != rng_name:
             raise ValueError('Not a ' + rng_name + ' RNG state')
-        print(state['state'])
+
         _set_state(&self.rng_state, state['state'])
         self.rng_state.has_gauss = state['gauss']['has_gauss']
         self.rng_state.gauss = state['gauss']['gauss']

randomstate/shims/mrg32k3a/mrg32k3a-shim.c

@@ -32,11 +32,12 @@ void set_seed(aug_state* state, uint64_t val)
 
 void init_state(aug_state* state, int64_t vals[6])
 {
-    state->rng->s10 =  vals[0];
-    state->rng->s11 =  vals[1];
-    state->rng->s12 =  vals[2];
-    state->rng->s20 =  vals[3];
-    state->rng->s21 =  vals[4];
-    state->rng->s22 =  vals[5];
+    state->rng->s1[0] =  vals[0];
+    state->rng->s1[1] =  vals[1];
+    state->rng->s1[2] =  vals[2];
+    state->rng->s2[0] =  vals[3];
+    state->rng->s2[1] =  vals[4];
+    state->rng->s2[2] =  vals[5];
+    state->rng->loc = 2;
 }
 

randomstate/shims/mrg32k3a/mrg32k3a.pxi

@@ -4,12 +4,9 @@ DEF RS_RNG_JUMPABLE = 1
 cdef extern from "distributions.h":
 
     cdef struct s_mrg32k3a_state:
-        int64_t s10
-        int64_t s11
-        int64_t s12
-        int64_t s20
-        int64_t s21
-        int64_t s22
+        int64_t s1[3]
+        int64_t s2[3]
+        int loc
 
     ctypedef s_mrg32k3a_state mrg32k3a_state
 
@@ -31,16 +28,18 @@ ctypedef mrg32k3a_state rng_t
 ctypedef uint64_t rng_state_t
 
 cdef object _get_state(aug_state state):
-    return (state.rng.s10, state.rng.s11, state.rng.s12,
-            state.rng.s20, state.rng.s21, state.rng.s22)
+    return (state.rng.s1[0], state.rng.s1[1], state.rng.s1[2],
+            state.rng.s2[0], state.rng.s2[1], state.rng.s2[2],
+            state.rng.loc)
 
 cdef object _set_state(aug_state *state, object state_info):
-    state.rng.s10 = state_info[0]
-    state.rng.s11 = state_info[1]
-    state.rng.s12 = state_info[2]
-    state.rng.s20 = state_info[3]
-    state.rng.s21 = state_info[4]
-    state.rng.s22 = state_info[5]
+    state.rng.s1[0] = state_info[0]
+    state.rng.s1[1] = state_info[1]
+    state.rng.s1[2] = state_info[2]
+    state.rng.s2[0] = state_info[3]
+    state.rng.s2[1] = state_info[4]
+    state.rng.s2[2] = state_info[5]
+    state.rng.loc = state_info[6]
 
 cdef object matrix_power_127(x, m):
     n = x.shape[0]
@@ -68,21 +67,39 @@ A2_127 = matrix_power_127(A2p, m2)
 
 cdef void jump_state(aug_state* state):
     # vectors s1 and s2
-    s1 = np.array([state.rng.s10,state.rng.s11,state.rng.s12], dtype=np.uint64)
-    s2 = np.array([state.rng.s20,state.rng.s21,state.rng.s22], dtype=np.uint64)
+    loc = state.rng.loc
+
+    if loc == 0:
+        loc_m1 = 2
+        loc_m2 = 1
+    elif loc == 1:
+        loc_m1 = 0
+        loc_m2 = 2
+    else:
+        loc_m1 = 1
+        loc_m2 = 0
+
+    s1 = np.array([state.rng.s1[loc_m2],
+                   state.rng.s1[loc_m1],
+                   state.rng.s1[loc]], dtype=np.uint64)
+    s2 = np.array([state.rng.s2[loc_m2],
+                   state.rng.s2[loc_m1],
+                   state.rng.s2[loc]], dtype=np.uint64)
 
     # Advance the state
     s1 = np.mod(A1_127.dot(s1), m1)
     s2 = np.mod(A1_127.dot(s2), m2)
 
     # Restore state
-    state.rng.s10 = s1[0]
-    state.rng.s11 = s1[1]
-    state.rng.s12 = s1[2]
+    state.rng.s1[0] = s1[0]
+    state.rng.s1[1] = s1[1]
+    state.rng.s1[2] = s1[2]
 
-    state.rng.s20 = s2[0]
-    state.rng.s21 = s2[1]
-    state.rng.s22 = s2[2]
+    state.rng.s2[0] = s2[0]
+    state.rng.s2[1] = s2[1]
+    state.rng.s2[2] = s2[2]
+
+    state.rng.loc = 2
 
 DEF CLASS_DOCSTRING = """
 RandomState(seed=None)

randomstate/src/mrg32k3a/mrg32k3a.c

@@ -38,10 +38,11 @@ void mrg32k3a_seed(mrg32k3a_state* state, uint64_t seed)
         seeds[i] = draw;
     }
 
-    state->s10 = seeds[0];
-    state->s11 = seeds[1];
-    state->s12 = seeds[2];
-    state->s20 = seeds[3];
-    state->s21 = seeds[4];
-    state->s22 = seeds[5];
+    state->s1[0] = seeds[0];
+    state->s1[1] = seeds[1];
+    state->s1[2] = seeds[2];
+    state->s2[0] = seeds[3];
+    state->s2[1] = seeds[4];
+    state->s2[2] = seeds[5];
+    state->loc = 2;
 }

randomstate/src/mrg32k3a/mrg32k3a.h

@@ -14,43 +14,41 @@
 
 typedef struct s_mrg32k3a_state
 {
-    int64_t s10;
-    int64_t s11;
-    int64_t s12;
-    int64_t s20;
-    int64_t s21;
-    int64_t s22;
+    int64_t s1[3];
+    int64_t s2[3];
+    int loc;
 } mrg32k3a_state;
 
 inline uint32_t mrg32k3a_random(mrg32k3a_state* state)
 {
-    int64_t k;
     int64_t p1, p2;
     /* Component 1 */
-    p1 = a12 * state->s11 - a13n * state->s10;
-    k = p1 / m1;
-    p1 -= k * m1;
-    if (p1 < 0)
-        p1 += m1;
-    state->s10 = state->s11;
-    state->s11 = state->s12;
-    state->s12 = p1;
+    switch (state->loc) {
+        case 0:
+            p1 = a12 * state->s1[2] - a13n * state->s1[1];
+            p2 = a21 * state->s2[0] - a23n * state->s2[1];
+            state->loc = 1;
+            break;
+        case 1:
+            p1 = a12 * state->s1[0] - a13n * state->s1[2];
+            p2 = a21 * state->s2[1] - a23n * state->s2[2];
+            state->loc = 2;
+            break;
+        case 2:
+            p1 = a12 * state->s1[1] - a13n * state->s1[0];
+            p2 = a21 * state->s2[2] - a23n * state->s2[0];
+            state->loc = 0;
+            break;
+    }
 
+    p1 -= (p1 >= 0) ?  (p1 / m1) * m1 : (p1 / m1) * m1 - m1;
+    state->s1[state->loc] = p1;
     /* Component 2 */
-    p2 = a21 * state->s22 - a23n * state->s20;
-    k = p2 / m2;
-    p2 -= k * m2;
-    if (p2 < 0)
-        p2 += m2;
-    state->s20 = state->s21;
-    state->s21 = state->s22;
-    state->s22 = p2;
+    p2 -= (p2 >= 0) ? (p2 / m2) * m2 : (p2 / m2) * m2 - m2;
+    state->s2[state->loc] = p2;
 
     /* Combination */
-    if (p1 <= p2)
-        return (uint32_t)(p1 - p2 + m1);
-    else
-        return (uint32_t)(p1 - p2);
+    return (uint32_t)((p1 <= p2) ? (p1 - p2 + m1) : (p1 - p2));
 }
 
 void mrg32k3a_seed(mrg32k3a_state* state, uint64_t seed);