Implement allocator context-locality

numpy · seberg · Oct 25, 2021 · Oct 11, 2020 · Oct 16, 2020 · Oct 18, 2020
commit 7291484f981dde634524656e2c2888ca2d20be6d
diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
@@ -351,7 +351,7 @@
     'PyArray_SetWritebackIfCopyBase':       (303,),
     # End 1.14 API
     'PyDataMem_SetHandler':                 (304,),
-    'PyDataMem_GetHandlerName':             (305,),
+    'PyDataMem_GetHandler':                 (305,),
     # End 1.21 API
 }
 

diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c
@@ -386,7 +386,7 @@ PyDataMem_Handler default_handler = {
     }
 };
 
-PyDataMem_Handler *current_handler = &default_handler;
+PyObject *current_handler;
 
 int uo_index=0;   /* user_override index */
 
@@ -470,50 +470,77 @@ PyDataMem_UserRENEW(void *ptr, size_t size, const PyDataMemAllocator *allocator)
 }
 
 /*NUMPY_API
- * Sets a new allocation policy. If the input value is NULL, will reset
- * the policy to the default. Returns the previous policy. We wrap
- * the user-provided functions so they will still call the python
- * and numpy memory management callback hooks.
+ * Set a new allocation policy. If the input value is NULL, will reset
+ * the policy to the default. Return the previous policy, or set an exception
+ * and 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.
  */
 NPY_NO_EXPORT const PyDataMem_Handler *
 PyDataMem_SetHandler(PyDataMem_Handler *handler)
 {
-    const PyDataMem_Handler *old = current_handler;
+    PyObject *capsule;
+    PyObject *old_capsule;
+    PyDataMem_Handler *old_handler;
+    PyObject *token;
+    if (PyContextVar_Get(current_handler, NULL, &old_capsule)) {
+        return NULL;
+    }
+    old_handler = (PyDataMem_Handler *) PyCapsule_GetPointer(old_capsule, NULL);
+    if (old_handler == NULL) {
+        return NULL;
+    }
     if (handler) {
-        current_handler = handler;
+        capsule = PyCapsule_New(handler, NULL, NULL);
+        if (capsule == NULL) {
+            return NULL;
+        }
     }
     else {
-        current_handler = &default_handler;
+        capsule = PyCapsule_New(&default_handler, NULL, NULL);
+        if (capsule == NULL) {
+            return NULL;
+        }
     }
-    return old;
+    token = PyContextVar_Set(current_handler, capsule);
+    if (token == NULL) {
+        Py_DECREF(capsule);
+        return NULL;
+    }
+    Py_DECREF(old_capsule);
+    Py_DECREF(token);
+    return old_handler;
 }
 
 /*NUMPY_API
- * Return the const char name of the PyDataMem_Handler used by the
- * PyArrayObject or its base. If neither the PyArrayObject owns its own data
- * nor its base is a PyArrayObject which owns its own data return an empty string.
- * If NULL, return the name of the current global policy that
- * will be used to allocate data for the next PyArrayObject.
+ * Return the PyDataMem_Handler used by the PyArrayObject. If NULL, return
+ * the current global policy that ill be used to allocate data
+ * for the next PyArrayObject. On failure, set an exception and return NULL.
  */
