From 34e5e0243da71954dcdf6cb6ba4c7bffe4d23589 Mon Sep 17 00:00:00 2001 From: Pieter Eendebak Date: Mon, 30 May 2022 20:46:45 +0200 Subject: [PATCH 1/3] PERF: Fast path for returning default memory allocator --- doc/source/reference/c-api/data_memory.rst | 4 +++ numpy/core/src/multiarray/alloc.c | 30 +++++++++++++++++++--- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/doc/source/reference/c-api/data_memory.rst b/doc/source/reference/c-api/data_memory.rst index 2084ab5d083f..e8e95a9210bd 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,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() diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c index 8b765aa954da..08b2bbbec65a 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,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; } From 81aa86e90571cbd2ac88aaa9a99bba9bc843bea8 Mon Sep 17 00:00:00 2001 From: Pieter Eendebak Date: Mon, 30 May 2022 21:25:06 +0200 Subject: [PATCH 2/3] Update doc/source/reference/c-api/data_memory.rst Co-authored-by: Sebastian Berg --- doc/source/reference/c-api/data_memory.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/doc/source/reference/c-api/data_memory.rst b/doc/source/reference/c-api/data_memory.rst index e8e95a9210bd..5cf439912699 100644 --- a/doc/source/reference/c-api/data_memory.rst +++ b/doc/source/reference/c-api/data_memory.rst @@ -89,7 +89,8 @@ For details see: :ref:`NEP 49 — Data allocation strategies `. 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), + 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() From 129340f91766b5b636038263081bdaf804815c43 Mon Sep 17 00:00:00 2001 From: Pieter Eendebak Date: Mon, 30 May 2022 21:25:23 +0200 Subject: [PATCH 3/3] Update numpy/core/src/multiarray/alloc.c Co-authored-by: Sebastian Berg --- numpy/core/src/multiarray/alloc.c | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c index 08b2bbbec65a..dc1a8c37cb27 100644 --- a/numpy/core/src/multiarray/alloc.c +++ b/numpy/core/src/multiarray/alloc.c @@ -549,7 +549,11 @@ 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 + /* + * 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;