[Numpy-discussion] Proposed change from POSIX to PyMem_RawXXX (plain text resend)

[Daniel Waddington]

Hi,
I'm working with Numpy in the context of supporting different memory types such as persistent memory and CXL attached.  I would like to propose a minor change, but figured I would get some initial feedback from the developer community before submitting a PR.
 
In multiarray/alloc.c the allocator (beneath the cache) using the POSIX malloc/calloc/realloc/free.  I propose that these should be changed to PyMem_RawXXX equivalents.  The reason for this is that by doing so, one can use the python custom allocator functions (e.g. PyMem_GetAllocator/PyMem_SetAllocator) to intercept the memory allocator for NumPy arrays.  This will be useful as heterogeneous memories need supporting. I don't think this will drastically change performance but it is an extra function redirection (and it will only impact when the cache can't deliver).
 
There are likely other places in NumPy that could do with a rinse and repeat - may be someone could advise?
 
Thanks,
Daniel Waddington
IBM Research
 
---
Example patch for 1.19.x (I'm building with Python3.6)
 
diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c
index 795fc7315..e9e888478 100644
--- a/numpy/core/src/multiarray/alloc.c
+++ b/numpy/core/src/multiarray/alloc.c
@@ -248,7 +248,7 @@ PyDataMem_NEW(size_t size)
     void *result;
 
     assert(size != 0);
-    result = malloc(size);
+    result = PyMem_RawMalloc(size);
     if (_PyDataMem_eventhook != NULL) {
         NPY_ALLOW_C_API_DEF
         NPY_ALLOW_C_API
@@ -270,7 +270,7 @@ PyDataMem_NEW_ZEROED(size_t size, size_t elsize)
 {
     void *result;
 
-    result = calloc(size, elsize);
+    result = PyMem_RawCalloc(size, elsize);
     if (_PyDataMem_eventhook != NULL) {
         NPY_ALLOW_C_API_DEF
         NPY_ALLOW_C_API
@@ -291,7 +291,7 @@ NPY_NO_EXPORT void
 PyDataMem_FREE(void *ptr)
 {
     PyTraceMalloc_Untrack(NPY_TRACE_DOMAIN, (npy_uintp)ptr);
-    free(ptr);
+    PyMem_RawFree(ptr);