ENH: add indexed loops for maximum, minimum, fmax, fmin (#23177)

Continuation of the ufunc.at optimizations:

add indexed loops for maximum, minimum, fmax, fmin ufuncs and a benchmark for maximum.at (performance increased by ~13x)

* BENCH: add np.maximum.at benchmark

* remove 'explain_chain'

* add a seed to `default_rng() (from review)

* MAINT: formatting and change comments, from review

5 files changed, 96 insertions, 16 deletions

diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
index 86d47b8a9..c7f4e9aee 100644
--- a/numpy/core/code_generators/generate_umath.py
+++ b/numpy/core/code_generators/generate_umath.py
@@ -599,7 +599,8 @@ defdict = {
           'PyUFunc_SimpleUniformOperationTypeResolver',
           TD('?', cfunc_alias='logical_or', dispatch=[('loops_logical', '?')]),
           TD(no_obj_bool, dispatch=[('loops_minmax', ints+'fdg')]),
-          TD(O, f='npy_ObjectMax')
+          TD(O, f='npy_ObjectMax'),
+          indexed=flts + ints,
           ),
 'minimum':
     Ufunc(2, 1, ReorderableNone,
@@ -608,7 +609,8 @@ defdict = {
           TD('?', cfunc_alias='logical_and',
                   dispatch=[('loops_logical', '?')]),
           TD(no_obj_bool, dispatch=[('loops_minmax', ints+'fdg')]),
-          TD(O, f='npy_ObjectMin')
+          TD(O, f='npy_ObjectMin'),
+          indexed=flts + ints,
           ),
 'clip':
     Ufunc(3, 1, ReorderableNone,
@@ -623,7 +625,8 @@ defdict = {
           'PyUFunc_SimpleUniformOperationTypeResolver',
           TD('?', cfunc_alias='logical_or', dispatch=[('loops_logical', '?')]),
           TD(no_obj_bool, dispatch=[('loops_minmax', 'fdg')]),
-          TD(O, f='npy_ObjectMax')
+          TD(O, f='npy_ObjectMax'),
+          indexed=flts + ints,
           ),
 'fmin':
     Ufunc(2, 1, ReorderableNone,
@@ -632,7 +635,8 @@ defdict = {
           TD('?', cfunc_alias='logical_and',
                   dispatch=[('loops_logical', '?')]),
           TD(no_obj_bool, dispatch=[('loops_minmax', 'fdg')]),
-          TD(O, f='npy_ObjectMin')
+          TD(O, f='npy_ObjectMin'),
+          indexed=flts + ints,
           ),
 'logaddexp':
     Ufunc(2, 1, MinusInfinity,
diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src
index 09268f681..d445fd700 100644
--- a/numpy/core/src/umath/loops.c.src
+++ b/numpy/core/src/umath/loops.c.src
@@ -455,7 +455,9 @@ 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@_fmax_indexed @TYPE@_maximum_indexed
 #define @TYPE@_fmin @TYPE@_minimum
+#define @TYPE@_fmin_indexed @TYPE@_minimum_indexed
 
 NPY_NO_EXPORT void
 @TYPE@__ones_like(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(data))
@@ -538,8 +540,8 @@ 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))
+         char **args, npy_intp const *dimensions, npy_intp const *steps,
+         void *NPY_UNUSED(func))
 {
     char *ip1 = args[0];
     char *indx = args[1];
@@ -902,6 +904,7 @@ NPY_NO_EXPORT void
         }
     }
 }
+
 /**end repeat1**/
 
 /**begin repeat1
@@ -1695,7 +1698,9 @@ HALF_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps, void
 }
 
 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))
+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];
     char *indx = args[1];
@@ -1812,6 +1817,27 @@ HALF_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps, void
     }
     /* npy_half_isnan will never set floatstatus_invalid, so do not clear */
 }
+
+NPY_NO_EXPORT int
+HALF_@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];
+    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);
+        npy_half v = *(npy_half *)value;
+        *indexed = (@OP@(*indexed, v) || npy_half_isnan(*indexed)) ? *indexed : v;
+    }
+    return 0;
+}
+
 /**end repeat**/
 
 /**begin repeat
@@ -1827,8 +1853,29 @@ HALF_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps, void
         const npy_half in2 = *(npy_half *)ip2;
         *((npy_half *)op1) = (@OP@(in1, in2) || npy_half_isnan(in2)) ? in1 : in2;
     }
-    /* npy_half_isnan will never set floatstatus_invalid, so do not clear */
+    /* no need to clear floatstatus_invalid */
+}
+
+NPY_NO_EXPORT int
+HALF_@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];
+    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);
+        npy_half v = *(npy_half *)value;
+        *indexed = (@OP@(*indexed, v) || npy_half_isnan(v)) ? *indexed: v;
+    }
+    return 0;
 }
+
 /**end repeat**/
 
 NPY_NO_EXPORT void
diff --git a/numpy/core/src/umath/loops.h.src b/numpy/core/src/umath/loops.h.src
index f1d91d1fe..f773d94bc 100644
--- a/numpy/core/src/umath/loops.h.src
+++ b/numpy/core/src/umath/loops.h.src
@@ -137,6 +137,8 @@ NPY_CPU_DISPATCH_DECLARE(NPY_NO_EXPORT void @TYPE@_@kind@,
 #define @S@@TYPE@_floor_divide_indexed @S@@TYPE@_divide_indexed
 #define @S@@TYPE@_fmax @S@@TYPE@_maximum
 #define @S@@TYPE@_fmin @S@@TYPE@_minimum
+#define @S@@TYPE@_fmax_indexed @S@@TYPE@_maximum_indexed
+#define @S@@TYPE@_fmin_indexed @S@@TYPE@_minimum_indexed
 
 NPY_NO_EXPORT void
 @S@@TYPE@__ones_like(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(data));
@@ -192,10 +194,13 @@ NPY_NO_EXPORT void
 
 /**begin repeat2
  * #kind = maximum, minimum#
- * #OP =  >, <#
  **/
 NPY_NO_EXPORT void
 @S@@TYPE@_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func));
+
+NPY_NO_EXPORT int
+@S@@TYPE@_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context), char *const *args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func));
+
 /**end repeat2**/
 
 NPY_NO_EXPORT void
@@ -475,18 +480,24 @@ NPY_NO_EXPORT void
 
 /**begin repeat1
  * #kind = maximum, minimum#
- * #OP =  >=, <=#
  **/
 NPY_NO_EXPORT void
 @TYPE@_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func));
+
+NPY_NO_EXPORT int
+@TYPE@_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context), char *const *args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func));
+
 /**end repeat1**/
 
 /**begin repeat1
  * #kind = fmax, fmin#
- * #OP =  >=, <=#
  **/
 NPY_NO_EXPORT void
 @TYPE@_@kind@(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func));
+
+NPY_NO_EXPORT int
+@TYPE@_@kind@_indexed(PyArrayMethod_Context *NPY_UNUSED(context), char *const *args, npy_intp const *dimensions, npy_intp const *steps, NpyAuxData *NPY_UNUSED(func));
+
 /**end repeat1**/
 
 NPY_NO_EXPORT void
diff --git a/numpy/core/src/umath/loops_minmax.dispatch.c.src b/numpy/core/src/umath/loops_minmax.dispatch.c.src
index 237c8e933..9d8667d38 100644
--- a/numpy/core/src/umath/loops_minmax.dispatch.c.src
+++ b/numpy/core/src/umath/loops_minmax.dispatch.c.src
@@ -451,6 +451,24 @@ clear_fp:
 #endif
 }
 
+
+NPY_NO_EXPORT int NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@_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];
+    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 = SCALAR_OP(*indexed, *(@type@ *)value);
+    }
+    return 0;
+}
+
 #undef SCALAR_OP
 
 #endif // !fp_only || (is_fp && fp_only)
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index b0c435f33..a3f56b7a3 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -1980,16 +1980,16 @@ class TestUfunc:
         assert_equal(aa, [0, 1, 202, 3, 4, 105, 6, 7, 8, 9])
 
         with pytest.raises(ValueError):
-            # extraneous second operand 
+            # extraneous second operand
             np.negative.at(a, [2, 5, 3], [1, 2, 3])
 
         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]
+    indexed_ufuncs = [np.add, np.subtract, np.multiply, np.floor_divide,
+                      np.maximum, np.minimum, np.fmax, np.fmin]
 
     @pytest.mark.parametrize(
                 "typecode", np.typecodes['AllInteger'] + np.typecodes['Float'])
@@ -2204,7 +2204,7 @@ class TestUfunc:
     def test_at_output_casting(self):
         arr = np.array([-1])
         np.equal.at(arr, [0], [0])
-        assert arr[0] == 0 
+        assert arr[0] == 0
 
     def test_at_broadcast_failure(self):
         arr = np.arange(5)