Skip to content

Commit

Permalink
BUG: Segfault in nditer buffer dealloc for Object arrays
Browse files Browse the repository at this point in the history
  • Loading branch information
@ahaldane
ahaldane committed Feb 22, 2021
1 parent 572d5a1 commit 5d864e8
Show file tree
Hide file tree
Showing 4 changed files with 35 additions and 8 deletions.
4 changes: 3 additions & 1 deletion numpy/core/src/multiarray/einsum.c.src
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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 */
Expand Down
4 changes: 3 additions & 1 deletion numpy/core/src/multiarray/nditer_api.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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) {
Expand Down
22 changes: 16 additions & 6 deletions numpy/core/src/umath/ufunc_object.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -1525,14 +1526,15 @@ 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);

/* Execute the loop */
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;
}
Expand Down Expand Up @@ -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);

Expand All @@ -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;

Expand Down Expand Up @@ -2973,14 +2976,15 @@ 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);
}
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;
Expand Down Expand Up @@ -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)) {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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 */
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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;
}
Expand Down
13 changes: 13 additions & 0 deletions numpy/core/tests/test_nditer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 5d864e8

Please sign in to comment.