Merge pull request #23 from bashtage/add-aligned-malloc

ENH: Add aligned malloc

Author: bashtage

randomstate/aligned_malloc.c

@@ -0,0 +1,9 @@
+#include "aligned_malloc.h"
+
+static NPY_INLINE void *PyArray_realloc_aligned(void *p, size_t n);
+
+static NPY_INLINE void *PyArray_malloc_aligned(size_t n);
+
+static NPY_INLINE void *PyArray_calloc_aligned(size_t n, size_t s);
+
+static NPY_INLINE void PyArray_free_aligned(void *p);

randomstate/aligned_malloc.h

@@ -0,0 +1,47 @@
+#include "Python.h"
+#include "numpy/npy_common.h"
+
+#define NPY_MEMALIGN 16   /* 16 for SSE2, 32 for AVX, 64 for Xeon Phi */
+
+static NPY_INLINE
+void *PyArray_realloc_aligned(void *p, size_t n)
+{
+    void *p1, **p2, *base;
+    size_t old_offs, offs = NPY_MEMALIGN - 1 + sizeof(void*);
+    if (NPY_UNLIKELY(p != NULL)) {
+        base = *(((void**)p)-1);
+        if (NPY_UNLIKELY((p1 = PyMem_Realloc(base,n+offs)) == NULL)) return NULL;
+        if (NPY_LIKELY(p1 == base)) return p;
+        p2 = (void**)(((Py_uintptr_t)(p1)+offs) & ~(NPY_MEMALIGN-1));
+        old_offs = (size_t)((Py_uintptr_t)p - (Py_uintptr_t)base);
+        /* TODO: This isn't right, removed void* to aloow msvc to do pointer math */
+        memmove((void*)p2,((char*)p1)+old_offs,n);
+    } else {
+        if (NPY_UNLIKELY((p1 = PyMem_Malloc(n + offs)) == NULL)) return NULL;
+        p2 = (void**)(((Py_uintptr_t)(p1)+offs) & ~(NPY_MEMALIGN-1));
+    }
+    *(p2-1) = p1;
+    return (void*)p2;
+}
+
+static NPY_INLINE
+void *PyArray_malloc_aligned(size_t n)
+{
+    return PyArray_realloc_aligned(NULL, n);
+}
+
+static NPY_INLINE
+void *PyArray_calloc_aligned(size_t n, size_t s)
+{
+    void *p;
+    if (NPY_UNLIKELY((p = PyArray_realloc_aligned(NULL,n*s)) == NULL)) return NULL;
+    memset(p, 0, n*s);
+    return p;
+}
+
+static NPY_INLINE
+void PyArray_free_aligned(void *p)
+{
+    void *base = *(((void**)p)-1);
+    PyMem_Free(base);
+}

randomstate/aligned_malloc.pxi

@@ -0,0 +1,5 @@
+cdef extern from "aligned_malloc.h":
+    cdef void *PyArray_realloc_aligned(void *p, size_t n);
+    cdef void *PyArray_malloc_aligned(size_t n);
+    cdef void *PyArray_calloc_aligned(size_t n, size_t s);
+    cdef void PyArray_free_aligned(void *p);

randomstate/interface.pyx

@@ -112,9 +112,9 @@ cdef extern from "distributions.h":
     cdef void random_gauss_fill(aug_state* state, int count, double *out) nogil
     cdef void random_gauss_zig_julia_fill(aug_state* state, int count, double *out) nogil
 
-
 include "array_utilities.pxi"
 include "bounded_integers.pxi"
+include "aligned_malloc.pxi"
 
 cdef double kahan_sum(double *darr, np.npy_intp n):
     cdef double c, y, t, sum
@@ -141,38 +141,25 @@ cdef class RandomState:
 
     IF RNG_SEED==1:
         def __init__(self, seed=None):
-            IF RNG_MOD_NAME == 'dsfmt':
-                cdef int8_t *iptr
-                cdef int8_t offset = 0
-                cdef intptr_t alignment = 0
-                self.rng_loc = PyMem_Malloc(sizeof(rng_t))
-                self.rng_state.rng = <rng_t *>self.rng_loc
-                alignment  = <intptr_t>(&(self.rng_state.rng.status[0].u32[0]))
-                if (alignment % 16) != 0:
-                    iptr = <int8_t *>self.rng_state.rng
-                    offset = 16 - (alignment % 16)
-                    if offset < 0:
-                        offset += 16
-                    self.rng_state.rng = <rng_t *>(iptr + offset)
-            ELSE:
-                self.rng_loc = PyMem_Malloc(sizeof(rng_t))
-                self.rng_state.rng = <rng_t *>self.rng_loc
-
+            self.rng_state.rng = <rng_t *>PyArray_malloc_aligned(sizeof(rng_t))
             self.rng_state.binomial = &self.binomial_info
