Function count_nonzero_int

numpy/core/src/multiarray/item_selection.c:2585–2628

* Counts the number of True values in a raw boolean array. This * is a low-overhead function which does no heap allocations. * * Returns -1 on error. */

 * Returns -1 on error.
 */
static NPY_GCC_OPT_3 npy_intp
count_nonzero_int(int ndim, char *data, const npy_intp *ashape, const npy_intp *astrides, int elsize)
{
    assert(elsize <= 8);
    int idim;
    npy_intp shape[NPY_MAXDIMS], strides[NPY_MAXDIMS];
    npy_intp coord[NPY_MAXDIMS];

    // Use raw iteration with no heap memory allocation
    if (PyArray_PrepareOneRawArrayIter(
                    ndim, ashape,
                    data, astrides,
                    &ndim, shape,
                    &data, strides) < 0) {
        return -1;
    }

    // Handle zero-sized array
    if (shape[0] == 0) {
        return 0;
    }

    NPY_BEGIN_THREADS_DEF;
    NPY_BEGIN_THREADS_THRESHOLDED(shape[0]);

    #define NONZERO_CASE(LEN, SFX) \
        case LEN: \
            NPY_RAW_ITER_START(idim, ndim, coord, shape) { \
                count += count_nonzero_##SFX(data, strides[0], shape[0]); \
            } NPY_RAW_ITER_ONE_NEXT(idim, ndim, coord, shape, data, strides); \
            break

    npy_intp count = 0;
    switch(elsize) {
        NONZERO_CASE(1, u8);
        NONZERO_CASE(2, u16);
        NONZERO_CASE(4, u32);
        NONZERO_CASE(8, u64);
    }
    #undef NONZERO_CASE

    NPY_END_THREADS;
    return count;
}
/*
 * Counts the number of True values in a raw boolean array. This
 * is a low-overhead function which does no heap allocations.