MAINT, BUG: fixes from review and testing

8 files changed, 177 insertions, 81 deletions

diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
index ff32cf1b5..b45f89344 100644
--- a/numpy/core/code_generators/generate_umath.py
+++ b/numpy/core/code_generators/generate_umath.py
@@ -415,6 +415,7 @@ defdict = {
            TypeDescription('m', FullTypeDescr, 'mm', 'd', cfunc_alias='divide'),
           ],
           TD(O, f='PyNumber_TrueDivide'),
+          indexed=flts
           ),
 'conjugate':
     Ufunc(1, 1, None,
@@ -1298,15 +1299,11 @@ def make_ufuncs(funcdict):
                 /* info is borrowed, no need to decref*/
             }}
             """)
-            cname = name
-            if cname == "floor_divide":
-                # XXX there must be a better way...
-                cname = "divide"
             mlist.append(fmt.format(
                 typenum=f"NPY_{english_upper(chartoname[c])}",
                 count=uf.nin+uf.nout,
                 name=name,
-                funcname = f"{english_upper(chartoname[c])}_{cname}_indexed",
+                funcname = f"{english_upper(chartoname[c])}_{name}_indexed",
             ))
 
         mlist.append(r"""PyDict_SetItemString(dictionary, "%s", f);""" % name)
@@ -1335,44 +1332,10 @@ def make_code(funcdict, filename):
     #include "_umath_doc_generated.h"
 
     %s
-    /*
-     * Return the PyArrayMethodObject or PyCapsule that matches a registered
-     * tuple of identical dtypes. Return a borrowed ref of the first match.
-     */
-    static PyObject *
+    /* Returns a borrowed ref of the second value in the matching info tuple */
+    PyObject *
     get_info_no_cast(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtype,
-                     int ndtypes)
-    {
-
-        PyObject *t_dtypes = PyTuple_New(ndtypes);
-        if (t_dtypes == NULL) {
-            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_GetItem(loops, i);
-            PyObject *cur_DType_tuple = PyTuple_GetItem(item, 0);
-            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_DECREF(t_dtypes);
-            return PyTuple_GetItem(item, 1);
-        }
-        Py_DECREF(t_dtypes);
-        Py_RETURN_NONE;
-    }
-
+                     int ndtypes);
 
     static int
     InitOperators(PyObject *dictionary) {
diff --git a/numpy/core/src/umath/dispatching.c b/numpy/core/src/umath/dispatching.c
index e09568d69..fee77b2d6 100644
--- a/numpy/core/src/umath/dispatching.c
+++ b/numpy/core/src/umath/dispatching.c
@@ -1240,3 +1240,40 @@ install_logical_ufunc_promoter(PyObject *ufunc)
 
     return PyUFunc_AddLoop((PyUFuncObject *)ufunc, info, 0);
 }
+
+/*
+ * Return the PyArrayMethodObject or PyCapsule that matches a registered
+ * tuple of identical dtypes. Return a borrowed ref of the first match.
+ */
+NPY_NO_EXPORT PyObject *
+get_info_no_cast(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtype,
+                 int ndtypes)
+{
+    PyObject *t_dtypes = PyTuple_New(ndtypes);
+    if (t_dtypes == NULL) {
+        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_GetItem(loops, i);
+        PyObject *cur_DType_tuple = PyTuple_GetItem(item, 0);
+        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_DECREF(t_dtypes);
+        return PyTuple_GetItem(item, 1);
+    }
+    Py_DECREF(t_dtypes);
+    Py_RETURN_NONE;
+}
diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src
index cd5a431cd..87ef2fca9 100644
--- a/numpy/core/src/umath/loops.c.src
+++ b/numpy/core/src/umath/loops.c.src
@@ -460,6 +460,7 @@ BOOL_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps, void
  */
 
 #define @TYPE@_floor_divide @TYPE@_divide
+#define @TYPE@_floor_divide_indexed @TYPE@_divide_indexed
 #define @TYPE@_fmax @TYPE@_maximum
 #define @TYPE@_fmin @TYPE@_minimum
 
@@ -531,7 +532,8 @@ NPY_NO_EXPORT NPY_GCC_OPT_3 @ATTR@ void
 
 #if @CHK@
 NPY_NO_EXPORT NPY_GCC_OPT_3 @ATTR@ void
-@TYPE@_@kind@@isa@(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+@TYPE@_@kind@@isa@(char **args, npy_intp const *dimensions,
+                   npy_intp const *steps, void *NPY_UNUSED(func))
 {
     if (IS_BINARY_REDUCE) {
         BINARY_REDUCE_LOOP_FAST(@type@, io1 @OP@= in2);
@@ -542,17 +544,19 @@ NPY_NO_EXPORT NPY_GCC_OPT_3 @ATTR@ void
 }
 
 NPY_NO_EXPORT NPY_GCC_OPT_3 @ATTR@ int
-@TYPE@_@kind@@isa@_indexed(PyArrayMethod_Context *NPY_UNUSED(context), char * const*args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
+@TYPE@_@kind@@isa@_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
+                           char * const*args, npy_intp const *dimensions,
+                           npy_intp const *steps, NpyAuxData *NPY_UNUSED(func))
 {
     char *ip1 = args[0];
-    npy_intp *indx = (npy_intp *)args[1];
+    char *indx = args[1];
     char *value = args[2];
-    npy_intp is1 = steps[0], is2 = steps[1] / sizeof(npy_intp), os1 = steps[2];
+    npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
     npy_intp n = dimensions[0];
     npy_intp i;
     @type@ *indexed;
-    for(i = 0; i < n; i++, indx += is2, value += os1) {
-        indexed = (@type@ *)(ip1 + is1 * indx[0]);
+    for(i = 0; i < n; i++, indx += isindex, value += isb) {
+        indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx);
         indexed[0] = indexed[0] @OP@ *(@type@ *)value;
     }
     return 0;
@@ -1418,6 +1422,25 @@ NPY_NO_EXPORT void
     }
 }
 
+NPY_NO_EXPORT int
+@TYPE@_floor_divide_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
+         char **args, npy_intp const *dimensions, npy_intp const *steps,
+         void *NPY_UNUSED(func))
+{
+    char *ip1 = args[0];
+    char *indx = args[1];
+    char *value = args[2];
+    npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
+    npy_intp n = dimensions[0];
+    npy_intp i;
+    @type@ *indexed;
+    for(i = 0; i < n; i++, indx += isindex, value += isb) {
+        indexed = (@type@ *)(ip1 + is1 * *(npy_intp *)indx);
+        indexed[0] = npy_floor_divide@c@(indexed[0], *(@type@ *)value);
+    }
+    return 0;
+}
+
 NPY_NO_EXPORT void
 @TYPE@_remainder(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
 {
@@ -1553,20 +1576,23 @@ LONGDOUBLE_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps
     }
 }
 
-NPY_NO_EXPORT void
-LONGDOUBLE_@kind@_indexed(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+NPY_NO_EXPORT int
+LONGDOUBLE_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
+         char **args, npy_intp const *dimensions, npy_intp const *steps,
+         void *NPY_UNUSED(func))
 {
     char *ip1 = args[0];
-    npy_intp *indx = (npy_intp *)args[1];
+    char *indx = args[1];
     char *value = args[2];
-    npy_intp is1 = steps[0], is2 = steps[1] / sizeof(npy_intp), os1 = steps[2];
+    npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
     npy_intp n = dimensions[0];
     npy_intp i;
     npy_longdouble *indexed;
-    for(i = 0; i < n; i++, indx += is2, value += os1) {
-        indexed = (npy_longdouble *)(ip1 + is1 * indx[0]);
+    for(i = 0; i < n; i++, indx += isindex, value += isb) {
+        indexed = (npy_longdouble *)(ip1 + is1 * *(npy_intp *)indx);
         indexed[0] = indexed[0] @OP@ *(npy_longdouble *)value;
     }
+    return 0;
 }
 
 /**end repeat**/
@@ -1679,14 +1705,14 @@ NPY_NO_EXPORT int
 HALF_@kind@_indexed(void *NPY_UNUSED(context), char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
 {
     char *ip1 = args[0];
-    npy_intp *indx = (npy_intp *)args[1];
+    char *indx = args[1];
     char *value = args[2];
-    npy_intp is1 = steps[0], is2 = steps[1] / sizeof(npy_intp), os1 = steps[2];
+    npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
     npy_intp n = dimensions[0];
     npy_intp i;
     npy_half *indexed;
-    for(i = 0; i < n; i++, indx += is2, value += os1) {
-        indexed = (npy_half *)(ip1 + is1 * indx[0]);
+    for(i = 0; i < n; i++, indx += isindex, value += isb) {
+        indexed = (npy_half *)(ip1 + is1 * *(npy_intp *)indx);
         const float v = npy_half_to_float(*(npy_half *)value);
         indexed[0] = npy_float_to_half(npy_half_to_float(indexed[0]) @OP@ v);
     }
@@ -1828,6 +1854,27 @@ HALF_floor_divide(char **args, npy_intp const *dimensions, npy_intp const *steps
     }
 }
 
+NPY_NO_EXPORT int
+HALF_floor_divide_indexed(PyArrayMethod_Context *NPY_UNUSED(context),
+         char **args, npy_intp const *dimensions, npy_intp const *steps,
+         void *NPY_UNUSED(func))
+{
+    char *ip1 = args[0];
+    char *indx = args[1];
+    char *value = args[2];
+    npy_intp is1 = steps[0], isindex = steps[1], isb = steps[2];
+    npy_intp n = dimensions[0];
+    npy_intp i;
+    npy_half *indexed;
+    for(i = 0; i < n; i++, indx += isindex, value += isb) {
+        indexed = (npy_half *)(ip1 + is1 * *(npy_intp *)indx);
+        float v = npy_half_to_float(*(npy_half *)value);
+        float div = npy_floor_dividef(npy_half_to_float(indexed[0]), v);
+        indexed[0] = npy_float_to_half(div);
+    }
+    return 0;
+}
+
 NPY_NO_EXPORT void
 HALF_remainder(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
 {
diff --git a/numpy/core/src/umath/loops.h.src b/numpy/core/src/umath/loops.h.src
index ee69330dd..b25f6770f 100644
--- a/numpy/core/src/umath/loops.h.src
+++ b/numpy/core/src/umath/loops.h.src
@@ -134,6 +134,7 @@ NPY_CPU_DISPATCH_DECLARE(NPY_NO_EXPORT void @TYPE@_@kind@,
  */
 
 #define @S@@TYPE@_floor_divide @S@@TYPE@_divide
+#define @S@@TYPE@_floor_divide_indexed @S@@TYPE@_divide_indexed
 #define @S@@TYPE@_fmax @S@@TYPE@_maximum
 #define @S@@TYPE@_fmin @S@@TYPE@_minimum
 
@@ -491,6 +492,9 @@ NPY_NO_EXPORT void
 NPY_NO_EXPORT void
 @TYPE@_floor_divide(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func));
 
+NPY_NO_EXPORT int
+@TYPE@_floor_divide_indexed(PyArrayMethod_Context *NPY_UNUSED(context), char *const *args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func));
+
 NPY_NO_EXPORT void
 @TYPE@_remainder(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func));
 
diff --git a/numpy/core/src/umath/loops_arithm_fp.dispatch.c.src b/numpy/core/src/umath/loops_arithm_fp.dispatch.c.src
index d37facef8..a00ef1e4b 100644
--- a/numpy/core/src/umath/loops_arithm_fp.dispatch.c.src
+++ b/numpy/core/src/umath/loops_arithm_fp.dispatch.c.src
@@ -553,11 +553,11 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@_indexed)
     char *ip1 = args[0];
     npy_intp *indx = (npy_intp *)args[1];
     char *value = args[2];
-    npy_intp is1 = steps[0], is2 = steps[1] / sizeof(npy_intp), os1 = steps[2];
+    npy_intp is1 = steps[0], isindex = steps[1] / sizeof(npy_intp), isb = steps[2];
     npy_intp n = dimensions[0];
     npy_intp i;
     @type@ *indexed;
-    for(i = 0; i < n; i++, indx += is2, value += os1) {
+    for(i = 0; i < n; i++, indx += isindex, value += isb) {
         indexed = (@type@ *)(ip1 + is1 * indx[0]);
         indexed[0] = indexed[0] @OP@ *(@type@ *)value;
     }
diff --git a/numpy/core/src/umath/loops_arithmetic.dispatch.c.src b/numpy/core/src/umath/loops_arithmetic.dispatch.c.src
index cef2adb34..9f3043f48 100644
--- a/numpy/core/src/umath/loops_arithmetic.dispatch.c.src
+++ b/numpy/core/src/umath/loops_arithmetic.dispatch.c.src
@@ -402,11 +402,11 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_divide_indexed)
     char *ip1 = args[0];
     npy_intp *indx = (npy_intp *)args[1];
     char *value = args[2];
-    npy_intp is1 = steps[0], is2 = steps[1] / sizeof(npy_intp), os1 = steps[2];
+    npy_intp is1 = steps[0], isindex = steps[1] / sizeof(npy_intp), isb = steps[2];
     npy_intp n = dimensions[0];
     npy_intp i;
     @type@ *indexed;
-    for(i = 0; i < n; i++, indx += is2, value += os1) {
+    for(i = 0; i < n; i++, indx += isindex, value += isb) {
         indexed = (@type@ *)(ip1 + is1 * indx[0]);
         indexed[0] = floor_div_@TYPE@(indexed[0], *(@type@ *)value);
     }
@@ -488,11 +488,11 @@ NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_divide_indexed)
     char *ip1 = args[0];
     npy_intp *indx = (npy_intp *)args[1];
     char *value = args[2];
-    npy_intp is1 = steps[0], is2 = steps[1] / sizeof(npy_intp), os1 = steps[2];
+    npy_intp is1 = steps[0], isindex = steps[1] / sizeof(npy_intp), isb = steps[2];
     npy_intp n = dimensions[0];
     npy_intp i;
     @type@ *indexed;
-    for(i = 0; i < n; i++, indx += is2, value += os1) {
+    for(i = 0; i < n; i++, indx += isindex, value += isb) {
         indexed = (@type@ *)(ip1 + is1 * indx[0]);
         @type@ in2 = *(@type@ *)value;
         if (NPY_UNLIKELY(in2 == 0)) {
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index d088e22de..5f199129f 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -5894,26 +5894,22 @@ new_array_op(PyArrayObject *op_array, char *data)
 }
 
 /*
- * If operands are 1D we can use the indexed loop to do the work
- * Returns 0 if successful, -1 otherwise
+ * Use an indexed loop to do the work
+ * Returns 0 if successful
  */
 static int
-try_trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
+trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
                     PyArrayMapIterObject *iter,
                     PyArrayObject *op1_array, PyArrayObject *op2_array,
                     PyArrayMethod_Context *context)
 {
-    if (PyArray_NDIM(op1_array) != 1) {
-        return -1;
-    }
-    /* check that nop == 3 */
-    if (ufuncimpl->nin >1 && PyArray_NDIM(op2_array) != 1) {
-        return -1;
-    }
-    if (iter->numiter > 1) {
+    int buffersize=0, errormask = 0;
+    int res;
+    char *args[3];
+    const char * ufunc_name = ufunc_get_name_cstr((PyUFuncObject *)context->caller);
+    if (_get_bufsize_errmask(NULL, ufunc_name, &buffersize, &errormask) < 0) {
         return -1;
     }
-    char *args[3];
     npy_intp dimensions = iter->size;
     npy_intp steps[3];
     args[0] = (char *) iter->baseoffset;
@@ -5922,8 +5918,20 @@ try_trivial_at_loop(PyArrayMethodObject *ufuncimpl, NPY_ARRAYMETHOD_FLAGS flags,
     steps[0] = iter->fancy_strides[0];
     steps[1] = iter->outer_strides[0];
     steps[2] = PyArray_STRIDE(op2_array, 0);
-    ufuncimpl->contiguous_indexed_loop(context, args, &dimensions, steps, NULL);
-    return 0;
+    res = ufuncimpl->contiguous_indexed_loop(context, args, &dimensions, steps, NULL);
+    /*
+     * An error should only be possible if `res != 0` is already set.
+     * But this is not strictly correct since old-style ufuncs (e.g. `power`
+     * released the GIL but manually set an Exception).
+     */
+    if (PyErr_Occurred()) {
+        res = -1;
+    }
+
+    if (res == 0 && !(flags & NPY_METH_NO_FLOATINGPOINT_ERRORS)) {
+        res = _check_ufunc_fperr(errormask, NULL, ufunc_name);
+    }
+    return res;
 }
 
 static int
@@ -6421,12 +6429,15 @@ ufunc_at(PyUFuncObject *ufunc, PyObject *args)
          * - the matching info has a indexed loop
          * - all operands are 1d
          */
-        if (ufuncimpl->contiguous_indexed_loop != NULL) {
-            res = try_trivial_at_loop(ufuncimpl, flags, iter, op1_array,
+        if ((ufuncimpl->contiguous_indexed_loop != NULL) &&
+                (PyArray_NDIM(op1_array) == 1)  &&
+                (ufuncimpl->nin == 1 || PyArray_NDIM(op2_array) == 1) && 
+                (iter->numiter == 1)) {
+            res = trivial_at_loop(ufuncimpl, flags, iter, op1_array,
                         op2_array, &context);
 
         }
-        if (res == -1) {
+        else {
             /* Couldn't use the fastest path, use the faster path */
             res = ufunc_at__fast_iter(ufunc, flags, iter, iter2, op1_array,
                         op2_array, strided_loop, &context, auxdata);
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index 55122d697..22370bbd8 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -1986,6 +1986,40 @@ class TestUfunc:
         with pytest.raises(ValueError):
             # second operand cannot be converted to an array
             np.add.at(a, [2, 5, 3], [[1, 2], 1])
+    
+    # ufuncs with indexed loops for perfomance in ufunc.at
+    indexed_ufuncs = [np.add, np.subtract, np.multiply, np.floor_divide, np.divide]
+    @pytest.mark.parametrize(
+                "typecode", np.typecodes['AllInteger'] + np.typecodes['Float'])
+    @pytest.mark.parametrize("ufunc", indexed_ufuncs)
+    def test_ufunc_at_inner_loops(self, typecode, ufunc):
+        if ufunc is np.divide and typecode in  np.typecodes['AllInteger']:
+            # Avoid divide-by-zero and inf for integer divide
+            a = np.ones(100, dtype=typecode)
+            indx = np.random.randint(100, size=30, dtype=np.intp)
+            vals = np.arange(1, 31, dtype=typecode)
+        else:
+            a = np.ones(1000, dtype=typecode)
+            indx = np.random.randint(1000, size=3000, dtype=np.intp)
+            vals = np.arange(3000, dtype=typecode)
+        atag = a.copy()
+        # Do the calculation twice and compare the answers
+        with warnings.catch_warnings(record=True) as w_at:
+            warnings.simplefilter('always')
+            ufunc.at(a, indx, vals)
+        with warnings.catch_warnings(record=True) as w_loop:
+            warnings.simplefilter('always')
+            for i,v in zip(indx, vals):
+                # Make sure all the work happens inside the ufunc
+                # in order to duplicate error/warning handling
+                ufunc(atag[i], v, out=atag[i:i+1], casting="unsafe")
+        assert_equal(atag, a)
+        # If w_loop warned, make sure w_at warned as well
+        if len(w_loop) > 0:
+            #
+            assert len(w_at) > 0
+            assert w_at[0].category == w_loop[0].category
+            assert str(w_at[0].message)[:10] == str(w_loop[0].message)[:10]
 
     def test_ufunc_at_multiD(self):
         a = np.arange(9).reshape(3, 3)