BUG: nditer: The nditer was reusing the reduce loop inappropriately (#1938)

3 files changed, 29 insertions, 7 deletions

diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c
index 0dccd6148..6db3279d6 100644
--- a/numpy/core/src/multiarray/nditer_api.c
+++ b/numpy/core/src/multiarray/nditer_api.c
@@ -2565,7 +2565,7 @@ npyiter_checkreducesize(NpyIter *iter, npy_intp count,
     npy_intp sizeof_axisdata;
     npy_intp coord, shape, *strides;
     npy_intp reducespace = 1, factor;
-    npy_bool nonzerocoord = 0;
+    npy_bool nonzerocoord;
 
     char *op_itflags = NIT_OPITFLAGS(iter);
     char stride0op[NPY_MAXARGS];
@@ -2596,6 +2596,9 @@ npyiter_checkreducesize(NpyIter *iter, npy_intp count,
     factor = shape;
     NIT_ADVANCE_AXISDATA(axisdata, 1);
 
+    /* Initialize nonzerocoord based on the first coordinate */
+    nonzerocoord = (coord != 0);
+
     /* Go forward through axisdata, calculating the space available */
     for (idim = 1; idim < ndim && reducespace < count;
                                 ++idim, NIT_ADVANCE_AXISDATA(axisdata, 1)) {
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index b80d3b5af..c99ce78c2 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -46,10 +46,10 @@
 #define NPY_UF_DBG_TRACING 0
 
 #if NPY_UF_DBG_TRACING
-#define NPY_UF_DBG_PRINT(s) printf("%s", s)
-#define NPY_UF_DBG_PRINT1(s, p1) printf(s, p1)
-#define NPY_UF_DBG_PRINT2(s, p1, p2) printf(s, p1, p2)
-#define NPY_UF_DBG_PRINT3(s, p1, p2, p3) printf(s, p1, p2, p3)
+#define NPY_UF_DBG_PRINT(s) {printf("%s", s);fflush(stdout);}
+#define NPY_UF_DBG_PRINT1(s, p1) {printf((s), (p1));fflush(stdout);}
+#define NPY_UF_DBG_PRINT2(s, p1, p2) {printf(s, p1, p2);fflush(stdout);}
+#define NPY_UF_DBG_PRINT3(s, p1, p2, p3) {printf(s, p1, p2, p3);fflush(stdout);}
 #else
 #define NPY_UF_DBG_PRINT(s)
 #define NPY_UF_DBG_PRINT1(s, p1)
@@ -2508,7 +2508,7 @@ get_masked_binary_op_function(PyUFuncObject *self, PyArrayObject *arr,
     char *ufunc_name = self->name ? self->name : "(unknown)";
 
     NPY_UF_DBG_PRINT1("Getting masked binary op function for type number %d\n",
-                                *otype);
+                                otype);
 
     *out_dtype = NULL;
 
@@ -2682,6 +2682,8 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
         retcode = get_binary_op_function(self, &otype_final,
                                 &innerloop, &innerloopdata);
 
+        NPY_UF_DBG_PRINT2("Loop retcode %d, otype final %d\n",
+                                                retcode, otype_final);
         /*
          * Set up the output data type, using the input's exact
          * data type if the type number didn't change to preserve
@@ -3030,7 +3032,7 @@ PyUFunc_Accumulate(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
     }
 
 #if NPY_UF_DBG_TRACING
-    printf("Found %s.%s inner loop with dtype :  ", ufunc_name, opname);
+    printf("Found %s.accumulate inner loop with dtype :  ", ufunc_name);
     PyObject_Print((PyObject *)op_dtypes[0], stdout, 0);
     printf("\n");
 #endif
diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
index cb562ba3f..b13adf6d6 100644
--- a/numpy/core/tests/test_nditer.py
+++ b/numpy/core/tests/test_nditer.py
@@ -2295,6 +2295,23 @@ def test_iter_buffering_reduction():
     it.reset()
     assert_equal(it[0], [1,2,1,2])
 
+def test_iter_buffering_reduction_reuse_reduce_loops():
+    # There was a bug triggering reuse of the reduce loop inappropriately,
+    # which caused processing to happen in unnecessarily small chunks
+    # and overran the buffer.
+
+    a = np.zeros((2,7))
+    b = np.zeros((1,7))
+    it = np.nditer([a,b], flags=['reduce_ok', 'external_loop', 'buffered'],
+                    op_flags=[['readonly'], ['readwrite']],
+                    buffersize = 5)
+
+    bufsizes = []
+    for x, y in it:
+        bufsizes.append(x.shape[0])
+    assert_equal(bufsizes, [5,2,5,2])
+    assert_equal(sum(bufsizes), a.size)
+
 def test_iter_writemasked_badinput():
     a = np.zeros((2,3))
     b = np.zeros((3,))