Merge pull request #18450 from ahaldane/fix_nditer_od_segfault

BUG: Segfault in nditer buffer dealloc for Object arrays

4 files changed, 35 insertions, 8 deletions

diff --git a/numpy/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src
index 6ad375f67..85806fab3 100644
--- a/numpy/core/src/multiarray/einsum.c.src
+++ b/numpy/core/src/multiarray/einsum.c.src
@@ -1100,6 +1100,7 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
         char **dataptr;
         npy_intp *stride;
         npy_intp *countptr;
+        int needs_api;
         NPY_BEGIN_THREADS_DEF;
 
         iternext = NpyIter_GetIterNext(iter, NULL);
@@ -1110,12 +1111,13 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
         dataptr = NpyIter_GetDataPtrArray(iter);
         stride = NpyIter_GetInnerStrideArray(iter);
         countptr = NpyIter_GetInnerLoopSizePtr(iter);
+        needs_api = NpyIter_IterationNeedsAPI(iter);
 
         NPY_BEGIN_THREADS_NDITER(iter);
         NPY_EINSUM_DBG_PRINT("Einsum loop\n");
         do {
             sop(nop, dataptr, stride, *countptr);
-        } while(iternext(iter));
+        } while (!(needs_api && PyErr_Occurred()) && iternext(iter));
         NPY_END_THREADS;
 
         /* If the API was needed, it may have thrown an error */
diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c
index fceb58f33..3403ce98b 100644
--- a/numpy/core/src/multiarray/nditer_api.c
+++ b/numpy/core/src/multiarray/nditer_api.c
@@ -2632,6 +2632,7 @@ npyiter_clear_buffers(NpyIter *iter)
     /* Cleanup any buffers with references */
     char **buffers = NBF_BUFFERS(bufferdata);
     PyArray_Descr **dtypes = NIT_DTYPES(iter);
+    npyiter_opitflags *op_itflags = NIT_OPITFLAGS(iter);
     for (int iop = 0; iop < nop; ++iop, ++buffers) {
         /*
          * We may want to find a better way to do this, on the other hand,
@@ -2640,7 +2641,8 @@ npyiter_clear_buffers(NpyIter *iter)
          * a well defined state (either NULL or owning the reference).
          * Only we implement cleanup
          */
-        if (!PyDataType_REFCHK(dtypes[iop])) {
+        if (!PyDataType_REFCHK(dtypes[iop]) ||
+                !(op_itflags[iop]&NPY_OP_ITFLAG_USINGBUFFER)) {
             continue;
         }
         if (*buffers == 0) {
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 1a035eb61..cd6e27a35 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -1431,6 +1431,7 @@ iterator_loop(PyUFuncObject *ufunc,
     char **dataptr;
     npy_intp *stride;
     npy_intp *count_ptr;
+    int needs_api;
 
     PyArrayObject **op_it;
     npy_uint32 iter_flags;
@@ -1525,6 +1526,7 @@ iterator_loop(PyUFuncObject *ufunc,
         dataptr = NpyIter_GetDataPtrArray(iter);
         stride = NpyIter_GetInnerStrideArray(iter);
         count_ptr = NpyIter_GetInnerLoopSizePtr(iter);
+        needs_api = NpyIter_IterationNeedsAPI(iter);
 
         NPY_BEGIN_THREADS_NDITER(iter);
 
@@ -1532,7 +1534,7 @@ iterator_loop(PyUFuncObject *ufunc,
         do {
             NPY_UF_DBG_PRINT1("iterator loop count %d\n", (int)*count_ptr);
             innerloop(dataptr, count_ptr, stride, innerloopdata);
-        } while (iternext(iter));
+        } while (!(needs_api && PyErr_Occurred()) && iternext(iter));
 
         NPY_END_THREADS;
     }
@@ -1859,6 +1861,7 @@ execute_fancy_ufunc_loop(PyUFuncObject *ufunc,
         dataptr = NpyIter_GetDataPtrArray(iter);
         strides = NpyIter_GetInnerStrideArray(iter);
         countptr = NpyIter_GetInnerLoopSizePtr(iter);
+        needs_api = NpyIter_IterationNeedsAPI(iter);
 
         NPY_BEGIN_THREADS_NDITER(iter);
 
@@ -1869,7 +1872,7 @@ execute_fancy_ufunc_loop(PyUFuncObject *ufunc,
             innerloop(dataptr, strides,
                         dataptr[nop], strides[nop],
                         *countptr, innerloopdata);
-        } while (iternext(iter));
+        } while (!(needs_api && PyErr_Occurred()) && iternext(iter));
 
         NPY_END_THREADS;
 
@@ -2973,6 +2976,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
         }
         dataptr = NpyIter_GetDataPtrArray(iter);
         count_ptr = NpyIter_GetInnerLoopSizePtr(iter);
+        needs_api = NpyIter_IterationNeedsAPI(iter);
 
         if (!needs_api && !NpyIter_IterationNeedsAPI(iter)) {
             NPY_BEGIN_THREADS_THRESHOLDED(total_problem_size);
@@ -2980,7 +2984,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
         do {
             inner_dimensions[0] = *count_ptr;
             innerloop(dataptr, inner_dimensions, inner_strides, innerloopdata);
-        } while (iternext(iter));
+        } while (!(needs_api && PyErr_Occurred()) && iternext(iter));
 
         if (!needs_api && !NpyIter_IterationNeedsAPI(iter)) {
             NPY_END_THREADS;
@@ -3520,6 +3524,10 @@ reduce_loop(NpyIter *iter, char **dataptrs, npy_intp const *strides,
             innerloop(dataptrs_copy, &count,
                         strides_copy, innerloopdata);
 
+            if (needs_api && PyErr_Occurred()) {
+                break;
+            }
+
             /* Jump to the faster loop when skipping is done */
             if (skip_first_count == 0) {
                 if (iternext(iter)) {
@@ -3569,7 +3577,7 @@ reduce_loop(NpyIter *iter, char **dataptrs, npy_intp const *strides,
                 n = 1;
             }
         }
-    } while (iternext(iter));
+    } while (!(needs_api && PyErr_Occurred()) && iternext(iter));
 
 finish_loop:
     NPY_END_THREADS;
@@ -3882,6 +3890,7 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
             goto fail;
         }
         dataptr = NpyIter_GetDataPtrArray(iter);
+        needs_api = NpyIter_IterationNeedsAPI(iter);
 
 
         /* Execute the loop with just the outer iterator */
@@ -3932,7 +3941,7 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
                 innerloop(dataptr_copy, &count_m1,
                             stride_copy, innerloopdata);
             }
-        } while (iternext(iter));
+        } while (!(needs_api && PyErr_Occurred()) && iternext(iter));
 
         NPY_END_THREADS;
     }
