Merge pull request numpy#21499 from mattip/avoid-0dim-reduce

ENH: avoid looping when dimensions[0] == 0 or array.size == 0

Author: mattip

numpy/core/src/multiarray/array_method.c

@@ -863,15 +863,17 @@ generic_masked_strided_loop(PyArrayMethod_Context *context,
 
         /* Process unmasked values */
         mask = npy_memchr(mask, 0, mask_stride, N, &subloopsize, 0);
-        int res = strided_loop(context,
-                dataptrs, &subloopsize, strides, strided_loop_auxdata);
-        if (res != 0) {
-            return res;
-        }
-        for (int i = 0; i < nargs; i++) {
-            dataptrs[i] += subloopsize * strides[i];
+        if (subloopsize > 0) {
+            int res = strided_loop(context,
+                    dataptrs, &subloopsize, strides, strided_loop_auxdata);
+            if (res != 0) {
+                return res;
+            }
+            for (int i = 0; i < nargs; i++) {
+                dataptrs[i] += subloopsize * strides[i];
+            }
+            N -= subloopsize;
         }
-        N -= subloopsize;
     } while (N > 0);
 
     return 0;

numpy/core/src/umath/fast_loop_macros.h

@@ -10,6 +10,8 @@
 #ifndef _NPY_UMATH_FAST_LOOP_MACROS_H_
 #define _NPY_UMATH_FAST_LOOP_MACROS_H_
 
+#include <assert.h>
+
 /*
  * MAX_STEP_SIZE is used to determine if we need to use SIMD version of the ufunc.
  * Very large step size can be as slow as processing it using scalar. The
@@ -99,12 +101,19 @@ abs_ptrdiff(char *a, char *b)
 
 #define IS_OUTPUT_CONT(tout) (steps[1] == sizeof(tout))
 
-#define IS_BINARY_REDUCE ((args[0] == args[2])\
+/*
+ * Make sure dimensions is non-zero with an assert, to allow subsequent code
+ * to ignore problems of accessing invalid memory
+ */
+
+#define IS_BINARY_REDUCE (assert(dimensions[0] != 0), \
+        (args[0] == args[2])\
         && (steps[0] == steps[2])\
         && (steps[0] == 0))
 
 /* input contiguous (for binary reduces only) */
-#define IS_BINARY_REDUCE_INPUT_CONT(tin) (steps[1] == sizeof(tin))
+#define IS_BINARY_REDUCE_INPUT_CONT(tin) (assert(dimensions[0] != 0), \
+         steps[1] == sizeof(tin))
 
 /* binary loop input and output contiguous */
 #define IS_BINARY_CONT(tin, tout) (steps[0] == sizeof(tin) && \

numpy/core/src/umath/ufunc_object.c

@@ -1321,6 +1321,10 @@ try_trivial_single_output_loop(PyArrayMethod_Context *context,
      */
     char *data[NPY_MAXARGS];
     npy_intp count = PyArray_MultiplyList(operation_shape, operation_ndim);
+    if (count == 0) {
+        /* Nothing to do */
+        return 0;
+    }
     NPY_BEGIN_THREADS_DEF;
 
     PyArrayMethod_StridedLoop *strided_loop;
@@ -2819,7 +2823,7 @@ reduce_loop(PyArrayMethod_Context *context,
         npy_intp const *countptr, NpyIter_IterNextFunc *iternext,
         int needs_api, npy_intp skip_first_count)
 {
-    int retval;
+    int retval = 0;
     char *dataptrs_copy[4];
     npy_intp strides_copy[4];
     npy_bool masked;
@@ -2849,19 +2853,20 @@ reduce_loop(PyArrayMethod_Context *context,
                     count = 0;
                 }
             }
-
-            /* Turn the two items into three for the inner loop */
-            dataptrs_copy[0] = dataptrs[0];
-            dataptrs_copy[1] = dataptrs[1];
-            dataptrs_copy[2] = dataptrs[0];
-            strides_copy[0] = strides[0];
-            strides_copy[1] = strides[1];
-            strides_copy[2] = strides[0];
-
-            retval = strided_loop(context,
-                    dataptrs_copy, &count, strides_copy, auxdata);
-            if (retval < 0) {
-                goto finish_loop;
+            if (count > 0) {
+                /* Turn the two items into three for the inner loop */
+                dataptrs_copy[0] = dataptrs[0];
+                dataptrs_copy[1] = dataptrs[1];
+                dataptrs_copy[2] = dataptrs[0];
+                strides_copy[0] = strides[0];
+                strides_copy[1] = strides[1];
+                strides_copy[2] = strides[0];
+
+                retval = strided_loop(context,
+                        dataptrs_copy, &count, strides_copy, auxdata);
+                if (retval < 0) {
+                    goto finish_loop;
+                }
             }
 
             /* Advance loop, and abort on error (or finish) */