-            self._reset_state_variables()
+            IF RNG_MOD_NAME == 'dsfmt':
+                self.rng_state.buffered_uniforms = <double *>PyArray_malloc_aligned(2 * DSFMT_N * sizeof(double))
             self.lock = Lock()
+            self._reset_state_variables()
             self.seed(seed)
-
     ELSE:
         def __init__(self, seed=None, inc=None):
-            self.rng_state.rng = <rng_t *>PyMem_Malloc(sizeof(rng_t)) # &self.rng
+            self.rng_state.rng = <rng_t *>PyArray_malloc_aligned(sizeof(rng_t))
             self.rng_state.binomial = &self.binomial_info
-            self._reset_state_variables()
             self.lock = Lock()
+            self._reset_state_variables()
             self.seed(seed, inc)
 
     def __dealloc__(self):
-        PyMem_Free(self.rng_loc)
+        PyArray_free_aligned(self.rng_state.rng)
+        IF RNG_MOD_NAME == 'dsfmt':
+            PyArray_free_aligned(self.rng_state.buffered_uniforms)
 
     # Pickling support:
     def __getstate__(self):
@@ -496,7 +483,7 @@ cdef class RandomState:
             if isinstance(state, tuple):
                 if state[0] != 'MT19937':
                     raise ValueError('Not a ' + rng_name + ' RNG state')
-                _set_state(self.rng_state, (state[1], state[2]))
+                _set_state(&self.rng_state, (state[1], state[2]))
                 self.rng_state.has_gauss = state[3]
                 self.rng_state.gauss = state[4]
                 self.rng_state.has_uint32 = 0
@@ -505,7 +492,8 @@ cdef class RandomState:
 
         if state['name'] != rng_name:
             raise ValueError('Not a ' + rng_name + ' RNG state')
