Function generic_masked_strided_loop

numpy/core/src/multiarray/array_method.c:859–902

* This function wraps a regular unmasked strided-loop as a * masked strided-loop, only calling the function for elements * where the mask is True. * * TODO: Reductions also use this code to implement masked reductions. * Before consolidating them, reductions had a special case for * broadcasts: when the mask stride was 0 the code does not check all * elements as `npy_memch

 *       masks are common enough.
 */
static int
generic_masked_strided_loop(PyArrayMethod_Context *context,
        char *const *data, const npy_intp *dimensions,
        const npy_intp *strides, NpyAuxData *_auxdata)
{
    _masked_stridedloop_data *auxdata = (_masked_stridedloop_data *)_auxdata;
    int nargs = auxdata->nargs;
    PyArrayMethod_StridedLoop *strided_loop = auxdata->unmasked_stridedloop;
    NpyAuxData *strided_loop_auxdata = auxdata->unmasked_auxdata;

    char **dataptrs = auxdata->dataptrs;
    memcpy(dataptrs, data, nargs * sizeof(char *));
    char *mask = data[nargs];
    npy_intp mask_stride = strides[nargs];

    npy_intp N = dimensions[0];
    /* Process the data as runs of unmasked values */
    do {
        Py_ssize_t subloopsize;

        /* Skip masked values */
        mask = npy_memchr(mask, 0, mask_stride, N, &subloopsize, 1);
        for (int i = 0; i < nargs; i++) {
            dataptrs[i] += subloopsize * strides[i];
        }
        N -= subloopsize;

        /* Process unmasked values */
        mask = npy_memchr(mask, 0, mask_stride, N, &subloopsize, 0);
        if (subloopsize > 0) {
            int res = strided_loop(context,
                    dataptrs, &subloopsize, strides, strided_loop_auxdata);
            if (res != 0) {
                return res;
            }
            for (int i = 0; i < nargs; i++) {
                dataptrs[i] += subloopsize * strides[i];
            }
            N -= subloopsize;
        }
    } while (N > 0);

    return 0;
}


/*