Merge pull request #4606 from juliantaylor/small-memcache

ENH: add small memory chunk cache

10 files changed, 279 insertions, 133 deletions

diff --git a/numpy/core/bscript b/numpy/core/bscript
index 3f1ec149b..872d4d2ba 100644
--- a/numpy/core/bscript
+++ b/numpy/core/bscript
@@ -444,6 +444,7 @@ def pre_build(context):
         bld(target="multiarray_templates", source=multiarray_templates)
         if ENABLE_SEPARATE_COMPILATION:
             sources = [pjoin('src', 'multiarray', 'arrayobject.c'),
+                pjoin('src', 'multiarray', 'alloc.c'),
                 pjoin('src', 'multiarray', 'arraytypes.c.src'),
                 pjoin('src', 'multiarray', 'array_assign.c'),
                 pjoin('src', 'multiarray', 'array_assign_array.c'),
diff --git a/numpy/core/code_generators/genapi.py b/numpy/core/code_generators/genapi.py
index c73209d04..5ab60a37c 100644
--- a/numpy/core/code_generators/genapi.py
+++ b/numpy/core/code_generators/genapi.py
@@ -23,7 +23,8 @@ from os.path import join
 __docformat__ = 'restructuredtext'
 
 # The files under src/ that are scanned for API functions
-API_FILES = [join('multiarray', 'array_assign_array.c'),
+API_FILES = [join('multiarray', 'alloc.c'),
+             join('multiarray', 'array_assign_array.c'),
              join('multiarray', 'array_assign_scalar.c'),
              join('multiarray', 'arrayobject.c'),
              join('multiarray', 'arraytypes.c.src'),
diff --git a/numpy/core/setup.py b/numpy/core/setup.py
index c28de7c11..3cde96d77 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -792,6 +792,7 @@ def configuration(parent_package='',top_path=None):
             ] + npysort_sources + npymath_sources
 
     multiarray_src = [
+            join('src', 'multiarray', 'alloc.c'),
             join('src', 'multiarray', 'arrayobject.c'),
             join('src', 'multiarray', 'arraytypes.c.src'),
             join('src', 'multiarray', 'array_assign.c'),
diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c
new file mode 100644
index 000000000..06c9ff296
--- /dev/null
+++ b/numpy/core/src/multiarray/alloc.c
@@ -0,0 +1,241 @@
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include "structmember.h"
+
+#define NPY_NO_DEPRECATED_API NPY_API_VERSION
+#define _MULTIARRAYMODULE
+#include <numpy/ndarraytypes.h>
+#include "numpy/arrayobject.h"
+#include <numpy/npy_common.h>
+#include "npy_config.h"
+
+#include <assert.h>
+
+#define NBUCKETS 1024 /* number of buckets for data*/
+#define NBUCKETS_DIM 16 /* number of buckets for dimensions/strides */
+#define NCACHE 8 /* 1 + number of cache entries per bucket */
+static void * datacache[NBUCKETS][NCACHE];
+static void * dimcache[NBUCKETS_DIM][NCACHE];
+
+/*
+ * very simplistic small memory block cache to avoid more expensive libc
+ * allocations
+ * base function for data cache with 1 byte buckets and dimension cache with
+ * sizeof(npy_intp) byte buckets
+ */
+static NPY_INLINE void *
+_npy_alloc_cache(npy_uintp nelem, npy_uintp esz, npy_uint msz,
+                 void * (*cache)[NCACHE], void * (*alloc)(size_t))
+{
+    assert((esz == 1 && cache == datacache) ||
+           (esz == sizeof(npy_intp) && cache == dimcache));
+    if (nelem < msz) {
+        /* first entry is used as fill counter */
+        npy_uintp * idx = (npy_uintp *)&(cache[nelem][0]);
+        if (*idx > 0) {
+            return cache[nelem][(*idx)--];
+        }
+    }
+    return alloc(nelem * esz);
+}
+
+/*
+ * return pointer p to cache, nelem is number of elements of the cache bucket
+ * size (1 or sizeof(npy_intp)) of the block pointed too
+ */
+static NPY_INLINE void
+_npy_free_cache(void * p, npy_uintp nelem, npy_uint msz,
+                void * (*cache)[NCACHE], void (*dealloc)(void *))
+{
+    if (p != NULL && nelem < msz) {
+        /* first entry is used as fill counter */
+        npy_uintp * idx = (npy_uintp *)&(cache[nelem][0]);
+        if (*idx < NCACHE - 1) {
+            cache[nelem][++(*idx)] = p;
+            return;
+        }
+    }
+    dealloc(p);
+}
+
+
+/*
+ * array data cache, sz is number of bytes to allocate
+ */
+NPY_NO_EXPORT void *
+npy_alloc_cache(npy_uintp sz)
+{
+    return _npy_alloc_cache(sz, 1, NBUCKETS, datacache, &PyDataMem_NEW);
+}
+
+/* zero initialized data, sz is number of bytes to allocate */
+NPY_NO_EXPORT void *
+npy_alloc_cache_zero(npy_uintp sz)
+{
+    void * p;
+    NPY_BEGIN_THREADS_DEF;
+    if (sz < NBUCKETS) {
+        p = _npy_alloc_cache(sz, 1, NBUCKETS, datacache, &PyDataMem_NEW);
+        memset(p, 0, sz);
+        return p;
+    }
+    NPY_BEGIN_THREADS;
+    p = PyDataMem_NEW_ZEROED(sz, 1);
+    NPY_END_THREADS;
+    return p;
+}
+
+NPY_NO_EXPORT void
+npy_free_cache(void * p, npy_uintp sz)
+{
+    _npy_free_cache(p, sz, NBUCKETS, datacache, &PyDataMem_FREE);
+}
+
+/*
+ * dimension/stride cache, uses a different allocator and is always a multiple
+ * of npy_intp
+ */
+NPY_NO_EXPORT void *
+npy_alloc_cache_dim(npy_uintp sz)
+{
+    /* dims + strides */
+    if (NPY_UNLIKELY(sz < 2)) {
+        sz = 2;
+    }
+    return _npy_alloc_cache(sz, sizeof(npy_intp), NBUCKETS_DIM, dimcache,
+                            &PyArray_malloc);
+}
+
+NPY_NO_EXPORT void
+npy_free_cache_dim(void * p, npy_uintp sz)
+{
+    /* dims + strides */
+    if (NPY_UNLIKELY(sz < 2)) {
+        sz = 2;
+    }
+    _npy_free_cache(p, sz, NBUCKETS_DIM, dimcache,
+                    &PyArray_free);
+}
+
+
+/* malloc/free/realloc hook */
+NPY_NO_EXPORT PyDataMem_EventHookFunc *_PyDataMem_eventhook;
+NPY_NO_EXPORT void *_PyDataMem_eventhook_user_data;
+
+/*NUMPY_API
+ * Sets the allocation event hook for numpy array data.
+ * Takes a PyDataMem_EventHookFunc *, which has the signature:
+ *        void hook(void *old, void *new, size_t size, void *user_data).
+ *   Also takes a void *user_data, and void **old_data.
+ *
+ * Returns a pointer to the previous hook or NULL.  If old_data is
+ * non-NULL, the previous user_data pointer will be copied to it.
+ *
+ * If not NULL, hook will be called at the end of each PyDataMem_NEW/FREE/RENEW:
+ *   result = PyDataMem_NEW(size)        -> (*hook)(NULL, result, size, user_data)
+ *   PyDataMem_FREE(ptr)                 -> (*hook)(ptr, NULL, 0, user_data)
+ *   result = PyDataMem_RENEW(ptr, size) -> (*hook)(ptr, result, size, user_data)
+ *
+ * When the hook is called, the GIL will be held by the calling
+ * thread.  The hook should be written to be reentrant, if it performs
+ * operations that might cause new allocation events (such as the
+ * creation/descruction numpy objects, or creating/destroying Python
+ * objects which might cause a gc)
+ */
+NPY_NO_EXPORT PyDataMem_EventHookFunc *
+PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook,
+                       void *user_data, void **old_data)
+{
+    PyDataMem_EventHookFunc *temp;
+    NPY_ALLOW_C_API_DEF
+    NPY_ALLOW_C_API
+    temp = _PyDataMem_eventhook;
+    _PyDataMem_eventhook = newhook;
+    if (old_data != NULL) {
+        *old_data = _PyDataMem_eventhook_user_data;
+    }
+    _PyDataMem_eventhook_user_data = user_data;
+    NPY_DISABLE_C_API
+    return temp;
+}
+
+/*NUMPY_API
+ * Allocates memory for array data.
+ */
+NPY_NO_EXPORT void *
+PyDataMem_NEW(size_t size)
+{
+    void *result;
+
+    result = malloc(size);
+    if (_PyDataMem_eventhook != NULL) {
+        NPY_ALLOW_C_API_DEF
+        NPY_ALLOW_C_API
+        if (_PyDataMem_eventhook != NULL) {
+            (*_PyDataMem_eventhook)(NULL, result, size,
+                                    _PyDataMem_eventhook_user_data);
+        }
+        NPY_DISABLE_C_API
+    }
+    return result;
+}
+
+/*NUMPY_API
+ * Allocates zeroed memory for array data.
+ */
+NPY_NO_EXPORT void *
+PyDataMem_NEW_ZEROED(size_t size, size_t elsize)
+{
+    void *result;
+
+    result = calloc(size, elsize);
+    if (_PyDataMem_eventhook != NULL) {
+        NPY_ALLOW_C_API_DEF
+        NPY_ALLOW_C_API
+        if (_PyDataMem_eventhook != NULL) {
+            (*_PyDataMem_eventhook)(NULL, result, size * elsize,
+                                    _PyDataMem_eventhook_user_data);
+        }
+        NPY_DISABLE_C_API
+    }
+    return result;
+}
+
+/*NUMPY_API
+ * Free memory for array data.
+ */
+NPY_NO_EXPORT void
+PyDataMem_FREE(void *ptr)
+{
+    free(ptr);
+    if (_PyDataMem_eventhook != NULL) {
+        NPY_ALLOW_C_API_DEF
+        NPY_ALLOW_C_API
+        if (_PyDataMem_eventhook != NULL) {
+            (*_PyDataMem_eventhook)(ptr, NULL, 0,
+                                    _PyDataMem_eventhook_user_data);
+        }
+        NPY_DISABLE_C_API
+    }
+}
+
+/*NUMPY_API
+ * Reallocate/resize memory for array data.
+ */
+NPY_NO_EXPORT void *
+PyDataMem_RENEW(void *ptr, size_t size)
+{
+    void *result;
+
+    result = realloc(ptr, size);
+    if (_PyDataMem_eventhook != NULL) {
+        NPY_ALLOW_C_API_DEF
+        NPY_ALLOW_C_API
+        if (_PyDataMem_eventhook != NULL) {
+            (*_PyDataMem_eventhook)(ptr, result, size,
+                                    _PyDataMem_eventhook_user_data);
+        }
+        NPY_DISABLE_C_API
+    }
+    return result;
+}
diff --git a/numpy/core/src/multiarray/alloc.h b/numpy/core/src/multiarray/alloc.h
new file mode 100644
index 000000000..8f6b167d0
--- /dev/null
+++ b/numpy/core/src/multiarray/alloc.h
@@ -0,0 +1,22 @@
+#ifndef _NPY_ARRAY_ALLOC_H_
+#define _NPY_ARRAY_ALLOC_H_
+#define NPY_NO_DEPRECATED_API NPY_API_VERSION
+#define _MULTIARRAYMODULE
+#include <numpy/ndarraytypes.h>
+
+NPY_NO_EXPORT void *
+npy_alloc_cache(npy_uintp sz);
+
+NPY_NO_EXPORT void *
+npy_alloc_cache_zero(npy_uintp sz);
+
+NPY_NO_EXPORT void
+npy_free_cache(void * p, npy_uintp sd);
+
+NPY_NO_EXPORT void *
+npy_alloc_cache_dim(npy_uintp sz);
+
+NPY_NO_EXPORT void
+npy_free_cache_dim(void * p, npy_uintp sd);
+
+#endif
diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
index cd0912510..d92522a9f 100644
--- a/numpy/core/src/multiarray/arrayobject.c
+++ b/numpy/core/src/multiarray/arrayobject.c
@@ -50,6 +50,7 @@ maintainer email:  oliphant.travis@ieee.org
 #include "sequence.h"
 #include "buffer.h"
 #include "array_assign.h"
+#include "alloc.h"
 
 /*NUMPY_API
   Compute the size of an array (in number of items)
@@ -419,10 +420,11 @@ array_dealloc(PyArrayObject *self)
              * self already...
              */
         }
-        PyDataMem_FREE(fa->data);
+        npy_free_cache(fa->data, PyArray_NBYTES(self));
     }
 