-NPY_NO_EXPORT const char *
-PyDataMem_GetHandlerName(PyArrayObject *obj)
+NPY_NO_EXPORT PyDataMem_Handler *
+PyDataMem_GetHandler(PyArrayObject *obj)
 {
+    PyObject *base;
+    PyObject *capsule;
+    PyDataMem_Handler *handler;
     if (obj == NULL) {
-        return current_handler->name;
+        if (PyContextVar_Get(current_handler, NULL, &capsule)) {
+            return NULL;
+        }
+        return (PyDataMem_Handler *) PyCapsule_GetPointer(capsule, NULL);
     }
-    PyDataMem_Handler *handler;
+    /* If there's a handler, the array owns its own datay */
     handler = PyArray_HANDLER(obj);
-    if (handler != NULL) {
-        return handler->name;
-    }
-    PyObject *base = PyArray_BASE(obj);
-    if (base != NULL && PyArray_Check(base)) {
-        handler = PyArray_HANDLER((PyArrayObject *) base);
-        if (handler != NULL) {
-             return handler->name;
+    if (handler == NULL) {
+        /*
+         * If the base is an array which owns its own data, return its allocator.
+         */
+        base = PyArray_BASE(obj);
+        if (base != NULL && PyArray_Check(base) && PyArray_CHKFLAGS(base, NPY_ARRAY_OWNDATA)) {
+            return PyArray_HANDLER(base);
         }
     }
-    return "";
+    return handler;
 }
 
 NPY_NO_EXPORT PyObject *
@@ -527,12 +554,9 @@ get_handler_name(PyObject *NPY_UNUSED(self), PyObject *args)
          PyErr_SetString(PyExc_ValueError, "if supplied, argument must be an ndarray");
          return NULL;
     }
-    const char * name = PyDataMem_GetHandlerName((PyArrayObject *)arr);
-    if (name == NULL) {
+    const PyDataMem_Handler * mem_handler = PyDataMem_GetHandler((PyArrayObject *)arr);
+    if (mem_handler == NULL) {
         return NULL;
     }
-    else if (strlen(name) == 0) {
-        Py_RETURN_NONE;
-    }
-    return PyUnicode_FromString(name);
+    return PyUnicode_FromString(mem_handler->name);
 }
diff --git a/numpy/core/src/multiarray/alloc.h b/numpy/core/src/multiarray/alloc.h
@@ -39,7 +39,7 @@ npy_free_cache_dim_array(PyArrayObject * arr)
     npy_free_cache_dim(PyArray_DIMS(arr), PyArray_NDIM(arr));
 }
 
-extern PyDataMem_Handler *current_handler;
+extern PyObject *current_handler; /* PyContextVar/PyCapsule */
 extern PyDataMem_Handler default_handler;
 
 NPY_NO_EXPORT PyObject *

