diff --git a/doc/source/reference/c-api/data_memory.rst b/doc/source/reference/c-api/data_memory.rst index 2084ab5d083f..5cf439912699 100644 --- a/doc/source/reference/c-api/data_memory.rst +++ b/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 `. + .. c:type:: PyDataMem_Handler A struct to hold function pointers used to manipulate memory @@ -87,6 +89,9 @@ 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() diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c index 8b765aa954da..dc1a8c37cb27 100644 --- a/numpy/core/src/multiarray/alloc.c +++ b/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,13 @@ 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 until then, + * this a is a micro-optimization to avoid the `PyContextVar_Get` + */ + default_allocator_policy = 0; + PyObject *old_handler; PyObject *token; if (PyContextVar_Get(current_handler, NULL, &old_handler)) { @@ -560,6 +581,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; }