-        _set_state(self.rng_state, state['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']
         self.rng_state.has_uint32 = state['uint32']['has_uint32']

randomstate/shims/dSFMT/dSFMT-shim.c

@@ -6,13 +6,25 @@ extern inline uint64_t random_uint64(aug_state* state);
 
 extern inline double random_double(aug_state* state);
 
+void reset_buffer(aug_state* state)
+{
+    int i = 0;
+    for (i = 0; i < (2 * DSFMT_N); i++)
+    {
+        state->buffered_uniforms[i] = 0.0;
+    }
+    state->buffer_loc = 2 * DSFMT_N;
+}
+
 extern void set_seed_by_array(aug_state* state, uint32_t init_key[], int key_length)
 {
+    reset_buffer(state);
     dsfmt_init_by_array(state->rng, init_key, key_length);
 }
 
 void set_seed(aug_state* state, uint32_t seed)
 {
+    reset_buffer(state);
     dsfmt_init_gen_rand(state->rng, seed);
 }
 

randomstate/shims/dSFMT/dSFMT-shim.h

@@ -19,31 +19,48 @@ typedef struct s_aug_state {
     double gauss;
     uint32_t uinteger;
     uint64_t zig_random_int;
+
+    double *buffered_uniforms;
+    int buffer_loc;
 } aug_state;
 
+static inline double random_double_from_buffer(aug_state *state)
+{
+    double out;
+    if (state->buffer_loc >= (2 * DSFMT_N))
+    {
+        state->buffer_loc = 0;
+        dsfmt_fill_array_close1_open2(state->rng, state->buffered_uniforms, 2 * DSFMT_N);
+    }
+    out = state->buffered_uniforms[state->buffer_loc];
+    state->buffer_loc++;
+    return  out;
+}
+
 static inline uint32_t random_uint32(aug_state* state)
 {
-    double d = dsfmt_genrand_close1_open2(state->rng);
+    double d = random_double_from_buffer(state);//dsfmt_genrand_close1_open2(state->rng);
     uint64_t *out = (uint64_t *)&d;
     return (uint32_t)(*out & 0xffffffff);
 }
 
 static inline uint64_t random_uint64(aug_state* state)
 {
-    double d = dsfmt_genrand_close1_open2(state->rng);
+    double d = random_double_from_buffer(state);//dsfmt_genrand_close1_open2(state->rng);
     uint64_t out;
     uint64_t *tmp;
     tmp = (uint64_t *)&d;
     out = *tmp << 32;
-    d = dsfmt_genrand_close1_open2(state->rng);
+    d = random_double_from_buffer(state);//dsfmt_genrand_close1_open2(state->rng);
     tmp =  (uint64_t *)&d;
     out |= *tmp & 0xffffffff;
     return out;
 }
 
 static inline double random_double(aug_state* state)
 {
-    return dsfmt_genrand_close1_open2(state->rng) - 1.0;
+    return random_double_from_buffer(state) - 1.0;
+    // return dsfmt_genrand_close1_open2(state->rng) - 1.0;
 }
 
 extern void entropy_init(aug_state* state);

randomstate/shims/dSFMT/dSFMT.pxi

@@ -34,6 +34,9 @@ cdef extern from "distributions.h":
         uint64_t zig_random_int
         uint32_t uinteger
 
+        double *buffered_uniforms
+        int buffer_loc
+
     ctypedef s_aug_state aug_state
 
     cdef void set_seed(aug_state* state, uint32_t seed)
@@ -43,25 +46,38 @@ cdef extern from "distributions.h":
 ctypedef dsfmt_t rng_t
 
 cdef object _get_state(aug_state state):
-    cdef uint32_t [:] key = np.zeros(4 * DSFMT_N_PLUS_1, dtype=np.uint32)
+    cdef uint32_t [::1] key = np.zeros(4 * DSFMT_N_PLUS_1, dtype=np.uint32)
+    cdef double [::1] buf = np.zeros(2 * DSFMT_N, dtype=np.double)
     cdef Py_ssize_t i, j, key_loc = 0
     cdef w128_t state_val
     for i in range(DSFMT_N_PLUS_1):
         state_val = state.rng.status[i]
         for j in range(4):
             key[key_loc] = state_val.u32[j]
             key_loc += 1
-    return (np.asarray(key), state.rng.idx)
+    for i in range(2 * DSFMT_N):
+        buf[i] = state.buffered_uniforms[i]
+
+    return (np.asarray(key), state.rng.idx,
+            np.asarray(buf), state.buffer_loc)
 
-cdef object _set_state(aug_state state, object state_info):
-    cdef uint32_t [:] key = state_info[0]
+cdef object _set_state(aug_state *state, object state_info):
     cdef Py_ssize_t i, j, key_loc = 0
+    cdef uint32_t [::1] key = state_info[0]
+    state.rng.idx = state_info[1]
+
+
     for i in range(DSFMT_N_PLUS_1):
         for j in range(4):
             state.rng.status[i].u32[j] = key[key_loc]
             key_loc += 1
 
-    state.rng.idx = state_info[1]
+    state.buffer_loc = <int>state_info[3]
+    for i in range(2 * DSFMT_N):
+        state.buffered_uniforms[i] = state_info[2][i]
+
+
+
 
 DEF CLASS_DOCSTRING = """
 RandomState(seed=None)

randomstate/shims/mlfg-1279-861/mlfg-1279-861.pxi

@@ -39,7 +39,7 @@ cdef object _get_state(aug_state state):
         key[i] = state.rng.lags[i]
     return (np.asanyarray(key), state.rng.pos, state.rng.lag_pos)
 
-cdef object _set_state(aug_state state, object state_info):
+cdef object _set_state(aug_state *state, object state_info):
     cdef uint64_t [:] key = state_info[0]
     cdef Py_ssize_t i
     for i in range(MLFG_STATE_LEN):

randomstate/shims/mrg32k3a/mrg32k3a.pxi

@@ -38,7 +38,7 @@ 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)
 
-cdef object _set_state(aug_state state, object state_info):
+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]

randomstate/shims/pcg-32/pcg-32.pxi

@@ -34,7 +34,7 @@ ctypedef pcg32_random_t rng_t
 cdef object _get_state(aug_state state):
     return (state.rng.state, state.rng.inc)
 
-cdef object _set_state(aug_state state, object state_info):
+cdef object _set_state(aug_state *state, object state_info):
     state.rng.state = state_info[0]
     state.rng.inc = state_info[1]