MAINT: elide the generic_wrapped_legacy_loop wrapper out of the hot loop

2 files changed, 24 insertions, 8 deletions

diff --git a/numpy/core/src/umath/legacy_array_method.c b/numpy/core/src/umath/legacy_array_method.c
index 96cc7a2f1..a0b12f7f2 100644
--- a/numpy/core/src/umath/legacy_array_method.c
+++ b/numpy/core/src/umath/legacy_array_method.c
@@ -95,6 +95,19 @@ int
 is_generic_wrapped_legacy_loop(PyArrayMethod_StridedLoop *strided_loop) {
     return strided_loop == generic_wrapped_legacy_loop;
 }
+
+PyUFuncGenericFunction
+get_inner_loop(void * auxdata) {
+    legacy_array_method_auxdata *ldata = (legacy_array_method_auxdata *)auxdata;
+    return ldata->loop;
+}    
+
+int
+get_pyerr_check(void * auxdata) {
+    legacy_array_method_auxdata *ldata = (legacy_array_method_auxdata *)auxdata;
+    return ldata->pyerr_check;
+}    
+
 /*
  * Signal that the old type-resolution function must be used to resolve
  * the descriptors (mainly/only used for datetimes due to the unit).
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 664ff8855..e4d0653cd 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -5944,15 +5944,15 @@ new_array_op(PyArrayObject *op_array, char *data)
 }
 
 int is_generic_wrapped_legacy_loop(PyArrayMethod_StridedLoop *strided_loop);
+PyUFuncGenericFunction get_inner_loop(void * auxdata);
+int get_pyerr_check(void * auxdata);
 
 static int
 ufunc_at__fast_iter(PyUFuncObject *ufunc, NPY_ARRAYMETHOD_FLAGS flags,
                     PyArrayMapIterObject *iter, PyArrayIterObject *iter2,
                     PyArrayObject *op1_array, PyArrayObject *op2_array,
-                    PyArrayMethod_StridedLoop *strided_loop,
-                    PyArrayMethod_Context *context,
-                    npy_intp strides[3],
-                    NpyAuxData *auxdata
+                    PyUFuncGenericFunction loop, int pyerr_check,
+                    npy_intp strides[3]
                     )
 {
     int buffersize;
@@ -5998,8 +5998,9 @@ ufunc_at__fast_iter(PyUFuncObject *ufunc, NPY_ARRAYMETHOD_FLAGS flags,
             dataptr[2] = NULL;
         }
 
-        res = strided_loop(context, dataptr, &count, strides, auxdata);
-        if (res != 0) {
+        loop(dataptr, &count, strides, NULL);
+        if (pyerr_check && PyErr_Occurred()) {
+            res = -1;
             break;
         }
 
@@ -6435,8 +6436,10 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args)
         }
     }
     if (fast_path) {
-        res = ufunc_at__fast_iter(ufunc, flags, iter, iter2, op1_array, op2_array,
-                                  strided_loop, &context, strides, auxdata);
+        PyUFuncGenericFunction loop = get_inner_loop(auxdata);
+        
+        res = ufunc_at__fast_iter(ufunc, flags, iter, iter2, op1_array,
+                                  op2_array, loop, get_pyerr_check(auxdata), strides);
     } else {
         res = ufunc_at__slow_iter(ufunc, flags, iter, iter2, op1_array, op2_array,
                                   operation_descrs, strided_loop, &context, strides, auxdata);