BUG: umath: in reduceat, arrays must be copied on overlap even if they are the same

5 files changed, 50 insertions, 3 deletions

diff --git a/doc/source/reference/c-api.iterator.rst b/doc/source/reference/c-api.iterator.rst
index b5d00f4be..5761e56c2 100644
--- a/doc/source/reference/c-api.iterator.rst
+++ b/doc/source/reference/c-api.iterator.rst
@@ -474,7 +474,8 @@ Construction and Destruction
               via *different* index/dtype/shape combinations.
 
             - In particular, unless the arrays have the same shape, dtype,
-              strides, and start address, any shared common data byte accessible
+              strides, start address, and NPY_ITER_OVERLAP_NOT_SAME is not specified,
+              any shared common data byte accessible
               by indexing implies overlap.
 
             Because exact overlap detection has exponential runtime
@@ -618,6 +619,15 @@ Construction and Destruction
             returns true from the corresponding element in the ARRAYMASK
             operand.
 
+        .. c:var:: NPY_ITER_OVERLAP_NOT_SAME
+
+            In the memory overlap checks done when ``NPY_ITER_COPY_IF_OVERLAP``
+            is specified, consider this array as overlapping even if it is
+            exactly the same as another array.
+
+            This flag should be set on arrays that are not accessed in the
+            iterator order.
+
 .. c:function:: NpyIter* NpyIter_AdvancedNew(npy_intp nop, PyArrayObject** op, npy_uint32 flags, NPY_ORDER order, NPY_CASTING casting, npy_uint32* op_flags, PyArray_Descr** op_dtypes, int oa_ndim, int** op_axes, npy_intp* itershape, npy_intp buffersize)
 
     Extends :c:func:`NpyIter_MultiNew` with several advanced options providing
diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index 06d351a6a..3c5af9408 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -1045,6 +1045,11 @@ typedef void (NpyIter_GetMultiIndexFunc)(NpyIter *iter,
 #define NPY_ITER_WRITEMASKED                0x10000000
 /* This array is the mask for all WRITEMASKED operands */
 #define NPY_ITER_ARRAYMASK                  0x20000000
+/*
+ * Consider this array as overlapping for COPY_IF_OVERLAP,
+ * even if it is exactly the same as another array.
+ */
+#define NPY_ITER_OVERLAP_NOT_SAME           0x40000000
 
 #define NPY_ITER_GLOBAL_FLAGS               0x0000ffff
 #define NPY_ITER_PER_OP_FLAGS               0xffff0000
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index 2bd349f95..f8829d0b9 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -2750,7 +2750,9 @@ npyiter_allocate_arrays(NpyIter *iter,
                  * to make copies, because ufunc inner loops are assumed to
                  * deal with that
                  */
-                if (PyArray_BYTES(op[iop]) == PyArray_BYTES(op[iother]) &&
+                if (!(op_flags[iop] & NPY_ITER_OVERLAP_NOT_SAME) &&
+                    !(op_flags[iother] & NPY_ITER_OVERLAP_NOT_SAME) &&
+                    PyArray_BYTES(op[iop]) == PyArray_BYTES(op[iother]) &&
                     PyArray_NDIM(op[iop]) == PyArray_NDIM(op[iother]) &&
                     PyArray_CompareLists(PyArray_DIMS(op[iop]),
                                          PyArray_DIMS(op[iother]),
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 3c0c54222..4ade0e34c 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -3627,7 +3627,8 @@ PyUFunc_Reduceat(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *ind,
                       NPY_ITER_ALIGNED;
         op_flags[1] = NPY_ITER_READONLY|
                       NPY_ITER_COPY|
-                      NPY_ITER_ALIGNED;
+                      NPY_ITER_ALIGNED|
+                      NPY_ITER_OVERLAP_NOT_SAME;
         op_flags[2] = NPY_ITER_READONLY;
 
         op_dtypes[1] = op_dtypes[0];
diff --git a/numpy/core/tests/test_mem_overlap.py b/numpy/core/tests/test_mem_overlap.py
index a5cb5a4f5..0e9b4d50f 100644
--- a/numpy/core/tests/test_mem_overlap.py
+++ b/numpy/core/tests/test_mem_overlap.py
@@ -703,6 +703,20 @@ class TestUFunc(object):
         self.check_unary_fuzz(do_reduceat, get_out_axis_size,
                               dtype=np.int16, count=500)
 
+    def test_binary_ufunc_reduceat_manual(self):
+        def check(ufunc, a, ind, out):
+            c1 = ufunc.reduceat(a.copy(), ind.copy(), out=out.copy())
+            c2 = ufunc.reduceat(a, ind, out=out)
+            assert_array_equal(c1, c2)
+
+        # Exactly same input/output arrays
+        a = np.arange(10000, dtype=np.int16)
+        check(np.add, a, a[::-1].copy(), a)
+
+        # Overlap with index
+        a = np.arange(10000, dtype=np.int16)
+        check(np.add, a, a[::-1], a)
+
     def test_unary_gufunc_fuzz(self):
         shapes = [7, 13, 8, 21, 29, 32]
         gufunc = umath_tests.euclidean_pdist
@@ -844,6 +858,21 @@ class TestUFunc(object):
                         check(x[:1].reshape([]), y[::-1])
                         check(x[-1:].reshape([]), y[::-1])
 
+    def test_unary_ufunc_where_same(self):
+        # Check behavior at wheremask overlap
+        ufunc = np.invert
+
+        def check(a, out, mask):
+            c1 = ufunc(a, out=out.copy(), where=mask.copy())
+            c2 = ufunc(a, out=out, where=mask)
+            assert_array_equal(c1, c2)
+
+        # Check behavior with same input and output arrays
+        x = np.arange(100).astype(np.bool_)
+        check(x, x, x)
+        check(x, x.copy(), x)
+        check(x, x, x.copy())
+
     def test_binary_ufunc_1d_manual(self):
         ufunc = np.add