ENH: Reduce overhead of configurable data allocation strategy (NEP49)

doc/source/reference/c-api/data_memory.rst

@@ -55,6 +55,8 @@ functions may change during the lifetime of the process, each ``ndarray``
 carries with it the functions used at the time of its instantiation, and these
 will be used to reallocate or free the data memory of the instance.
 
+For details see: :ref:`NEP 49 — Data allocation strategies <NEP49>`.
+
 .. c:type:: PyDataMem_Handler
 
     A struct to hold function pointers used to manipulate memory
@@ -87,6 +89,8 @@ will be used to reallocate or free the data memory of the instance.
    return ``NULL`` if an error has occurred. We wrap the user-provided functions
    so they will still call the python and numpy memory management callback
    hooks.
+   The handlers are stored in a Python context variable (see https://docs.python.org/3/library/contextvars.html),
+   so there can be multiple handlers in a Python session.
 
 .. c:function:: PyObject * PyDataMem_GetHandler()
 

numpy/core/src/multiarray/alloc.c

@@ -430,11 +430,25 @@ int uo_index=0;   /* user_override index */
 
 /* Wrappers for the default or any user-assigned PyDataMem_Handler */
 
+int default_allocator_policy = 1;
+
+static inline PyDataMem_Handler *
+_PyDataMem_GetHandler_Internal(PyObject *mem_handler)
+{
+    if (PyDataMem_DefaultHandler == mem_handler)
+        // fast path for default allocator
+        return &default_handler;
+
+    PyDataMem_Handler *handler = (PyDataMem_Handler *)PyCapsule_GetPointer(
+            mem_handler, "mem_handler");
+    return handler;
+}
+
 NPY_NO_EXPORT void *
 PyDataMem_UserNEW(size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = _PyDataMem_GetHandler_Internal(mem_handler);
     if (handler == NULL) {
         return NULL;
     }
@@ -457,7 +471,7 @@ NPY_NO_EXPORT void *
 PyDataMem_UserNEW_ZEROED(size_t nmemb, size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = _PyDataMem_GetHandler_Internal(mem_handler);
     if (handler == NULL) {
         return NULL;
     }
@@ -479,7 +493,7 @@ PyDataMem_UserNEW_ZEROED(size_t nmemb, size_t size, PyObject *mem_handler)
 NPY_NO_EXPORT void
 PyDataMem_UserFREE(void *ptr, size_t size, PyObject *mem_handler)
 {
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = _PyDataMem_GetHandler_Internal(mem_handler);
     if (handler == NULL) {
         WARN_NO_RETURN(PyExc_RuntimeWarning,
                      "Could not get pointer to 'mem_handler' from PyCapsule");
@@ -502,7 +516,7 @@ NPY_NO_EXPORT void *
 PyDataMem_UserRENEW(void *ptr, size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = _PyDataMem_GetHandler_Internal(mem_handler);
     if (handler == NULL) {
         return NULL;
     }
@@ -535,6 +549,9 @@ PyDataMem_UserRENEW(void *ptr, size_t size, PyObject *mem_handler)
 NPY_NO_EXPORT PyObject *
 PyDataMem_SetHandler(PyObject *handler)
 {
+    // once the user sets an allocation policy, we cannot guarantee the default allocator without checking the context
+    default_allocator_policy = 0;
+
     PyObject *old_handler;
     PyObject *token;
     if (PyContextVar_Get(current_handler, NULL, &old_handler)) {
@@ -560,6 +577,11 @@ NPY_NO_EXPORT PyObject *
 PyDataMem_GetHandler()
 {
     PyObject *handler;
+
+    if (default_allocator_policy) {
+        Py_INCREF(PyDataMem_DefaultHandler);
+        return PyDataMem_DefaultHandler;
+    }
     if (PyContextVar_Get(current_handler, NULL, &handler)) {
         return NULL;
     }