-    PyDimMem_FREE(fa->dimensions);
+    /* must match allocation in PyArray_NewFromDescr */
+    npy_free_cache_dim(fa->dimensions, 2 * fa->nd);
     Py_DECREF(fa->descr);
     Py_TYPE(self)->tp_free((PyObject *)self);
 }
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index 5f8fbc4dc..a6837169b 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -26,6 +26,7 @@
 #include "array_assign.h"
 #include "mapping.h" /* for array_item_asarray */
 #include "scalarmathmodule.h" /* for npy_mul_with_overflow_intp */
+#include "alloc.h"
 #include <assert.h>
 
 /*
@@ -886,15 +887,12 @@ PyArray_NewFromDescr_int(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
     int i;
     size_t sd;
     npy_intp size;
-    assert(dims != NULL || (nd == 0));
 
     if (descr->subarray) {
         PyObject *ret;
         npy_intp newdims[2*NPY_MAXDIMS];
         npy_intp *newstrides = NULL;
-        if (nd > 0) {
-            memcpy(newdims, dims, nd * sizeof(npy_intp));
-        }
+        memcpy(newdims, dims, nd*sizeof(npy_intp));
         if (strides) {
             newstrides = newdims + NPY_MAXDIMS;
             memcpy(newstrides, strides, nd*sizeof(npy_intp));
@@ -994,7 +992,7 @@ PyArray_NewFromDescr_int(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
     fa->weakreflist = (PyObject *)NULL;
 
     if (nd > 0) {
-        fa->dimensions = PyDimMem_NEW(3*nd);
+        fa->dimensions = npy_alloc_cache_dim(2 * nd);
         if (fa->dimensions == NULL) {
             PyErr_NoMemory();
             goto fail;
@@ -1034,10 +1032,10 @@ PyArray_NewFromDescr_int(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
          * which could also be sub-fields of a VOID array
          */
         if (zeroed || PyDataType_FLAGCHK(descr, NPY_NEEDS_INIT)) {
-            data = PyDataMem_NEW_ZEROED(sd, 1);
+            data = npy_alloc_cache_zero(sd);
         }
         else {
-            data = PyDataMem_NEW(sd);
+            data = npy_alloc_cache(sd);
         }
         if (data == NULL) {
             PyErr_NoMemory();
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index a05e87e84..3f673927b 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -3604,128 +3604,6 @@ test_interrupt(PyObject *NPY_UNUSED(self), PyObject *args)
     return PyInt_FromLong(a);
 }
 