diff --git a/numpy/core/src/multiarray/arraytypes.c.src b/numpy/core/src/multiarray/arraytypes.c.src
@@ -3082,6 +3082,7 @@ UNICODE_compare(npy_ucs4 *ip1, npy_ucs4 *ip2,
 static int
 VOID_compare(char *ip1, char *ip2, PyArrayObject *ap)
 {
+    PyDataMem_Handler *mem_handler;
     PyArray_Descr *descr;
     PyObject *names, *key;
     PyObject *tup;
@@ -3093,6 +3094,12 @@ VOID_compare(char *ip1, char *ip2, PyArrayObject *ap)
     if (!PyArray_HASFIELDS(ap)) {
         return STRING_compare(ip1, ip2, ap);
     }
+    mem_handler = PyDataMem_GetHandler(NULL);
+    if (mem_handler == NULL) {
+        /* Better fallback to default than goto finish. */
+        PyErr_Clear();
+        mem_handler = &default_handler;
+    }
     descr = PyArray_DESCR(ap);
     /*
      * Compare on the first-field.  If equal, then
@@ -3119,7 +3126,7 @@ VOID_compare(char *ip1, char *ip2, PyArrayObject *ap)
                  * create temporary buffer and copy,
                  * always use the current handler for internal allocations
                  */
-                nip1 = PyDataMem_UserNEW(new->elsize, &current_handler->allocator);
+                nip1 = PyDataMem_UserNEW(new->elsize, &mem_handler->allocator);
                 if (nip1 == NULL) {
                     goto finish;
                 }
@@ -3132,12 +3139,12 @@ VOID_compare(char *ip1, char *ip2, PyArrayObject *ap)
                  * create temporary buffer and copy,
                  * always use the current handler for internal allocations
                  */
-                nip2 = PyDataMem_UserNEW(new->elsize, &current_handler->allocator);
+                nip2 = PyDataMem_UserNEW(new->elsize, &mem_handler->allocator);
                 if (nip2 == NULL) {
                     if (nip1 != ip1 + offset) {
                         /* destroy temporary buffer */
                         PyDataMem_UserFREE(nip1, new->elsize,
-                                           &current_handler->allocator);
+                                           &mem_handler->allocator);
                     }
                     goto finish;
                 }
@@ -3150,11 +3157,11 @@ VOID_compare(char *ip1, char *ip2, PyArrayObject *ap)
         if (swap || new->alignment > 1) {
             if (nip1 != ip1 + offset) {
                 /* destroy temporary buffer */
-                PyDataMem_UserFREE(nip1, new->elsize, &current_handler->allocator);
+                PyDataMem_UserFREE(nip1, new->elsize, &mem_handler->allocator);
             }
             if (nip2 != ip2 + offset) {
                 /* destroy temporary buffer */
-                PyDataMem_UserFREE(nip2, new->elsize, &current_handler->allocator);
+                PyDataMem_UserFREE(nip2, new->elsize, &mem_handler->allocator);
             }
         }
         if (res != 0) {

diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
@@ -807,7 +807,10 @@ PyArray_NewFromDescr_int(
 
     if (data == NULL) {
         /* Store the functions in case the global handler is modified */
-        fa->mem_handler = current_handler;
+        fa->mem_handler = PyDataMem_GetHandler(NULL);
+        if (fa->mem_handler == NULL) {
+            goto fail;
+        }
         /*
          * Allocate something even for zero-space arrays
          * e.g. shape=(0,) -- otherwise buffer exposure

diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
@@ -951,6 +951,7 @@ _new_sortlike(PyArrayObject *op, int axis, PyArray_SortFunc *sort,
     PyArray_CopySwapNFunc *copyswapn = PyArray_DESCR(op)->f->copyswapn;
     char *buffer = NULL;
 
+    PyDataMem_Handler *mem_handler;
     PyArrayIterObject *it;
     npy_intp size;
 
@@ -963,14 +964,18 @@ _new_sortlike(PyArrayObject *op, int axis, PyArray_SortFunc *sort,
         return 0;
     }
 
+    mem_handler = PyDataMem_GetHandler(NULL);
+    if (mem_handler == NULL) {
+        return -1;
+    }
     it = (PyArrayIterObject *)PyArray_IterAllButAxis((PyObject *)op, &axis);
     if (it == NULL) {
         return -1;
     }
     size = it->size;
 
     if (needcopy) {
-        buffer = PyDataMem_UserNEW(N * elsize, &current_handler->allocator);
+        buffer = PyDataMem_UserNEW(N * elsize, &mem_handler->allocator);
         if (buffer == NULL) {
             ret = -1;
             goto fail;
@@ -1055,7 +1060,7 @@ _new_sortlike(PyArrayObject *op, int axis, PyArray_SortFunc *sort,
 fail:
     NPY_END_THREADS_DESCR(PyArray_DESCR(op));
     /* cleanup internal buffer */
-    PyDataMem_UserFREE(buffer, N * elsize, &current_handler->allocator);
+    PyDataMem_UserFREE(buffer, N * elsize, &mem_handler->allocator);
     if (ret < 0 && !PyErr_Occurred()) {
         /* Out of memory during sorting or buffer creation */
         PyErr_NoMemory();
@@ -1082,6 +1087,7 @@ _new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort,
     char *valbuffer = NULL;
     npy_intp *idxbuffer = NULL;
 
+    PyDataMem_Handler *mem_handler;
     PyArrayObject *rop;
     npy_intp rstride;
 
@@ -1092,6 +1098,10 @@ _new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort,
 
     NPY_BEGIN_THREADS_DEF;
 
+    mem_handler = PyDataMem_GetHandler(NULL);
+    if (mem_handler == NULL) {
+        return NULL;
+    }
     rop = (PyArrayObject *)PyArray_NewFromDescr(
             Py_TYPE(op), PyArray_DescrFromType(NPY_INTP),
             PyArray_NDIM(op), PyArray_DIMS(op), NULL, NULL,
@@ -1117,7 +1127,7 @@ _new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort,
     size = it->size;
 
     if (needcopy) {
-        valbuffer = PyDataMem_UserNEW(N * elsize, &current_handler->allocator);
+        valbuffer = PyDataMem_UserNEW(N * elsize, &mem_handler->allocator);
         if (valbuffer == NULL) {
             ret = -1;
             goto fail;
@@ -1126,7 +1136,7 @@ _new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort,
 
     if (needidxbuffer) {
         idxbuffer = (npy_intp *)PyDataMem_UserNEW(N * sizeof(npy_intp),
-                                                  &current_handler->allocator);
+                                                  &mem_handler->allocator);
         if (idxbuffer == NULL) {
             ret = -1;
             goto fail;
@@ -1216,8 +1226,8 @@ _new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort,
 fail:
     NPY_END_THREADS_DESCR(PyArray_DESCR(op));
     /* cleanup internal buffers */
-    PyDataMem_UserFREE(valbuffer, N * elsize, &current_handler->allocator);
-    PyDataMem_UserFREE(idxbuffer, N * sizeof(npy_intp), &current_handler->allocator);
+    PyDataMem_UserFREE(valbuffer, N * elsize, &mem_handler->allocator);
+    PyDataMem_UserFREE(idxbuffer, N * sizeof(npy_intp), &mem_handler->allocator);
     if (ret < 0) {
         if (!PyErr_Occurred()) {
             /* Out of memory during sorting or buffer creation */

diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
@@ -2099,7 +2099,11 @@ array_setstate(PyArrayObject *self, PyObject *args)
                 Py_RETURN_NONE;
             }
             /* Store the functions in case the global handler is modified */
-            fa->mem_handler = current_handler;
+            fa->mem_handler = PyDataMem_GetHandler(NULL);
+            if (fa->mem_handler == NULL) {
+                Py_DECREF(rawdata);
+                return NULL;
+            }
             fa->data = PyDataMem_UserNEW(num, &PyArray_HANDLER(fa)->allocator);
             if (PyArray_DATA(self) == NULL) {
                 Py_DECREF(rawdata);
@@ -2153,7 +2157,10 @@ array_setstate(PyArrayObject *self, PyObject *args)
             Py_RETURN_NONE;
         }
         /* Store the functions in case the global handler is modified */
-        fa->mem_handler = current_handler;
+        fa->mem_handler = PyDataMem_GetHandler(NULL);
+        if (fa->mem_handler == NULL) {
+            return NULL;
+        }
         fa->data = PyDataMem_UserNEW(num, &PyArray_HANDLER(fa)->allocator);
         if (PyArray_DATA(self) == NULL) {
             return PyErr_NoMemory();

diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
@@ -4908,6 +4908,17 @@ PyMODINIT_FUNC PyInit__multiarray_umath(void) {
     if (initumath(m) != 0) {
         goto err;
     }
+    /*
+     * Initialize the context-local PyDataMem_Handler capsule.
+     */
+    c_api = PyCapsule_New(&default_handler, NULL, NULL);
+    if (c_api == NULL) {
+        goto err;
+    }
+    current_handler = PyContextVar_New("current_allocator", c_api);
+    if (current_handler == NULL) {
+        goto err;
+    }
     return m;
 
  err:

diff --git a/numpy/core/tests/test_mem_policy.py b/numpy/core/tests/test_mem_policy.py
@@ -18,7 +18,7 @@ def get_module(tmp_path):
                         "must be called with a numpy scalar or ndarray");
             }
             return PyUnicode_FromString(
-                            PyDataMem_GetHandlerName((PyArrayObject*)args));
+                            PyDataMem_GetHandler((PyArrayObject*)args)->name);
         """
         ),
         ("set_new_policy", "METH_NOARGS",