decouple the cache allocation and memory allocation.

numpy/core/src/multiarray/alloc.c

@@ -11,79 +11,49 @@
 
 #include <assert.h>
 
-#define NBUCKETS 1024 /* number of buckets for data*/
+#define NBUCKETS_DATA 1024 /* number of buckets for data*/
 #define NBUCKETS_DIM 16 /* number of buckets for dimensions/strides */
-#define NCACHE 7 /* number of cache entries per bucket */
+#define NCACHE_DATA 7 /* number of cache entries per bucket */
+#define NCACHE_DIM 15 /* number of cache entries per bucket */
+
 /* this structure fits neatly into a cacheline */
 typedef struct {
     npy_uintp available; /* number of cached pointers */
-    void * ptrs[NCACHE];
-} cache_bucket;
-static cache_bucket datacache[NBUCKETS];
-static cache_bucket dimcache[NBUCKETS_DIM];
-
-/*
- * 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,
-                 cache_bucket * cache, void * (*alloc)(size_t))
-{
-    assert((esz == 1 && cache == datacache) ||
-           (esz == sizeof(npy_intp) && cache == dimcache));
-    if (nelem < msz) {
-        if (cache[nelem].available > 0) {
-            return cache[nelem].ptrs[--(cache[nelem].available)];
-        }
-    }
-    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,
-                cache_bucket * cache, void (*dealloc)(void *))
-{
-    if (p != NULL && nelem < msz) {
-        if (cache[nelem].available < NCACHE) {
-            cache[nelem].ptrs[cache[nelem].available++] = p;
-            return;
-        }
-    }
-    dealloc(p);
-}
+    void * ptrs[NCACHE_DATA];
+} data_cache_bucket;
 
+typedef struct {
+    npy_uintp available; /* number of cached pointers */
+    void * ptrs[NCACHE_DIM];
+} dim_cache_bucket;
 
-static NPY_INLINE void *
-_PyDataMem_NEW_ZEROED_RAW(size_t size) {
-    return PyDataMem_NEW_ZEROED(size, 1);
-}
+static data_cache_bucket datacache[NBUCKETS_DATA];
+static dim_cache_bucket dimcache[NBUCKETS_DIM];
 
 /*
  * 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);
+    if (sz < NBUCKETS_DATA && datacache[sz].available > 0) {
+        return datacache[sz].ptrs[--(datacache[sz].available)];
+    }
+    return PyDataMem_NEW(sz);
 }
 
-
 /* 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) {
-        return _npy_alloc_cache(sz, 1, NBUCKETS, datacache, &_PyDataMem_NEW_ZEROED_RAW);
+    if (sz < NBUCKETS_DATA && datacache[sz].available > 0) {
+        p = datacache[sz].ptrs[--(datacache[sz].available)];
+        memset(p, 0, sz);
+        return p;
     }
+
     NPY_BEGIN_THREADS;
     p = PyDataMem_NEW_ZEROED(sz, 1);
     NPY_END_THREADS;
@@ -93,7 +63,11 @@ npy_alloc_cache_zero(npy_uintp sz)
 NPY_NO_EXPORT void
 npy_free_cache(void * p, npy_uintp sz)
 {
-    _npy_free_cache(p, sz, NBUCKETS, datacache, &PyDataMem_FREE);
+    if (p != NULL && sz < NBUCKETS_DATA && datacache[sz].available < NCACHE_DATA) {
+        datacache[sz].ptrs[datacache[sz].available++] = p;
+        return ;
+    }
+    PyDataMem_FREE(p);
 }
 
 /*
@@ -103,12 +77,18 @@ npy_free_cache(void * p, npy_uintp sz)
 NPY_NO_EXPORT void *
 npy_alloc_cache_dim(npy_uintp sz)
 {
+    void * p;
     /* dims + strides */
     if (NPY_UNLIKELY(sz < 2)) {
         sz = 2;
     }
-    return _npy_alloc_cache(sz, sizeof(npy_intp), NBUCKETS_DIM, dimcache,
-                            &PyArray_malloc);
+
+    if (sz < NBUCKETS_DIM && dimcache[sz].available > 0) {
+            return dimcache[sz].ptrs[--(dimcache[sz].available)];
+    }
+
+    p = PyArray_malloc(sz * sizeof(npy_intp));
+    return p;
 }
 
 NPY_NO_EXPORT void
@@ -118,8 +98,12 @@ npy_free_cache_dim(void * p, npy_uintp sz)
     if (NPY_UNLIKELY(sz < 2)) {
         sz = 2;
     }
-    _npy_free_cache(p, sz, NBUCKETS_DIM, dimcache,
-                    &PyArray_free);
+
+    if (p != NULL && sz < NBUCKETS_DIM && dimcache[sz].available < NCACHE_DIM) {
+        dimcache[sz].ptrs[dimcache[sz].available++] = p;
+        return ;
+    }
+    PyArray_free(p);
 }