@@ -4263,6 +4272,7 @@ PyUFunc_Reduceat(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *ind,
         npy_intp stride0_ind = PyArray_STRIDE(op[0], axis);
 
         int itemsize = op_dtypes[0]->elsize;
+        int needs_api = NpyIter_IterationNeedsAPI(iter);
 
         /* Get the variables needed for the loop */
         iternext = NpyIter_GetIterNext(iter, NULL);
@@ -4327,7 +4337,7 @@ PyUFunc_Reduceat(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *ind,
                                 stride_copy, innerloopdata);
                 }
             }
-        } while (iternext(iter));
+        } while (!(needs_api && PyErr_Occurred()) && iternext(iter));
 
         NPY_END_THREADS;
     }
diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
index 411095199..f34547f9c 100644
--- a/numpy/core/tests/test_nditer.py
+++ b/numpy/core/tests/test_nditer.py
@@ -2869,6 +2869,19 @@ def test_0d_iter():
     assert_equal(vals['c'], [[(0.5)]*3]*2)
     assert_equal(vals['d'], 0.5)
 
+def test_object_iter_cleanup():
+    # see gh-18450
+    # object arrays can raise a python exception in ufunc inner loops using
+    # nditer, which should cause iteration to stop & cleanup. There were bugs
+    # in the nditer cleanup when decref'ing object arrays.
+    # This test would trigger valgrind "uninitialized read" before the bugfix.
+    assert_raises(TypeError, lambda: np.zeros((17000, 2), dtype='f4') * None)
+
+    # this more explicit code also triggers the invalid access
+    arr = np.arange(np.BUFSIZE * 10).reshape(10, -1).astype(str)
+    oarr = arr.astype(object)
+    oarr[:, -1] = None
+    assert_raises(TypeError, lambda: np.add(oarr[:, ::-1], arr[:, ::-1]))
 
 def test_iter_too_large():
     # The total size of the iterator must not exceed the maximum intp due