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

Code before:
                size_t r = lin;
                for (int d = 0; d < rank; ++d)
                {
                    int idx = static_cast<int>(r / strides[d]);
                    y[d * (size_t)nnz + col] = static_cast<int64>(idx);
                    r -= (size_t)idx * (size_t)strides[d];
                }
                ++col;
            };

            switch (depth)
            {
            case CV_32F: { const float* p = X.ptr<float>();
                for (size_t i = 0; i < total; ++i) { if (isNonZero<float>(p[i])) emit_idx(i); } break; }
            case CV_64F: { const double* p = X.ptr<double>();
                for (size_t i = 0; i < total; ++i) { if (isNonZero<double>(p[i])) emit_idx(i); } break; }
            case CV_32S: { const int* p = X.ptr<int>();
                for (size_t i = 0; i < total; ++i) { if (p[i] != 0) emit_idx(i); } break; }
            case CV_16F: { const hfloat* p = X.ptr<hfloat>();
                for (size_t i = 0; i < total; ++i) { if (isNonZero<hfloat, float>(p[i])) emit_idx(i); } break; }
            case CV_16BF:{ const bfloat* p = X.ptr<bfloat>();
                for (size_t i = 0; i < total; ++i) { if (isNonZero<bfloat, float>(p[i])) emit_idx(i); } break; }
            case CV_8U:
            case CV_Bool: { const uchar* p = X.ptr<uchar>();
                for (size_t i = 0; i < total; ++i) { if (p[i] != 0) emit_idx(i); } break; }
            default: break;
            }
        } 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);
            if (nnz == 0) return;

            Mat Y(outShape, CV_64S);
            int64* y = Y.ptr<int64>();
            size_t col = 0;

            auto emit_idx = [&](size_t lin) {
                size_t r = lin;
                for (int d = 0; d < rank; ++d)
                {
                    int idx = static_cast<int>(r / strides[d]);
                    y[d * (size_t)nnz + col] = static_cast<int64>(idx);
                    r -= (size_t)idx * (size_t)strides[d];
                }
                ++col;
            };

            switch (depth)
            {
            case CV_32F: { const float* p = X.ptr<float>();

Review comment (suggestion):
put this into a template function with parallel loop inside. Compute number of non-zero elements for each stripe.

After you computed number of non-zero elements in each stripe, compute cumsum:

```
constexpr int nstripes = 16;
std::vector<size_t> nzcounts(nstripes, 0);
std::vector<size_t> nzstart(nstripes+1, 0);

parallel_for_(Range(0, nstripes), [&](const Range& range) {
        size_t i0 = total*range.start/nstripes;
        size_t i1 = total*range.end/nstripes;
        size_t nz = 0;
        for (...) { nz += ... }
        nzcounts[range.start] = nz;
});

// nzstart = { 0, nzcounts[0], nzcounts[0] + nzcounts[1], nzcounts[0] + ... + nzcounts[2], ... }
for (int i = 1; i <= nstripes; i++) {
    nzstart[i] = nzstart[i-1] + nzcounts[i-1];
}

// allocate output of shape nzstart[nstripes] = nzcounts[0] + ... + nzcounts[nstripes-1]
...

parallel_for_(Range(0, nstripes), [&](const Range& range) {
        size_t i0 = total*range.start/nstripes;
        size_t i1 = total*range.end/nstripes;
        size_t nzcurrstart = nzstart[range.start];
        for (...) { if (nonZero(inp[i]) { outarray[nzcurrstart] = <index>; nzcurrstart++;} }
});


Code after:
    }

    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) {
            strides.back() = 1;
            for (int i = rank - 2; i >= 0; --i)
                strides[i] = strides[i + 1] * dims[i + 1];
        }

        size_t total = X.total();
        const int nstripes = 16;

        const int depth = CV_MAT_DEPTH(X.type());
        std::vector<size_t> nzcounts;
        std::vector<size_t> nzstart;
        size_t nnz = 0;
        switch (depth)
        {
        case CV_Bool:
        case CV_8U:
        case CV_8S:  { nnz = computeNonZeroCountsPerStripe<uchar>(X.ptr<uchar>(), total, nstripes, nzcounts); break; }
        case CV_16U:
        case CV_16S: { nnz = computeNonZeroCountsPerStripe<uint16_t>(X.ptr<uint16_t>(), total, nstripes, nzcounts); break; }
        case CV_32U:
        case CV_32S: { nnz = computeNonZeroCountsPerStripe<uint32_t>(X.ptr<uint32_t>(), total, nstripes, nzcounts); break; }
        case CV_64U:
        case CV_64S: { nnz = computeNonZeroCountsPerStripe<uint64_t>(X.ptr<uint64_t>(), total, nstripes, nzcounts); break; }
        case CV_32F: { nnz = computeNonZeroCountsPerStripe<float>(X.ptr<float>(), total, nstripes, nzcounts); break; }
        case CV_64F: { nnz = computeNonZeroCountsPerStripe<double>(X.ptr<double>(), total, nstripes, nzcounts); break; }
        case CV_16F:
        case CV_16BF:{ nnz = computeNonZeroCountsPerStripe<uint16_t>(X.ptr<uint16_t>(), total, nstripes, nzcounts, [](uint16_t v){ return isNonZeroF16(v); }); break; }
        default:
            CV_Error_(Error::StsError, ("NonZero: Unsupported input depth=%d", depth));
        }

        nzstart.assign(nzcounts.size() + 1, 0);
        for (size_t i = 1; i <= nzcounts.size(); ++i)
            nzstart[i] = nzstart[i - 1] + nzcounts[i - 1];

        MatShape outShape = shape(rank, (int)nnz);
        auto kind = out_arr.kind();
        std::vector<Mat>* out_mats = nullptr;