-/* malloc/free/realloc hook */
-NPY_NO_EXPORT PyDataMem_EventHookFunc *_PyDataMem_eventhook;
-NPY_NO_EXPORT void *_PyDataMem_eventhook_user_data;
-
-/*NUMPY_API
- * Sets the allocation event hook for numpy array data.
- * Takes a PyDataMem_EventHookFunc *, which has the signature:
- *        void hook(void *old, void *new, size_t size, void *user_data).
- *   Also takes a void *user_data, and void **old_data.
- *
- * Returns a pointer to the previous hook or NULL.  If old_data is
- * non-NULL, the previous user_data pointer will be copied to it.
- *
- * If not NULL, hook will be called at the end of each PyDataMem_NEW/FREE/RENEW:
- *   result = PyDataMem_NEW(size)        -> (*hook)(NULL, result, size, user_data)
- *   PyDataMem_FREE(ptr)                 -> (*hook)(ptr, NULL, 0, user_data)
- *   result = PyDataMem_RENEW(ptr, size) -> (*hook)(ptr, result, size, user_data)
- *
- * When the hook is called, the GIL will be held by the calling
- * thread.  The hook should be written to be reentrant, if it performs
- * operations that might cause new allocation events (such as the
- * creation/descruction numpy objects, or creating/destroying Python
- * objects which might cause a gc)
- */
-NPY_NO_EXPORT PyDataMem_EventHookFunc *
-PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook,
-                       void *user_data, void **old_data)
-{
-    PyDataMem_EventHookFunc *temp;
-    NPY_ALLOW_C_API_DEF
-    NPY_ALLOW_C_API
-    temp = _PyDataMem_eventhook;
-    _PyDataMem_eventhook = newhook;
-    if (old_data != NULL) {
-        *old_data = _PyDataMem_eventhook_user_data;
-    }
-    _PyDataMem_eventhook_user_data = user_data;
-    NPY_DISABLE_C_API
-    return temp;
-}
-
-/*NUMPY_API
- * Allocates memory for array data.
- */
-NPY_NO_EXPORT void *
-PyDataMem_NEW(size_t size)
-{
-    void *result;
-
-    result = malloc(size);
-    if (_PyDataMem_eventhook != NULL) {
-        NPY_ALLOW_C_API_DEF
-        NPY_ALLOW_C_API
-        if (_PyDataMem_eventhook != NULL) {
-            (*_PyDataMem_eventhook)(NULL, result, size,
-                                    _PyDataMem_eventhook_user_data);
-        }
-        NPY_DISABLE_C_API
-    }
-    return result;
-}
-
-/*NUMPY_API
- * Allocates zeroed memory for array data.
- */
-NPY_NO_EXPORT void *
-PyDataMem_NEW_ZEROED(size_t size, size_t elsize)
-{
-    void *result;
-
-    result = calloc(size, elsize);
-    if (_PyDataMem_eventhook != NULL) {
-        NPY_ALLOW_C_API_DEF
-        NPY_ALLOW_C_API
-        if (_PyDataMem_eventhook != NULL) {
-            (*_PyDataMem_eventhook)(NULL, result, size * elsize,
-                                    _PyDataMem_eventhook_user_data);
-        }
-        NPY_DISABLE_C_API
-    }
-    return result;
-}
-
-/*NUMPY_API
- * Free memory for array data.
- */
-NPY_NO_EXPORT void
-PyDataMem_FREE(void *ptr)
-{
-    free(ptr);
-    if (_PyDataMem_eventhook != NULL) {
-        NPY_ALLOW_C_API_DEF
-        NPY_ALLOW_C_API
-        if (_PyDataMem_eventhook != NULL) {
-            (*_PyDataMem_eventhook)(ptr, NULL, 0,
-                                    _PyDataMem_eventhook_user_data);
-        }
-        NPY_DISABLE_C_API
-    }
-}
-
-/*NUMPY_API
- * Reallocate/resize memory for array data.
- */
-NPY_NO_EXPORT void *
-PyDataMem_RENEW(void *ptr, size_t size)
-{
-    void *result;
-
-    result = realloc(ptr, size);
-    if (_PyDataMem_eventhook != NULL) {
-        NPY_ALLOW_C_API_DEF
-        NPY_ALLOW_C_API
-        if (_PyDataMem_eventhook != NULL) {
-            (*_PyDataMem_eventhook)(ptr, result, size,
-                                    _PyDataMem_eventhook_user_data);
-        }
-        NPY_DISABLE_C_API
-    }
-    return result;
-}
-
 static PyObject *
 array_may_share_memory(PyObject *NPY_UNUSED(ignored), PyObject *args)
 {
diff --git a/numpy/core/src/multiarray/multiarraymodule_onefile.c b/numpy/core/src/multiarray/multiarraymodule_onefile.c
index 9410263e4..2d05c20ef 100644
--- a/numpy/core/src/multiarray/multiarraymodule_onefile.c
+++ b/numpy/core/src/multiarray/multiarraymodule_onefile.c
@@ -21,6 +21,7 @@
 
 #include "descriptor.c"
 #include "flagsobject.c"
+#include "alloc.c"
 #include "ctors.c"
 #include "iterators.c"
 #include "mapping.c"
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 7ba6f0779..8b2bc1d82 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -3990,7 +3990,8 @@ class TestMemEventHook(TestCase):
         # multiarray/multiarray_tests.c.src
         test_pydatamem_seteventhook_start()
         # force an allocation and free of a numpy array
-        a = np.zeros(10)
+        # needs to be larger then limit of small memory cacher in ctors.c
+        a = np.zeros(1000)
         del a
         test_pydatamem_seteventhook_end()