Repository: opencv/opencv
File: modules/dnn/src/layers/nonzero_layer.cpp (C++)

Code before:
        outputs.assign(1, out);
        return false;
    }

    void getTypes(const std::vector<MatType>& /*inputs*/, const int requiredOutputs,
                  const int requiredInternals, std::vector<MatType>& outputs,
                  std::vector<MatType>& internals) const CV_OVERRIDE
    {
        outputs.assign(requiredOutputs, CV_64S);
    }

    void forward(InputArrayOfArrays in_arr, OutputArrayOfArrays out_arr, OutputArrayOfArrays) CV_OVERRIDE
    {
        CV_Assert(in_arr.size().area() == 1);
        const Mat X = in_arr.getMat(0);
        CV_Assert(X.data != nullptr);

        const int rank = X.dims;
        std::vector<int> dims(rank), strides(rank);
        for (int i = 0; i < rank; ++i) dims[i] = X.size[i];
        if (rank == 0) {
            const bool nz = X.type() == CV_32F ? isNonZero<float>(X.at<float>(0))
                          : X.type() == CV_64F ? isNonZero<double>(X.at<double>(0))
                          : X.type() == CV_32S ? (X.at<int>(0) != 0)
                          : X.type() == CV_16F ? isNonZero<hfloat, float>(X.at<hfloat>(0))
                          : X.type() == CV_16BF ? isNonZero<bfloat, float>(X.at<bfloat>(0))
                          : X.at<uchar>(0) != 0;
            MatShape outShape = shape(0, nz ? 1 : 0);
            auto kind = out_arr.kind();
            if (kind == _InputArray::STD_VECTOR_MAT) {
                std::vector<Mat>& outs_m = out_arr.getMatVecRef();
                outs_m.resize(1);
                outs_m[0].fit(outShape, CV_64S);
            } else if (kind == _InputArray::STD_VECTOR_UMAT) {
                std::vector<UMat>& outs_u = out_arr.getUMatVecRef();
                outs_u.resize(1);
                outs_u[0].fit(outShape, CV_64S);
            } else {
                CV_Error(Error::StsNotImplemented, "NonZero: unsupported OutputArrayOfArrays kind");
            }
            return;
        }

        strides.back() = 1;
        for (int i = rank - 2; i >= 0; --i)
            strides[i] = strides[i + 1] * dims[i + 1];

        size_t total = X.total();
        size_t nnz = 0;

        const int depth = CV_MAT_DEPTH(X.type());

Review comment (suggestion):
hfloat/bfloat should be handled with custom implementations of isNonZero: they should just check for 0x0/0x8000 case (one function is enough for both hfloat and bfloat):

```
static inline isNonZeroF16(uint16_t val) { return (val&0x7fff) != 0; }

...

X.type() == CV_16F || X.type() == CV_16BF ? isNonZeroF16(*X.ptr<uint16_t>())
```

Since hfloat/bfloat is the only case that requires `WT!=T` and this case can be solved as shown above, WT parameter in isNonZero is not needed anymore. Maybe the whole function is not needed either, just put the expression (`x!=0`) in-place.

also, all other types: uint8/int8, uint16/int16, uint32/int32, uint64/int64 must be supported. It's enough to have 1 

Code after:

template <typename T, typename PT>
static void emitIndicesStripes(const T* data,
                               size_t total,
                               int nstripes,
                               int rank,
                               const std::vector<int>& dims,
                               const std::vector<int>& strides,
                               const std::vector<size_t>& nzstart,
                               int64* y,
                               PT isNonZero)
{
    const size_t nnz_total = nzstart.back();
    if (rank == 0) {
        return;
    }
    const int lastDimSize = dims[rank - 1];

    cv::parallel_for_(cv::Range(0, nstripes), [&](const cv::Range& range) {
        size_t i = total * (size_t)range.start / (size_t)nstripes;
        const size_t i1 = total * (size_t)range.end / (size_t)nstripes;
        size_t col = nzstart[(size_t)range.start];

        std::array<int, CV_MAX_DIM> coord;
        std::fill_n(coord.begin(), rank, 0);
        if (rank > 0) coord[rank - 1] = lastDimSize - 1;

        for (; i < i1; ++i)
        {
            if (rank > 0 && ++coord[rank - 1] >= lastDimSize)
            {
                size_t idxLinear = i;
                for (int d = rank - 1; d >= 0; --d)
                {
                    coord[d] = static_cast<int>(idxLinear % (size_t)dims[d]);
                    idxLinear /= (size_t)dims[d];
                }
            }
            if (isNonZero(data[i]))
            {
                for (int d = 0; d < rank; ++d)
                    y[(size_t)d * nnz_total + col] = (int64)coord[d];
                ++col;
            }
        }
    });
}

template <typename T>
static void emitIndicesStripes(const T* data,
                               size_t total,