MAINT: elide the generic_wrapped_legacy_loop wrapper out of the hot loop

mattip · mattip · commit 154c293786a2 · 2022-12-28T00:09:15.000+02:00
diff --git 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
@@ -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);

Original file line number	Diff line number	Diff line change
`@@ -5944,15 +5944,15 @@ new_array_op(PyArrayObject op_array, char data)`
`5944`	`5944`	`}`
`5945`	`5945`
`5946`	`5946`	`int is_generic_wrapped_legacy_loop(PyArrayMethod_StridedLoop *strided_loop);`
	`5947`	`+PyUFuncGenericFunction get_inner_loop(void * auxdata);`
	`5948`	`+int get_pyerr_check(void * auxdata);`
`5947`	`5949`
`5948`	`5950`	`static int`
`5949`	`5951`	`ufunc_at__fast_iter(PyUFuncObject *ufunc, NPY_ARRAYMETHOD_FLAGS flags,`
`5950`	`5952`	`PyArrayMapIterObject iter, PyArrayIterObject iter2,`
`5951`	`5953`	`PyArrayObject op1_array, PyArrayObject op2_array,`
`5952`		`- PyArrayMethod_StridedLoop *strided_loop,`
`5953`		`- PyArrayMethod_Context *context,`
`5954`		`- npy_intp strides[3],`
`5955`		`- NpyAuxData *auxdata`
	`5954`	`+ PyUFuncGenericFunction loop, int pyerr_check,`
	`5955`	`+ npy_intp strides[3]`
`5956`	`5956`	`)`
`5957`	`5957`	`{`
`5958`	`5958`	`int buffersize;`
`@@ -5998,8 +5998,9 @@ ufunc_at__fast_iter(PyUFuncObject *ufunc, NPY_ARRAYMETHOD_FLAGS flags,`
`5998`	`5998`	`dataptr[2] = NULL;`
`5999`	`5999`	`}`
`6000`	`6000`
`6001`		`- res = strided_loop(context, dataptr, &count, strides, auxdata);`
`6002`		`- if (res != 0) {`
	`6001`	`+ loop(dataptr, &count, strides, NULL);`
	`6002`	`+ if (pyerr_check && PyErr_Occurred()) {`
	`6003`	`+ res = -1;`
`6003`	`6004`	`break;`
`6004`	`6005`	`}`
`6005`	`6006`
`@@ -6435,8 +6436,10 @@ ufunc_at(PyUFuncObject ufunc, PyObject args)`
`6435`	`6436`	`}`
`6436`	`6437`	`}`
`6437`	`6438`	`if (fast_path) {`
`6438`		`- res = ufunc_at__fast_iter(ufunc, flags, iter, iter2, op1_array, op2_array,`
`6439`		`- strided_loop, &context, strides, auxdata);`
	`6439`	`+ PyUFuncGenericFunction loop = get_inner_loop(auxdata);`
	`6440`	`+`
	`6441`	`+ res = ufunc_at__fast_iter(ufunc, flags, iter, iter2, op1_array,`
	`6442`	`+ op2_array, loop, get_pyerr_check(auxdata), strides);`
`6440`	`6443`	`} else {`
`6441`	`6444`	`res = ufunc_at__slow_iter(ufunc, flags, iter, iter2, op1_array, op2_array,`
`6442`	`6445`	`operation_descrs, strided_loop, &context, strides, auxdata);`