diff options
| -rw-r--r-- | numpy/core/src/multiarray/item_selection.c | 308 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/reduction.c | 233 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/reduction.h | 113 |
3 files changed, 443 insertions, 211 deletions
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index 28fc7e7fa..88b81f43d 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -19,6 +19,7 @@ #include "ctors.h" #include "lowlevel_strided_loops.h" #include "na_singleton.h" +#include "reduction.h" #include "item_selection.h" @@ -1891,6 +1892,90 @@ count_boolean_trues(int ndim, char *data, npy_intp *ashape, npy_intp *astrides) return count; } +static int +assign_reduce_unit_zero(PyArrayObject *result, int preservena, void *data) +{ + return PyArray_AssignZero(result, NULL, preservena, NULL); +} + +static void +reduce_count_nonzero_inner_loop(NpyIter *iter, + char **dataptr, + npy_intp *strideptr, + npy_intp *countptr, + NpyIter_IterNextFunc *iternext, + int needs_api, + void *data) +{ + PyArray_NonzeroFunc *nonzero = (PyArray_NonzeroFunc *)data; + PyArrayObject *arr = NpyIter_GetOperandArray(iter)[1]; + + NPY_BEGIN_THREADS_DEF; + + if (!needs_api) { + NPY_BEGIN_THREADS; + } + + do { + char *data0 = dataptr[0], *data1 = dataptr[1]; + npy_intp stride0 = strideptr[0], stride1 = strideptr[1]; + npy_intp count = *countptr; + + while (count--) { + if (nonzero(data1, arr)) { + ++(*(npy_intp *)data0); + } + data0 += stride0; + data1 += stride1; + } + } while (iternext(iter)); + + if (!needs_api) { + NPY_END_THREADS; + } +} + +static void +reduce_count_nonzero_masked_inner_loop(NpyIter *iter, + char **dataptr, + npy_intp *strideptr, + npy_intp *countptr, + NpyIter_IterNextFunc *iternext, + int needs_api, + void *data) +{ + PyArray_NonzeroFunc *nonzero = (PyArray_NonzeroFunc *)data; + PyArrayObject *arr = NpyIter_GetOperandArray(iter)[1]; + + NPY_BEGIN_THREADS_DEF; + + if (!needs_api) { + NPY_BEGIN_THREADS; + } + + do { + char *data0 = dataptr[0], *data1 = dataptr[1], *data2 = dataptr[2]; + npy_intp stride0 = strideptr[0], stride1 = strideptr[1], + stride2 = strideptr[2]; + npy_intp count = *countptr; + + while (count--) { + if (NpyMaskValue_IsExposed((npy_mask)*data2) && + nonzero(data1, arr)) { + ++(*(npy_intp *)data0); + } + data0 += stride0; + data1 += stride1; + data2 += stride2; + } + } while (iternext(iter)); + + if (!needs_api) { + NPY_END_THREADS; + } +} + + /* * A full reduction version of PyArray_CountNonzero, supporting * an 'out' parameter and doing the count as a reduction along @@ -1902,15 +1987,8 @@ PyArray_ReduceCountNonzero(PyArrayObject *arr, PyArrayObject *out, npy_bool *axis_flags, int skipna) { PyArray_NonzeroFunc *nonzero; - int use_maskna; + PyArrayObject *result; PyArray_Descr *dtype; - PyArrayObject *result = NULL; - - /* Iterator parameters */ - NpyIter *iter = NULL; - PyArrayObject *op[2]; - PyArray_Descr *op_dtypes[2]; - npy_uint32 flags, op_flags[2]; nonzero = PyArray_DESCR(arr)->f->nonzero; if (nonzero == NULL) { @@ -1920,217 +1998,25 @@ PyArray_ReduceCountNonzero(PyArrayObject *arr, PyArrayObject *out, return NULL; } - use_maskna = PyArray_HASMASKNA(arr); - - /* - * If 'arr' has an NA mask, but 'out' doesn't, validate that 'arr' - * contains no NA values so we can ignore the mask entirely. - */ - if (use_maskna && !skipna && out != NULL && !PyArray_HASMASKNA(out)) { - if (PyArray_ContainsNA(arr)) { - PyErr_SetString(PyExc_ValueError, - "Cannot assign NA value to an array which " - "does not support NAs"); - return NULL; - } - else { - use_maskna = 0; - } - } - - /* This reference gets stolen by PyArray_CreateReduceResult */ dtype = PyArray_DescrFromType(NPY_INTP); if (dtype == NULL) { return NULL; } - /* This either conforms 'out' to the ndim of 'arr', or allocates - * a new array appropriate for this reduction. - */ - result = PyArray_CreateReduceResult(arr, out, - dtype, axis_flags, !skipna && use_maskna, - "count_nonzero"); - if (result == NULL) { - return NULL; - } - /* - * Do the reduction on the NA mask before the data. This way - * we can avoid modifying the outputs which end up masked, obeying - * the required NA masking semantics. - */ - if (use_maskna && !skipna) { - if (PyArray_ReduceMaskNAArray(result, arr) < 0) { - Py_DECREF(result); - return NULL; - } - - /* Short circuit any calculation if the result is 0-dim NA */ - if (PyArray_SIZE(result) == 1 && - !NpyMaskValue_IsExposed( - (npy_mask)*PyArray_MASKNA_DATA(result))) { - goto finish; - } - } - - /* - * Initialize the result to all zeros, except for the NAs - * when skipna is False. - */ - if (PyArray_AssignZero(result, NULL, !skipna, NULL) < 0) { - Py_DECREF(result); - return NULL; - } - - /* Set up the iterator */ - op[0] = result; - op[1] = arr; - op_dtypes[0] = PyArray_DescrFromType(NPY_INTP); - if (op_dtypes[0] == NULL) { - Py_DECREF(result); - return NULL; - } - op_dtypes[1] = NULL; - - flags = NPY_ITER_BUFFERED | - NPY_ITER_EXTERNAL_LOOP | - NPY_ITER_DONT_NEGATE_STRIDES | - NPY_ITER_ZEROSIZE_OK | - NPY_ITER_REDUCE_OK | - NPY_ITER_REFS_OK; - op_flags[0] = NPY_ITER_READWRITE | - NPY_ITER_ALIGNED | - NPY_ITER_NO_SUBTYPE; - op_flags[1] = NPY_ITER_READONLY | - NPY_ITER_ALIGNED; - - /* Add mask-related flags */ - if (use_maskna) { - if (skipna) { - /* The output's mask has been set to all exposed already */ - op_flags[0] |= NPY_ITER_IGNORE_MASKNA; - /* Need the input's mask to determine what to skip */ - op_flags[1] |= NPY_ITER_USE_MASKNA; - } - else { - /* Iterate over the output's mask */ - op_flags[0] |= NPY_ITER_USE_MASKNA; - /* The input's mask is already incorporated in the output's mask */ - op_flags[1] |= NPY_ITER_IGNORE_MASKNA; - } - } - else { - /* - * If 'out' had no mask, and 'arr' did, we checked that 'arr' - * contains no NA values and can ignore the masks. - */ - op_flags[0] |= NPY_ITER_IGNORE_MASKNA; - op_flags[1] |= NPY_ITER_IGNORE_MASKNA; - } - - iter = NpyIter_MultiNew(2, op, flags, - NPY_KEEPORDER, NPY_UNSAFE_CASTING, - op_flags, - op_dtypes); - if (iter == NULL) { - Py_DECREF(result); - Py_DECREF(op_dtypes[0]); - return NULL; - } - - if (NpyIter_GetIterSize(iter) != 0) { - int needs_api; - NpyIter_IterNextFunc *iternext; - char **dataptr; - npy_intp *strideptr; - npy_intp *countptr; - NPY_BEGIN_THREADS_DEF; - - iternext = NpyIter_GetIterNext(iter, NULL); - if (iternext == NULL) { - Py_DECREF(result); - Py_DECREF(op_dtypes[0]); - NpyIter_Deallocate(iter); - return NULL; - } - dataptr = NpyIter_GetDataPtrArray(iter); - strideptr = NpyIter_GetInnerStrideArray(iter); - countptr = NpyIter_GetInnerLoopSizePtr(iter); - - needs_api = NpyIter_IterationNeedsAPI(iter) || - PyDataType_REFCHK(PyArray_DESCR(arr)); - - if (!needs_api) { - NPY_BEGIN_THREADS; - } - - /* Straightforward reduction */ - if (!use_maskna) { - do { - char *data0 = dataptr[0]; - char *data1 = dataptr[1]; - npy_intp stride0 = strideptr[0]; - npy_intp stride1 = strideptr[1]; - npy_intp count = *countptr; - - while (count--) { - if (nonzero(data1, arr)) { - ++(*(npy_intp *)data0); - } - data0 += stride0; - data1 += stride1; - } - } while (iternext(iter)); - } - /* Masked reduction */ - else { - do { - char *data0 = dataptr[0]; - char *data1 = dataptr[1]; - char *data2 = dataptr[2]; - npy_intp stride0 = strideptr[0]; - npy_intp stride1 = strideptr[1]; - npy_intp stride2 = strideptr[2]; - npy_intp count = *countptr; - - while (count--) { - if (NpyMaskValue_IsExposed((npy_mask)*data2) && - nonzero(data1, arr)) { - ++(*(npy_intp *)data0); - } - data0 += stride0; - data1 += stride1; - data2 += stride2; - } - } while (iternext(iter)); - } - - if (!needs_api) { - NPY_END_THREADS; - } - - if (needs_api && PyErr_Occurred()) { - Py_DECREF(result); - Py_DECREF(op_dtypes[0]); - NpyIter_Deallocate(iter); - return NULL; - } - } - - Py_DECREF(op_dtypes[0]); - NpyIter_Deallocate(iter); - -finish: - /* Strip out the extra 'one' dimensions in the result */ - if (out == NULL) { - PyArray_RemoveAxesInPlace(result, axis_flags); + result = PyArray_ReduceWrapper(arr, out, + PyArray_DESCR(arr), dtype, + axis_flags, skipna, + &assign_reduce_unit_zero, + &reduce_count_nonzero_inner_loop, + &reduce_count_nonzero_masked_inner_loop, + nonzero, "count_nonzero"); + Py_DECREF(dtype); + if (out == NULL && result != NULL) { + return PyArray_Return(result); } else { - Py_DECREF(result); - result = out; - Py_INCREF(result); + return (PyObject *)result; } - - return PyArray_Return(result); } /*NUMPY_API diff --git a/numpy/core/src/multiarray/reduction.c b/numpy/core/src/multiarray/reduction.c index 09070cb72..30fb33bed 100644 --- a/numpy/core/src/multiarray/reduction.c +++ b/numpy/core/src/multiarray/reduction.c @@ -400,3 +400,236 @@ PyArray_InitializeReduceResult( return op_view; } + +/* + * This function executes all the standard NumPy reduction function + * boilerplate code, just calling assign_unit and the appropriate + * inner loop function where necessary. + * + * operand : The array to be reduced. + * out : NULL, or the array into which to place the result. + * operand_dtype : The dtype the inner loop expects for the operand. + * result_dtype : The dtype the inner loop expects for the result. + * axis_flags : Flags indicating the reduction axes of 'operand'. + * skipna : If true, NAs are skipped instead of propagating. + * assign_unit : If NULL, PyArray_InitializeReduceResult is used, otherwise + * this function is called to initialize the result to + * the reduction's unit. + * inner_loop : The inner loop which does the reduction. + * masked_inner_loop: The inner loop which does the reduction with a mask. + * data : Data which is passed to assign_unit and the inner loop. + * funcname : The name of the reduction function, for error messages. + */ +NPY_NO_EXPORT PyArrayObject * +PyArray_ReduceWrapper(PyArrayObject *operand, PyArrayObject *out, + PyArray_Descr *operand_dtype, + PyArray_Descr *result_dtype, + npy_bool *axis_flags, int skipna, + PyArray_AssignReduceUnitFunc *assign_unit, + PyArray_ReduceInnerLoopFunc *inner_loop, + PyArray_ReduceInnerLoopFunc *masked_inner_loop, + void *data, const char *funcname) +{ + int use_maskna; + PyArrayObject *result = NULL, *op_view = NULL; + + /* Iterator parameters */ + NpyIter *iter = NULL; + PyArrayObject *op[2]; + PyArray_Descr *op_dtypes[2]; + npy_uint32 flags, op_flags[2]; + + use_maskna = PyArray_HASMASKNA(operand); + + /* + * If 'operand' has an NA mask, but 'out' doesn't, validate that 'operand' + * contains no NA values so we can ignore the mask entirely. + */ + if (use_maskna && !skipna && out != NULL && !PyArray_HASMASKNA(out)) { + if (PyArray_ContainsNA(operand)) { + PyErr_SetString(PyExc_ValueError, + "Cannot assign NA value to an array which " + "does not support NAs"); + goto fail; + } + else { + use_maskna = 0; + } + } + + /* + * This either conforms 'out' to the ndim of 'operand', or allocates + * a new array appropriate for this reduction. + */ + Py_INCREF(result_dtype); + result = PyArray_CreateReduceResult(operand, out, + result_dtype, axis_flags, !skipna && use_maskna, + funcname); + if (result == NULL) { + goto fail; + } + + /* + * Do the reduction on the NA mask before the data. This way + * we can avoid modifying the outputs which end up masked, obeying + * the required NA masking semantics. + */ + if (use_maskna && !skipna) { + if (PyArray_ReduceMaskNAArray(result, operand) < 0) { + goto fail; + } + + /* Short circuit any calculation if the result is 0-dim NA */ + if (PyArray_SIZE(result) == 1 && + !NpyMaskValue_IsExposed( + (npy_mask)*PyArray_MASKNA_DATA(result))) { + goto finish; + } + } + + /* + * Initialize the result to the reduction unit if possible, + * otherwise copy the initial values and get a view to the rest. + */ + if (assign_unit != NULL) { + if (assign_unit(result, !skipna, data) < 0) { + goto fail; + } + op_view = operand; + Py_INCREF(op_view); + } + else { + op_view = PyArray_InitializeReduceResult(result, operand, + axis_flags, skipna, funcname); + if (op_view == NULL) { + Py_DECREF(result); + return NULL; + } + } + + /* Set up the iterator */ + op[0] = result; + op[1] = op_view; + op_dtypes[0] = result_dtype; + op_dtypes[1] = operand_dtype; + + flags = NPY_ITER_BUFFERED | + NPY_ITER_EXTERNAL_LOOP | + NPY_ITER_DONT_NEGATE_STRIDES | + NPY_ITER_ZEROSIZE_OK | + NPY_ITER_REDUCE_OK | + NPY_ITER_REFS_OK; + op_flags[0] = NPY_ITER_READWRITE | + NPY_ITER_ALIGNED | + NPY_ITER_NO_SUBTYPE; + op_flags[1] = NPY_ITER_READONLY | + NPY_ITER_ALIGNED; + + /* Add mask-related flags */ + if (use_maskna) { + if (skipna) { + /* The output's mask has been set to all exposed already */ + op_flags[0] |= NPY_ITER_IGNORE_MASKNA; + /* Need the input's mask to determine what to skip */ + op_flags[1] |= NPY_ITER_USE_MASKNA; + } + else { + /* Iterate over the output's mask */ + op_flags[0] |= NPY_ITER_USE_MASKNA; + /* The input's mask is already incorporated in the output's mask */ + op_flags[1] |= NPY_ITER_IGNORE_MASKNA; + } + } + else { + /* + * If 'out' had no mask, and 'operand' did, we checked that 'operand' + * contains no NA values and can ignore the masks. + */ + op_flags[0] |= NPY_ITER_IGNORE_MASKNA; + op_flags[1] |= NPY_ITER_IGNORE_MASKNA; + } + + iter = NpyIter_MultiNew(2, op, flags, + NPY_KEEPORDER, NPY_SAME_KIND_CASTING, + op_flags, + op_dtypes); + if (iter == NULL) { + Py_DECREF(result); + Py_DECREF(op_dtypes[0]); + return NULL; + } + + if (NpyIter_GetIterSize(iter) != 0) { + NpyIter_IterNextFunc *iternext; + char **dataptr; + npy_intp *strideptr; + npy_intp *countptr; + int needs_api; + + iternext = NpyIter_GetIterNext(iter, NULL); + if (iternext == NULL) { + Py_DECREF(result); + Py_DECREF(op_dtypes[0]); + NpyIter_Deallocate(iter); + return NULL; + } + dataptr = NpyIter_GetDataPtrArray(iter); + strideptr = NpyIter_GetInnerStrideArray(iter); + countptr = NpyIter_GetInnerLoopSizePtr(iter); + + needs_api = NpyIter_IterationNeedsAPI(iter); + + /* Straightforward reduction */ + if (!use_maskna) { + if (inner_loop == NULL) { + PyErr_Format(PyExc_RuntimeError, + "reduction operation %s did not supply an " + "unmasked inner loop function", funcname); + goto fail; + } + + inner_loop(iter, dataptr, strideptr, countptr, + iternext, needs_api, data); + } + /* Masked reduction */ + else { + if (masked_inner_loop == NULL) { + PyErr_Format(PyExc_RuntimeError, + "reduction operation %s did not supply a " + "masked inner loop function", funcname); + goto fail; + } + + masked_inner_loop(iter, dataptr, strideptr, countptr, + iternext, needs_api, data); + } + + if (PyErr_Occurred()) { + goto fail; + } + } + + NpyIter_Deallocate(iter); + +finish: + /* Strip out the extra 'one' dimensions in the result */ + if (out == NULL) { + PyArray_RemoveAxesInPlace(result, axis_flags); + } + else { + Py_DECREF(result); + result = out; + Py_INCREF(result); + } + + return result; + +fail: + Py_XDECREF(result); + Py_XDECREF(op_view); + if (iter != NULL) { + NpyIter_Deallocate(iter); + } + + return NULL; +} diff --git a/numpy/core/src/multiarray/reduction.h b/numpy/core/src/multiarray/reduction.h index f19909dbc..b2334f896 100644 --- a/numpy/core/src/multiarray/reduction.h +++ b/numpy/core/src/multiarray/reduction.h @@ -1,4 +1,117 @@ #ifndef _NPY_PRIVATE__REDUCTION_H_ #define _NPY_PRIVATE__REDUCTION_H_ +/* + * This is a function for assigning a reduction unit to the result, + * before doing the reduction computation. If 'preservena' is True, + * any masked NA values in 'result' should not be overwritten. The + * value in 'data' is passed through from PyArray_ReduceWrapper. + * + * This function could, for example, simply be a call like + * return PyArray_AssignZero(result, NULL, preservena, NULL); + * + * It should return -1 on failure, or 0 on success. + */ +typedef int (PyArray_AssignReduceUnitFunc)(PyArrayObject *result, + int preservena, void *data); + +/* + * This is a function for the inner reduce loop. Both the unmasked and + * masked variants have the same prototype, but should behave differently. + * + * The needs_api parameter indicates whether it's ok to release the GIL during + * the inner loop, such as when the iternext() function never calls + * a function which could raise a Python exception. + * + * The unmasked inner loop gets two data pointers and two strides, and should + * look roughly like this: + * NPY_BEGIN_THREADS_DEF; + * if (!needs_api) { + * NPY_BEGIN_THREADS; + * } + * do { + * char *data0 = dataptr[0], *data1 = dataptr[1]; + * npy_intp stride0 = strideptr[0], stride1 = strideptr[1]; + * npy_intp count = *countptr; + * + * while (count--) { + * *(result_t *)data0 = my_reduce_op(*(result_t *)data0, + * *(operand_t *)data1); + * data0 += stride0; + * data1 += stride1; + * } + * } while (iternext(iter)); + * if (!needs_api) { + * NPY_END_THREADS; + * } + * + * The masked inner loop gets three data pointers and three strides, and + * should look roughly like this: + * NPY_BEGIN_THREADS_DEF; + * if (!needs_api) { + * NPY_BEGIN_THREADS; + * } + * do { + * char *data0 = dataptr[0], *data1 = dataptr[1], *data2 = dataptr[2]; + * npy_intp stride0 = strideptr[0], stride1 = strideptr[1], + * stride2 = strideptr[2]; + * npy_intp count = *countptr; + * + * while (count--) { + * if (NpyMaskValue_IsExposed((npy_mask)*data2)) { + * *(result_t *)data0 = my_reduce_op(*(result_t *)data0, + * *(operand_t *)data1); + * } + * data0 += stride0; + * data1 += stride1; + * data2 += stride2; + * } + * } while (iternext(iter)); + * if (!needs_api) { + * NPY_END_THREADS; + * } + * + * Once the inner loop is finished, PyArray_ReduceWrapper calls + * PyErr_Occurred to check if any exception was raised during the + * computation. + */ +typedef void (PyArray_ReduceInnerLoopFunc)(NpyIter *iter, + char **dataptr, + npy_intp *strideptr, + npy_intp *countptr, + NpyIter_IterNextFunc *iternext, + int needs_api, + void *data); + +/* + * This function executes all the standard NumPy reduction function + * boilerplate code, just calling assign_unit and the appropriate + * inner loop function where necessary. + * + * operand : The array to be reduced. + * out : NULL, or the array into which to place the result. + * operand_dtype : The dtype the inner loop expects for the operand. + * result_dtype : The dtype the inner loop expects for the result. + * axis_flags : Flags indicating the reduction axes of 'operand'. + * skipna : If true, NAs are skipped instead of propagating. + * assign_unit : If NULL, PyArray_InitializeReduceResult is used, otherwise + * this function is called to initialize the result to + * the reduction's unit. + * inner_loop : The inner loop which does the reduction. + * masked_inner_loop: The inner loop which does the reduction with a mask. + * data : Data which is passed to assign_unit and the inner loop. + * needs_api : Should be 1 if the inner loop might call a Python API + * function like PyErr_SetString. + * funcname : The name of the reduction function, for error messages. + */ +NPY_NO_EXPORT PyArrayObject * +PyArray_ReduceWrapper(PyArrayObject *operand, PyArrayObject *out, + PyArray_Descr *operand_dtype, + PyArray_Descr *result_dtype, + npy_bool *axis_flags, int skipna, + PyArray_AssignReduceUnitFunc *assign_unit, + PyArray_ReduceInnerLoopFunc *inner_loop, + PyArray_ReduceInnerLoopFunc *masked_inner_loop, + void *data, const char *funcname); + #endif |