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

Code before:
// This file is part of OpenCV project.
// It is subject to the license terms in the LICENSE file found in the
// top-level directory of this distribution and at http://opencv.org/license.html.

#include "../precomp.hpp"
#include "layers_common.hpp"
#include <opencv2/dnn/shape_utils.hpp>
#include <cstdint>

namespace cv {
namespace dnn {

// ONNX NonZero operator
// Spec: https://onnx.ai/onnx/operators/onnx__NonZero.html
// Supported opsets: 9-18 (no attributes)

namespace {
static inline bool isNonZeroF16(uint16_t val)
{
    return (val & 0x7fff) != 0;
}

template <typename T, typename PT>
static size_t computeNonZeroCountsPerStripe(const T* data, size_t total, std::vector<size_t>& nzcounts, PT isNonZero)
{
    const int nstripes = 16;
    nzcounts.assign(nstripes, 0);
    if (total == 0)
        return 0;
    cv::parallel_for_(cv::Range(0, nstripes), [&](const cv::Range& range) {
        const size_t i0 = total * (size_t)range.start / (size_t)nstripes;
        const size_t i1 = total * (size_t)range.end   / (size_t)nstripes;
        size_t nz = 0;
        for (size_t i = i0; i < i1; ++i)
            nz += isNonZero(data[i]);
        nzcounts[(size_t)range.start] = nz;
    });

    size_t total_nz = 0;
    for (size_t c : nzcounts) total_nz += c;
    return total_nz;
}

template <typename T>
static inline size_t computeNonZeroCountsPerStripe(const T* data, size_t total, std::vector<size_t>& nzcounts)
{
    return computeNonZeroCountsPerStripe<T>(data, total, nzcounts, [](T v){ return v != (T)0; });
}

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

Review comment (suggestion):
pass nstripes as function parameter, to make sure that all stages of the algorithm use exactly the same number of stripes.

Code after:
// This file is part of OpenCV project.
// It is subject to the license terms in the LICENSE file found in the
// top-level directory of this distribution and at http://opencv.org/license.html.

#include "../precomp.hpp"
#include "layers_common.hpp"
#include <opencv2/dnn/shape_utils.hpp>
#include <cstdint>
#include <array>

namespace cv {
namespace dnn {

// ONNX NonZero operator
// Spec: https://onnx.ai/onnx/operators/onnx__NonZero.html
// Supported opsets: 9-18 (no attributes)

namespace {
static inline bool isNonZeroF16(uint16_t val)
{
    return (val & 0x7fff) != 0;
}

template <typename T, typename PT>
static size_t computeNonZeroCountsPerStripe(const T* data, size_t total, int nstripes, std::vector<size_t>& nzcounts, PT isNonZero)
{
    nzcounts.assign(nstripes, 0);
    if (total == 0)
        return 0;
    cv::parallel_for_(cv::Range(0, nstripes), [&](const cv::Range& range) {
        const size_t i0 = total * (size_t)range.start / (size_t)nstripes;
        const size_t i1 = total * (size_t)range.end   / (size_t)nstripes;
        size_t nz = 0;
        for (size_t i = i0; i < i1; ++i)
            nz += isNonZero(data[i]);
        nzcounts[(size_t)range.start] = nz;
    });

    size_t total_nz = 0;
    for (size_t c : nzcounts) total_nz += c;
    return total_nz;
}

template <typename T>
static inline size_t computeNonZeroCountsPerStripe(const T* data, size_t total, int nstripes, std::vector<size_t>& nzcounts)
{
    return computeNonZeroCountsPerStripe<T>(data, total, nstripes, nzcounts, [](T v){ return v != (T)0; });
}

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