BUG: Use dict to store ufunc loops (fixes thread-safety) (numpy#31184)

Author: seberg

numpy/_core/code_generators/generate_umath.py

@@ -1589,7 +1589,7 @@ def make_ufuncs(funcdict):
             mlist.append(rf"((PyUFuncObject *)f)->type_resolver = &{uf.typereso};")
         for c in uf.indexed:
             # Handle indexed loops by getting the underlying ArrayMethodObject
-            # from the list in f._loops and setting its field appropriately
+            # from the dict in f._loops and setting its field appropriately
             fmt = textwrap.dedent("""
             {{
                 PyArray_DTypeMeta *dtype = PyArray_DTypeFromTypeNum({typenum});

numpy/_core/include/numpy/ufuncobject.h

@@ -223,7 +223,7 @@ typedef struct _tagPyUFuncObject {
     #if NPY_FEATURE_VERSION >= NPY_1_22_API_VERSION
         /* New private fields related to dispatching */
         void *_dispatch_cache;
-        /* A PyListObject of `(tuple of DTypes, ArrayMethod/Promoter)` */
+        /* Ordered dict `tuple of DTypes -> (tuple of DTypes, ArrayMethod/Promoter)` */
         PyObject *_loops;
     #endif
     #if NPY_FEATURE_VERSION >= NPY_2_1_API_VERSION

numpy/_core/src/umath/dispatching.cpp

@@ -6,8 +6,8 @@
  *
  * - operand_DTypes:  The datatypes as passed in by the user.
  * - signature: The DTypes fixed by the user with `dtype=` or `signature=`.
- * - ufunc._loops: A list of all ArrayMethods and promoters, it contains
- *   tuples `(dtypes, ArrayMethod)` or `(dtypes, promoter)`.
+ * - ufunc._loops: Ordered dict of all ArrayMethods and promoters, mapping
+ *   `dtypes` to tuples `(dtypes, ArrayMethod)` or `(dtypes, promoter)`.
  * - ufunc._dispatch_cache: A cache to store previous promotion and/or
  *   dispatching results.
  * - The actual arrays are used to support the old code paths where necessary.
@@ -70,8 +70,8 @@ promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
 
 
 /**
- * Function to add a new loop to the ufunc.  This mainly appends it to the
- * list (as it currently is just a list).
+ * Function to add a new loop to the ufunc.  This adds it to the
+ * _loops dict keyed by the DType tuple.
  *
  * @param ufunc The universal function to add the loop to.
  * @param info The tuple (dtype_tuple, ArrayMethod/promoter).
@@ -114,38 +114,16 @@ PyUFunc_AddLoop(PyUFuncObject *ufunc, PyObject *info, int ignore_duplicate)
         return -1;
     }
 
-    if (ufunc->_loops == NULL) {
-        ufunc->_loops = PyList_New(0);
-        if (ufunc->_loops == NULL) {
-            return -1;
-        }
+    int found = PyDict_SetDefaultRef(ufunc->_loops, DType_tuple, info, NULL);
+    if (found < 0) {
+        return -1;
     }
-
-    PyObject *loops = ufunc->_loops;
-    Py_ssize_t length = PyList_Size(loops);
-    for (Py_ssize_t i = 0; i < length; i++) {
-        PyObject *item = PyList_GetItemRef(loops, i);
-        PyObject *cur_DType_tuple = PyTuple_GetItem(item, 0);
-        Py_DECREF(item);
-        int cmp = PyObject_RichCompareBool(cur_DType_tuple, DType_tuple, Py_EQ);
-        if (cmp < 0) {
-            return -1;
-        }
-        if (cmp == 0) {
-            continue;
-        }
-        if (ignore_duplicate) {
-            return 0;
-        }
+    if (found && !ignore_duplicate) {
         PyErr_Format(PyExc_TypeError,
                 "A loop/promoter has already been registered with '%s' for %R",
                 ufunc_get_name_cstr(ufunc), DType_tuple);
         return -1;
     }
-
-    if (PyList_Append(loops, info) < 0) {
-        return -1;
-    }
     return 0;
 }
 
@@ -332,7 +310,13 @@ resolve_implementation_info(PyUFuncObject *ufunc,
         PyObject **out_info)
 {
     int nin = ufunc->nin, nargs = ufunc->nargs;
-    Py_ssize_t size = PySequence_Length(ufunc->_loops);
+    int ret = -1;
+    /* PyDict_Values returns a snapshot, safe against concurrent additions. */
+    PyObject *loops = PyDict_Values(ufunc->_loops);
+    if (loops == NULL) {
+        return -1;
+    }
+    Py_ssize_t size = PySequence_Length(loops);
     PyObject *best_dtypes = NULL;
     PyObject *best_resolver_info = NULL;
 
@@ -349,8 +333,7 @@ resolve_implementation_info(PyUFuncObject *ufunc,
 
     for (Py_ssize_t res_idx = 0; res_idx < size; res_idx++) {
         /* Test all resolvers  */
-        PyObject *resolver_info = PySequence_Fast_GET_ITEM(
-                ufunc->_loops, res_idx);
+        PyObject *resolver_info = PySequence_Fast_GET_ITEM(loops, res_idx);
 
         if (only_promoters && PyObject_TypeCheck(
                     PyTuple_GET_ITEM(resolver_info, 1), &PyArrayMethod_Type)) {
@@ -411,7 +394,7 @@ resolve_implementation_info(PyUFuncObject *ufunc,
             int subclass = PyObject_IsSubclass(
                     (PyObject *)given_dtype, (PyObject *)resolver_dtype);
             if (subclass < 0) {
-                return -1;
+                goto finish;
             }
             if (!subclass) {
                 matches = NPY_FALSE;
@@ -509,7 +492,7 @@ resolve_implementation_info(PyUFuncObject *ufunc,
                             "a better match is not yet implemented.  This "
                             "will pick the better (or bail) in the future.");
                     *out_info = NULL;
-                    return -1;
+                    goto finish;
                 }
 
                 if (best == -1) {
@@ -541,8 +524,9 @@ resolve_implementation_info(PyUFuncObject *ufunc,
                  * We just redo it anyway for simplicity.)
                  */
                 if (!only_promoters) {
-                    return resolve_implementation_info(ufunc,
-                            op_dtypes, NPY_TRUE, out_info);
+                    ret = resolve_implementation_info(
+                            ufunc, op_dtypes, NPY_TRUE, out_info);
+                    goto finish;
                 }
                 /*
                  * If this is already the retry, we are out of luck.  Promoters
@@ -564,7 +548,8 @@ resolve_implementation_info(PyUFuncObject *ufunc,
                     Py_DECREF(given);
                 }
                 *out_info = NULL;
-                return 0;
+                ret = 0;
+                goto finish;
             }
             else if (current_best == 0) {
                 /* The new match is not better, continue looking. */
@@ -578,11 +563,15 @@ resolve_implementation_info(PyUFuncObject *ufunc,
     if (best_dtypes == NULL) {
         /* The non-legacy lookup failed */
         *out_info = NULL;
-        return 0;
     }
+    else {
+        *out_info = best_resolver_info;
+    }
+    ret = 0;
 
-    *out_info = best_resolver_info;
-    return 0;
+finish:
+    Py_DECREF(loops);
+    return ret;
 }
 
 
@@ -816,7 +805,7 @@ legacy_promote_using_legacy_type_resolver(PyUFuncObject *ufunc,
 
 
 /*
- * Note, this function returns a BORROWED references to info since it adds
+ * Note, this function returns a BORROWED reference to info since it adds
  * it to the loops.
  */
 NPY_NO_EXPORT PyObject *
@@ -845,8 +834,11 @@ add_and_return_legacy_wrapping_ufunc_loop(PyUFuncObject *ufunc,
         Py_DECREF(info);
         return NULL;
     }
-    Py_DECREF(info);  /* now borrowed from the ufunc's list of loops */
-    return info;
+    /* Loop currently borrowed from the _loops (use original if not replaced) */
+    PyObject *result = PyDict_GetItemWithError(  // noqa: borrowed-ref OK
+        ufunc->_loops, PyTuple_GET_ITEM(info, 0));
+    Py_DECREF(info);
+    return result;
 }
 
 
@@ -1369,28 +1361,20 @@ get_info_no_cast(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtype,
         return NULL;
     }
     for (int i=0; i < ndtypes; i++) {
-        PyTuple_SetItem(t_dtypes, i, (PyObject *)op_dtype);
-    }
-    PyObject *loops = ufunc->_loops;
-    Py_ssize_t length = PyList_Size(loops);
-    for (Py_ssize_t i = 0; i < length; i++) {
-        PyObject *item = PyList_GetItemRef(loops, i);
-        PyObject *cur_DType_tuple = PyTuple_GetItem(item, 0);
-        Py_DECREF(item);
-        int cmp = PyObject_RichCompareBool(cur_DType_tuple,
-                                           t_dtypes, Py_EQ);
-        if (cmp < 0) {
-            Py_DECREF(t_dtypes);
-            return NULL;
-        }
-        if (cmp == 0) {
-            continue;
-        }
-        /* Got the match */
+        Py_INCREF(op_dtype);
+        PyTuple_SET_ITEM(t_dtypes, i, (PyObject *)op_dtype);
+    }
+    PyObject *info;
+    if (PyDict_GetItemRef(ufunc->_loops, t_dtypes, &info) < 0) {
         Py_DECREF(t_dtypes);
-        return PyTuple_GetItem(item, 1);
+        return NULL;
     }
     Py_DECREF(t_dtypes);
+    if (info != NULL) {
+        PyObject *result = PyTuple_GET_ITEM(info, 1);
+        Py_DECREF(info);
+        return result;
+    }
     Py_RETURN_NONE;
 }
 

numpy/_core/src/umath/ufunc_object.c

@@ -4921,7 +4921,7 @@ PyUFunc_FromFuncAndDataAndSignatureAndIdentity(PyUFuncGenericFunction *func, voi
          */
         ufunc->_dispatch_cache = NULL;
     }
-    ufunc->_loops = PyList_New(0);
+    ufunc->_loops = PyDict_New();
     if (ufunc->_loops == NULL) {
         Py_DECREF(ufunc);
         return NULL;
@@ -5249,21 +5249,18 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
      * A new-style loop should not be replaced by an old-style one.
      */
     int add_new_loop = 1;
-    for (Py_ssize_t j = 0; j < PyList_GET_SIZE(ufunc->_loops); j++) {
-        PyObject *item = PyList_GET_ITEM(ufunc->_loops, j); // noqa: borrowed-ref OK
-        PyObject *existing_tuple = PyTuple_GET_ITEM(item, 0);
-
-        int cmp = PyObject_RichCompareBool(existing_tuple, signature_tuple, Py_EQ);
-        if (cmp < 0) {
-            goto fail;
-        }
-        if (!cmp) {
-            continue;
-        }
-        PyObject *registered = PyTuple_GET_ITEM(item, 1);
-        if (!PyObject_TypeCheck(registered, &PyArrayMethod_Type) || (
-                (PyArrayMethodObject *)registered)->get_strided_loop !=
-                        &get_wrapped_legacy_ufunc_loop) {
+    PyObject *existing_item;
+    if (PyDict_GetItemRef(ufunc->_loops, signature_tuple, &existing_item) < 0) {
+        goto fail;
+    }
+    if (existing_item != NULL) {
+        PyObject *registered = PyTuple_GET_ITEM(existing_item, 1);
+        int not_compatible = (
+            !PyObject_TypeCheck(registered, &PyArrayMethod_Type) ||
+            ((PyArrayMethodObject *)registered)->get_strided_loop !=
+                &get_wrapped_legacy_ufunc_loop);
+        Py_DECREF(existing_item);
+        if (not_compatible) {
             PyErr_Format(PyExc_TypeError,
                     "A non-compatible loop was already registered for "
                     "ufunc %s and DTypes %S.",
@@ -5272,7 +5269,6 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
         }
         /* The loop was already added */
         add_new_loop = 0;
-        break;
     }
     if (add_new_loop) {
         PyObject *info = add_and_return_legacy_wrapping_ufunc_loop(

numpy/_core/src/umath/wrapping_array_method.c

@@ -235,6 +235,7 @@ PyUFunc_AddWrappingLoop(PyObject *ufunc_obj,
     PyObject *wrapped_dt_tuple = NULL;
     PyObject *new_dt_tuple = NULL;
     PyArrayMethodObject *meth = NULL;
+    PyObject *existing_info = NULL;
 
     if (!PyObject_TypeCheck(ufunc_obj, &PyUFunc_Type)) {
         PyErr_SetString(PyExc_TypeError,
@@ -249,28 +250,19 @@ PyUFunc_AddWrappingLoop(PyObject *ufunc_obj,
     }
 
     PyArrayMethodObject *wrapped_meth = NULL;
-    PyObject *loops = ufunc->_loops;
-    Py_ssize_t length = PyList_Size(loops);
-    for (Py_ssize_t i = 0; i < length; i++) {
-        PyObject *item = PyList_GetItemRef(loops, i);
-        PyObject *cur_DType_tuple = PyTuple_GetItem(item, 0);
-        Py_DECREF(item);
-        int cmp = PyObject_RichCompareBool(cur_DType_tuple, wrapped_dt_tuple, Py_EQ);
-        if (cmp < 0) {
-            goto finish;
-        }
-        if (cmp == 0) {
-            continue;
-        }
-        wrapped_meth = (PyArrayMethodObject *)PyTuple_GET_ITEM(item, 1);
-        if (!PyObject_TypeCheck(wrapped_meth, &PyArrayMethod_Type)) {
+    if (PyDict_GetItemRef(ufunc->_loops, wrapped_dt_tuple, &existing_info) < 0) {
+        goto finish;
+    }
+    if (existing_info != NULL) {
+        PyObject *existing_meth = PyTuple_GET_ITEM(existing_info, 1);
+        if (!PyObject_TypeCheck(existing_meth, &PyArrayMethod_Type)) {
             PyErr_SetString(PyExc_TypeError,
                     "Matching loop was not an ArrayMethod.");
             goto finish;
         }
-        break;
+        wrapped_meth = (PyArrayMethodObject *)existing_meth;
     }
-    if (wrapped_meth == NULL) {
+    else {
         PyErr_Format(PyExc_TypeError,
                 "Did not find the to-be-wrapped loop in the ufunc with given "
                 "DTypes. Received wrapping types: %S", wrapped_dt_tuple);
@@ -336,5 +328,6 @@ PyUFunc_AddWrappingLoop(PyObject *ufunc_obj,
     Py_XDECREF(wrapped_dt_tuple);
     Py_XDECREF(new_dt_tuple);
     Py_XDECREF(meth);
+    Py_XDECREF(existing_info);
     return res;
 }