diff --git a/numpy/core/code_generators/cversions.txt b/numpy/core/code_generators/cversions.txt index acfced81262a..f017db2ef193 100644 --- a/numpy/core/code_generators/cversions.txt +++ b/numpy/core/code_generators/cversions.txt @@ -29,3 +29,6 @@ # The interface has not changed, but the hash is different due to # the annotations, so keep the previous version number. 0x00000009 = 982c4ebb6e7e4c194bf46b1535b4ef1b + +# Version 10 (NumPy 1.10) Added configurable array data allocator. +0x0000000a = 7c6ff1780a699fd67e75dc0bc180a994 diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py index 972966627b99..539f97b03bab 100644 --- a/numpy/core/code_generators/numpy_api.py +++ b/numpy/core/code_generators/numpy_api.py @@ -344,6 +344,8 @@ # End 1.9 API 'PyArray_CheckAnyScalarExact': (300, NonNull(1)), # End 1.10 API + 'PyDataMem_GetAllocator': (301,), + 'PyDataMem_SetAllocator': (302,), } ufunc_types_api = { diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h index 78f79d5feaef..9c55cb3d3775 100644 --- a/numpy/core/include/numpy/ndarraytypes.h +++ b/numpy/core/include/numpy/ndarraytypes.h @@ -626,6 +626,26 @@ typedef struct _arr_descr { PyObject *shape; /* a tuple */ } PyArray_ArrayDescr; +/* + * This is a function for hooking into the PyDataMem_NEW/FREE/RENEW functions. + * See the documentation for PyDataMem_SetEventHook. + */ +typedef void (PyDataMem_EventHookFunc)(void *inp, void *outp, size_t size, + void *user_data); + +/* Allocator structure for array data */ +typedef void *(PyDataMem_AllocFunc)(size_t size); +typedef void *(PyDataMem_ZeroedAllocFunc)(size_t nelems, size_t elsize); +typedef void (PyDataMem_FreeFunc)(void *ptr); +typedef void *(PyDataMem_ReallocFunc)(void *ptr, size_t size); + +typedef struct { + PyDataMem_AllocFunc *alloc; + PyDataMem_ZeroedAllocFunc *zeroed_alloc; + PyDataMem_FreeFunc *free; + PyDataMem_ReallocFunc *realloc; +} PyDataMem_Allocator; + /* * The main array object structure. * @@ -675,6 +695,8 @@ typedef struct tagPyArrayObject_fields { int flags; /* For weak references */ PyObject *weakreflist; + /* For resizes and deallocation */ + const PyDataMem_Allocator *allocator; } PyArrayObject_fields; /* @@ -1536,6 +1558,12 @@ PyArray_SETITEM(PyArrayObject *arr, char *itemptr, PyObject *v) v, itemptr, arr); } +static NPY_INLINE const PyDataMem_Allocator * +PyArray_ALLOCATOR(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->allocator; +} + #else /* These macros are deprecated as of NumPy 1.7. */ @@ -1778,13 +1806,6 @@ typedef struct { */ } PyArrayInterface; -/* - * This is a function for hooking into the PyDataMem_NEW/FREE/RENEW functions. - * See the documentation for PyDataMem_SetEventHook. - */ -typedef void (PyDataMem_EventHookFunc)(void *inp, void *outp, size_t size, - void *user_data); - /* * Use the keyword NPY_DEPRECATED_INCLUDES to ensure that the header files * npy_*_*_deprecated_api.h are only included from here and nowhere else. diff --git a/numpy/core/setup_common.py b/numpy/core/setup_common.py index 0b18bc6c62fa..e7116c87b492 100644 --- a/numpy/core/setup_common.py +++ b/numpy/core/setup_common.py @@ -35,7 +35,8 @@ # 0x00000008 - 1.7.x # 0x00000009 - 1.8.x # 0x00000009 - 1.9.x -C_API_VERSION = 0x00000009 +# 0x0000000A - 1.10.x +C_API_VERSION = 0x0000000A class MismatchCAPIWarning(Warning): pass diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c index 77eb0416bc34..fa583c1ffe7d 100644 --- a/numpy/core/src/multiarray/alloc.c +++ b/numpy/core/src/multiarray/alloc.c @@ -11,6 +11,150 @@ #include +/* 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; +} + +/* Allocator structure */ +static PyDataMem_Allocator default_allocator = { + malloc, + calloc, + free, + realloc +}; +static const PyDataMem_Allocator *current_allocator = &default_allocator; + +static NPY_INLINE void * +allocator_new(const PyDataMem_Allocator *a, size_t size) +{ + void *result; + + result = a->alloc(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; +} + +static NPY_INLINE void * +allocator_new_zeroed(const PyDataMem_Allocator *a, size_t size, size_t elsize) +{ + void *result; + + result = a->zeroed_alloc(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; +} + +static NPY_INLINE void +allocator_free(const PyDataMem_Allocator *a, void *ptr) +{ + a->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 + } +} + +static NPY_INLINE void * +allocator_renew(const PyDataMem_Allocator *a, void *ptr, size_t size) +{ + void *result; + + result = a->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 NPY_INLINE void * +current_allocator_new(size_t size) +{ + return allocator_new(current_allocator, size); +} + +static NPY_INLINE void * +current_allocator_new_zeroed(size_t size, size_t elsize) +{ + return allocator_new_zeroed(current_allocator, size, elsize); +} + +static NPY_INLINE void +current_allocator_free(void *ptr) +{ + allocator_free(current_allocator, ptr); +} + +static NPY_INLINE void * +current_allocator_renew(void *ptr, size_t size) +{ + return allocator_renew(current_allocator, ptr, size); +} + + #define NBUCKETS 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 */ @@ -22,6 +166,7 @@ typedef struct { static cache_bucket datacache[NBUCKETS]; static cache_bucket dimcache[NBUCKETS_DIM]; + /* * very simplistic small memory block cache to avoid more expensive libc * allocations @@ -59,39 +204,61 @@ _npy_free_cache(void * p, npy_uintp nelem, npy_uint msz, dealloc(p); } +/* + * clear all cache data in the given cache + */ +static void +_npy_clear_cache(npy_uint msz, cache_bucket * cache, void (*dealloc)(void *)) +{ + npy_intp i, nelem; + for (nelem = 0; nelem < msz; nelem++) { + for (i = 0; i < cache[nelem].available; i++) { + dealloc(cache[nelem].ptrs[i]); + } + cache[nelem].available = 0; + } +} + /* * array data cache, sz is number of bytes to allocate */ NPY_NO_EXPORT void * -npy_alloc_cache(npy_uintp sz) +npy_alloc_cache(const PyDataMem_Allocator **a, npy_uintp sz) { - return _npy_alloc_cache(sz, 1, NBUCKETS, datacache, &PyDataMem_NEW); + *a = current_allocator; + return _npy_alloc_cache(sz, 1, NBUCKETS, datacache, ¤t_allocator_new); } /* zero initialized data, sz is number of bytes to allocate */ NPY_NO_EXPORT void * -npy_alloc_cache_zero(npy_uintp sz) +npy_alloc_cache_zero(const PyDataMem_Allocator **a, npy_uintp sz) { void * p; NPY_BEGIN_THREADS_DEF; + *a = current_allocator; if (sz < NBUCKETS) { - p = _npy_alloc_cache(sz, 1, NBUCKETS, datacache, &PyDataMem_NEW); + p = _npy_alloc_cache(sz, 1, NBUCKETS, datacache, ¤t_allocator_new); if (p) { memset(p, 0, sz); } return p; } NPY_BEGIN_THREADS; - p = PyDataMem_NEW_ZEROED(sz, 1); + p = current_allocator_new_zeroed(sz, 1); NPY_END_THREADS; return p; } NPY_NO_EXPORT void -npy_free_cache(void * p, npy_uintp sz) +npy_free_cache(const PyDataMem_Allocator *a, void * p, npy_uintp sz) { - _npy_free_cache(p, sz, NBUCKETS, datacache, &PyDataMem_FREE); + if (a == current_allocator) { + _npy_free_cache(p, sz, NBUCKETS, datacache, ¤t_allocator_free); + } + else { + allocator_free(a, p); + } } /* @@ -120,47 +287,9 @@ npy_free_cache_dim(void * p, npy_uintp sz) &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) +/* NOTE: the PyDataMem_ functions use the default allocator, as they are + * called for other things than just array data. */ -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. @@ -168,19 +297,7 @@ PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook, 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; + return allocator_new(&default_allocator, size); } /*NUMPY_API @@ -189,19 +306,7 @@ PyDataMem_NEW(size_t size) 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; + return allocator_new_zeroed(&default_allocator, size, elsize); } /*NUMPY_API @@ -210,16 +315,7 @@ PyDataMem_NEW_ZEROED(size_t size, size_t elsize) 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 - } + allocator_free(&default_allocator, ptr); } /*NUMPY_API @@ -228,17 +324,30 @@ PyDataMem_FREE(void *ptr) NPY_NO_EXPORT void * PyDataMem_RENEW(void *ptr, size_t size) { - void *result; + return allocator_renew(&default_allocator, ptr, size); +} - 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 + +/*NUMPY_API + * Get the current allocator for array data. + */ +NPY_NO_EXPORT const PyDataMem_Allocator * +PyDataMem_GetAllocator(void) +{ + return current_allocator; +} + +/*NUMPY_API + * Set the current allocator for array data. If the parameter is NULL, + * the allocator is reset to the default Numpy allocator. + */ +NPY_NO_EXPORT void +PyDataMem_SetAllocator(const PyDataMem_Allocator *new_allocator) +{ + /* We must deallocate all cached data areas before switching allocators */ + _npy_clear_cache(NBUCKETS, datacache, ¤t_allocator_free); + if (new_allocator == NULL) { + new_allocator = &default_allocator; } - return result; + current_allocator = new_allocator; } diff --git a/numpy/core/src/multiarray/alloc.h b/numpy/core/src/multiarray/alloc.h index 8f6b167d0380..6bf816a5e3ee 100644 --- a/numpy/core/src/multiarray/alloc.h +++ b/numpy/core/src/multiarray/alloc.h @@ -5,13 +5,13 @@ #include NPY_NO_EXPORT void * -npy_alloc_cache(npy_uintp sz); +npy_alloc_cache(const PyDataMem_Allocator **, npy_uintp sz); NPY_NO_EXPORT void * -npy_alloc_cache_zero(npy_uintp sz); +npy_alloc_cache_zero(const PyDataMem_Allocator **, npy_uintp sz); NPY_NO_EXPORT void -npy_free_cache(void * p, npy_uintp sd); +npy_free_cache(const PyDataMem_Allocator *, void * p, npy_uintp sd); NPY_NO_EXPORT void * npy_alloc_cache_dim(npy_uintp sz); diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c index 6e48ef381363..f5f5d591378b 100644 --- a/numpy/core/src/multiarray/arrayobject.c +++ b/numpy/core/src/multiarray/arrayobject.c @@ -420,7 +420,7 @@ array_dealloc(PyArrayObject *self) * self already... */ } - npy_free_cache(fa->data, PyArray_NBYTES(self)); + npy_free_cache(fa->allocator, fa->data, PyArray_NBYTES(self)); } /* must match allocation in PyArray_NewFromDescr */ diff --git a/numpy/core/src/multiarray/arraytypes.c.src b/numpy/core/src/multiarray/arraytypes.c.src index 685aba542f2d..69bb63831168 100644 --- a/numpy/core/src/multiarray/arraytypes.c.src +++ b/numpy/core/src/multiarray/arraytypes.c.src @@ -2692,6 +2692,7 @@ VOID_compare(char *ip1, char *ip2, PyArrayObject *ap) PyObject *tup; char *nip1, *nip2; int i, res = 0, swap=0; + const PyDataMem_Allocator *allocator = NULL; if (!PyArray_HASFIELDS(ap)) { return STRING_compare(ip1, ip2, ap); @@ -2721,7 +2722,7 @@ VOID_compare(char *ip1, char *ip2, PyArrayObject *ap) if ((swap) || (new->alignment > 1)) { if ((swap) || (!npy_is_aligned(nip1, new->alignment))) { /* create buffer and copy */ - nip1 = npy_alloc_cache(new->elsize); + nip1 = npy_alloc_cache(&allocator, new->elsize); if (nip1 == NULL) { goto finish; } @@ -2731,10 +2732,10 @@ VOID_compare(char *ip1, char *ip2, PyArrayObject *ap) } if ((swap) || (!npy_is_aligned(nip2, new->alignment))) { /* copy data to a buffer */ - nip2 = npy_alloc_cache(new->elsize); + nip2 = npy_alloc_cache(&allocator, new->elsize); if (nip2 == NULL) { if (nip1 != ip1 + offset) { - npy_free_cache(nip1, new->elsize); + npy_free_cache(allocator, nip1, new->elsize); } goto finish; } @@ -2746,10 +2747,10 @@ VOID_compare(char *ip1, char *ip2, PyArrayObject *ap) res = new->f->compare(nip1, nip2, ap); if ((swap) || (new->alignment > 1)) { if (nip1 != ip1 + offset) { - npy_free_cache(nip1, new->elsize); + npy_free_cache(allocator, nip1, new->elsize); } if (nip2 != ip2 + offset) { - npy_free_cache(nip2, new->elsize); + npy_free_cache(allocator, nip2, new->elsize); } } if (res != 0) { diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c index 010420826f1b..8c052b50e8ba 100644 --- a/numpy/core/src/multiarray/ctors.c +++ b/numpy/core/src/multiarray/ctors.c @@ -991,6 +991,7 @@ PyArray_NewFromDescr_int(PyTypeObject *subtype, PyArray_Descr *descr, int nd, fa->descr = descr; fa->base = (PyObject *)NULL; fa->weakreflist = (PyObject *)NULL; + fa->allocator = NULL; if (nd > 0) { fa->dimensions = npy_alloc_cache_dim(2 * nd); @@ -1033,17 +1034,16 @@ 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 = npy_alloc_cache_zero(sd); + data = npy_alloc_cache_zero(&fa->allocator, sd); } else { - data = npy_alloc_cache(sd); + data = npy_alloc_cache(&fa->allocator, sd); } if (data == NULL) { PyErr_NoMemory(); goto fail; } fa->flags |= NPY_ARRAY_OWNDATA; - } else { /* @@ -3226,7 +3226,7 @@ array_from_text(PyArray_Descr *dtype, npy_intp num, char *sep, size_t *nread, dptr += dtype->elsize; if (num < 0 && thisbuf == size) { totalbytes += bytes; - tmp = PyDataMem_RENEW(PyArray_DATA(r), totalbytes); + tmp = PyArray_ALLOCATOR(r)->realloc(PyArray_DATA(r), totalbytes); if (tmp == NULL) { err = 1; break; @@ -3240,7 +3240,8 @@ array_from_text(PyArray_Descr *dtype, npy_intp num, char *sep, size_t *nread, } } if (num < 0) { - tmp = PyDataMem_RENEW(PyArray_DATA(r), PyArray_MAX(*nread,1)*dtype->elsize); + tmp = PyArray_ALLOCATOR(r)->realloc(PyArray_DATA(r), + PyArray_MAX(*nread,1)*dtype->elsize); if (tmp == NULL) { err = 1; } @@ -3322,9 +3323,9 @@ PyArray_FromFile(FILE *fp, PyArray_Descr *dtype, npy_intp num, char *sep) if (((npy_intp) nread) < num) { /* Realloc memory for smaller number of elements */ const size_t nsize = PyArray_MAX(nread,1)*PyArray_DESCR(ret)->elsize; - char *tmp; + char *tmp = PyArray_ALLOCATOR(ret)->realloc(PyArray_DATA(ret), nsize); - if((tmp = PyDataMem_RENEW(PyArray_DATA(ret), nsize)) == NULL) { + if(tmp == NULL) { Py_DECREF(ret); return PyErr_NoMemory(); } @@ -3604,7 +3605,8 @@ PyArray_FromIter(PyObject *obj, PyArray_Descr *dtype, npy_intp count) */ elcount = (i >> 1) + (i < 4 ? 4 : 2) + i; if (elcount <= NPY_MAX_INTP/elsize) { - new_data = PyDataMem_RENEW(PyArray_DATA(ret), elcount * elsize); + new_data = PyArray_ALLOCATOR(ret)->realloc(PyArray_DATA(ret), + elcount * elsize); } else { new_data = NULL; @@ -3642,10 +3644,10 @@ PyArray_FromIter(PyObject *obj, PyArray_Descr *dtype, npy_intp count) * (assuming realloc is reasonably good about reusing space...) */ if (i == 0) { - /* The size cannot be zero for PyDataMem_RENEW. */ + /* The size cannot be zero for ->realloc(). */ i = 1; } - new_data = PyDataMem_RENEW(PyArray_DATA(ret), i * elsize); + new_data = PyArray_ALLOCATOR(ret)->realloc(PyArray_DATA(ret), i * elsize); if (new_data == NULL) { PyErr_SetString(PyExc_MemoryError, "cannot allocate array memory"); diff --git a/numpy/core/src/multiarray/getset.c b/numpy/core/src/multiarray/getset.c index 5374f42e7ec7..4930f67202a0 100644 --- a/numpy/core/src/multiarray/getset.c +++ b/numpy/core/src/multiarray/getset.c @@ -357,7 +357,7 @@ array_data_set(PyArrayObject *self, PyObject *op) } if (PyArray_FLAGS(self) & NPY_ARRAY_OWNDATA) { PyArray_XDECREF(self); - PyDataMem_FREE(PyArray_DATA(self)); + PyArray_ALLOCATOR(self)->free(PyArray_DATA(self)); } if (PyArray_BASE(self)) { if (PyArray_FLAGS(self) & NPY_ARRAY_UPDATEIFCOPY) { diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 851a374562f9..81dfc429bb2b 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -1701,7 +1701,7 @@ array_setstate(PyArrayObject *self, PyObject *args) } if ((PyArray_FLAGS(self) & NPY_ARRAY_OWNDATA)) { - PyDataMem_FREE(PyArray_DATA(self)); + PyArray_ALLOCATOR(self)->free(PyArray_DATA(self)); PyArray_CLEARFLAGS(self, NPY_ARRAY_OWNDATA); } Py_XDECREF(PyArray_BASE(self)); @@ -1744,7 +1744,7 @@ array_setstate(PyArrayObject *self, PyObject *args) if (!_IsAligned(self) || swap || (len <= 1000)) { #endif npy_intp num = PyArray_NBYTES(self); - fa->data = PyDataMem_NEW(num); + fa->data = npy_alloc_cache(&fa->allocator, num); if (PyArray_DATA(self) == NULL) { fa->nd = 0; PyDimMem_FREE(PyArray_DIMS(self)); @@ -1787,10 +1787,10 @@ array_setstate(PyArrayObject *self, PyObject *args) } } else { - fa->data = PyDataMem_NEW(PyArray_NBYTES(self)); + fa->data = npy_alloc_cache(&fa->allocator, PyArray_NBYTES(self)); if (PyArray_DATA(self) == NULL) { fa->nd = 0; - fa->data = PyDataMem_NEW(PyArray_DESCR(self)->elsize); + fa->data = npy_alloc_cache(&fa->allocator, PyArray_DESCR(self)->elsize); PyDimMem_FREE(PyArray_DIMS(self)); return PyErr_NoMemory(); } diff --git a/numpy/core/src/multiarray/multiarray_tests.c.src b/numpy/core/src/multiarray/multiarray_tests.c.src index a80d3af191f9..207c99573877 100644 --- a/numpy/core/src/multiarray/multiarray_tests.c.src +++ b/numpy/core/src/multiarray/multiarray_tests.c.src @@ -949,6 +949,114 @@ test_nditer_too_large(PyObject *NPY_UNUSED(self), PyObject *args) { return NULL; } +/* + * Test allocator to exercise setting the default PyDataMem_Allocator. + * It pads allocations and sets a canary to make sure another allocator + * can't be used on the same pointers. + */ + +static int alloc_counter = 0; +static int realloc_counter = 0; +static int free_counter = 0; + +#define ALLOC_PAD 16 + +static void +store_canary(void *ptr) +{ + *(npy_intp *) ptr = (npy_intp) ptr; +} + +static void +check_canary(void *ptr) +{ + if (*(npy_intp *) ptr != (npy_intp) ptr) { + char msg[300]; + sprintf(msg, "Canary check fail at %p: expected %p, got %p", + ptr, *(npy_intp *) ptr, (npy_intp) ptr); + Py_FatalError(msg); + } +} + +static void * +testalloc_alloc(size_t size) +{ + void *base; + alloc_counter += 1; + base = malloc(size + ALLOC_PAD); + store_canary(base); + return (void *) ((npy_intp) base + ALLOC_PAD); +} + +static void * +testalloc_zeroed_alloc(size_t nelems, size_t elsize) +{ + void *base; + alloc_counter += 1; + base = calloc(nelems * elsize + ALLOC_PAD, 1); + store_canary(base); + return (void *) ((npy_intp) base + ALLOC_PAD); +} + +static void +testalloc_free(void *ptr) +{ + void *base = (void *) ((npy_intp) ptr - ALLOC_PAD); + free_counter += 1; + check_canary(base); + free(base); +} + +static void * +testalloc_realloc(void *ptr, size_t size) +{ + void *base = (void *) ((npy_intp) ptr - ALLOC_PAD); + realloc_counter += 1; + check_canary(base); + base = realloc(base, size + ALLOC_PAD); + store_canary(base); + return (void *) ((npy_intp) base + ALLOC_PAD); +} + +static PyDataMem_Allocator test_allocator = { + testalloc_alloc, + testalloc_zeroed_alloc, + testalloc_free, + testalloc_realloc, +}; + +static PyObject * +enable_test_allocator(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) +{ + alloc_counter = 0; + realloc_counter = 0; + free_counter = 0; + PyDataMem_SetAllocator(&test_allocator); + if (PyDataMem_GetAllocator() != &test_allocator) { + PyDataMem_SetAllocator(NULL); + PyErr_SetString(PyExc_AssertionError, + "wrong PyDataMem_GetAllocator() result after " + "calling PyDataMem_SetAllocator()"); + return NULL; + } + Py_RETURN_NONE; +} + +static PyObject * +disable_test_allocator(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) +{ + PyDataMem_SetAllocator(NULL); + Py_RETURN_NONE; +} + +static PyObject * +get_test_allocator_counts(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) +{ + return Py_BuildValue("{sisisi}", + "allocs", alloc_counter, + "reallocs", realloc_counter, + "frees", free_counter); +} static PyMethodDef Multiarray_TestsMethods[] = { {"IsPythonScalar", @@ -992,6 +1100,15 @@ static PyMethodDef Multiarray_TestsMethods[] = { {"test_nditer_too_large", test_nditer_too_large, METH_VARARGS, NULL}, + {"enable_test_allocator", + enable_test_allocator, + METH_NOARGS, NULL}, + {"disable_test_allocator", + disable_test_allocator, + METH_NOARGS, NULL}, + {"get_test_allocator_counts", + get_test_allocator_counts, + METH_NOARGS, NULL}, {NULL, NULL, 0, NULL} /* Sentinel */ }; diff --git a/numpy/core/src/multiarray/shape.c b/numpy/core/src/multiarray/shape.c index f1e81ff6b5f9..76447cce6509 100644 --- a/numpy/core/src/multiarray/shape.c +++ b/numpy/core/src/multiarray/shape.c @@ -106,7 +106,7 @@ PyArray_Resize(PyArrayObject *self, PyArray_Dims *newshape, int refcheck, sd = newsize*PyArray_DESCR(self)->elsize; } /* Reallocate space if needed */ - new_data = PyDataMem_RENEW(PyArray_DATA(self), sd); + new_data = PyArray_ALLOCATOR(self)->realloc(PyArray_DATA(self), sd); if (new_data == NULL) { PyErr_SetString(PyExc_MemoryError, "cannot allocate memory for array"); diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py index b0d6770527f5..d0bb35292677 100644 --- a/numpy/core/tests/test_multiarray.py +++ b/numpy/core/tests/test_multiarray.py @@ -8,6 +8,7 @@ import warnings import operator import io +import pickle if sys.version_info[0] >= 3: import builtins else: @@ -23,7 +24,8 @@ from numpy.core.multiarray_tests import ( test_neighborhood_iterator, test_neighborhood_iterator_oob, test_pydatamem_seteventhook_start, test_pydatamem_seteventhook_end, - test_inplace_increment, get_buffer_info, test_as_c_array + test_inplace_increment, get_buffer_info, test_as_c_array, + enable_test_allocator, disable_test_allocator, get_test_allocator_counts ) from numpy.testing import ( TestCase, run_module_suite, assert_, assert_raises, @@ -4987,5 +4989,31 @@ def test_collections_hashable(self): self.assertFalse(isinstance(x, collections.Hashable)) +class TestDataMemAllocator(TestCase): + """ + Test setting the array data allocator. + """ + + def test_set_allocator(self): + x = np.empty(10) + enable_test_allocator() + try: + np.empty(10).resize((101,)) + np.zeros(11).resize((102,)) + a = np.empty(12) + b = np.zeros(13) + c = pickle.loads(pickle.dumps(b, protocol=-1)) + del x + finally: + disable_test_allocator() + a.resize((103,)) + b.resize((104,)) + del a, b, c + counts = get_test_allocator_counts() + self.assertEqual(counts['allocs'], 6) + self.assertEqual(counts['frees'], 6) + self.assertEqual(counts['reallocs'], 4) + + if __name__ == "__main__": run_module_suite()