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

Code before:
    {
        outputs.resize(2);
        outputs[0] = inputs.front();
        outputs[1] = CV_64S;
    }

private:
    template<class Comparator, typename T>
    void FindTopK(const Mat &input, Mat &output_vals, Mat &output_idxs, int normalized_axis)
    {
        const auto inShape = shape(input);
        size_t outer = std::accumulate(inShape.begin(), inShape.begin() + normalized_axis, 1, std::multiplies<int>());
        size_t inner = std::accumulate(inShape.begin() + normalized_axis + 1, inShape.end(), 1, std::multiplies<int>());
        int dimAxis = inShape[normalized_axis];

        auto worker = [&](const Range &r) {
            const T* inPtr = input.ptr<T>();
            T* valPtr = output_vals.ptr<T>();
            int64_t* idxPtr = output_idxs.ptr<int64_t>();
            Comparator cmp(inPtr, inner);
            AutoBuffer<size_t> indices(dimAxis);
            size_t* idxBuf = indices.data();
            for (size_t b = r.start; b < r.end; ++b) {
                for (size_t j = 0; j < inner; ++j) {
                    size_t offset = b * dimAxis * inner + j;
                    cmp.addOffset(offset);
                    std::iota(idxBuf, idxBuf + dimAxis, 0);
                    std::stable_sort(idxBuf, idxBuf + dimAxis, cmp);
                    // By the time we reach here, 'K' has been resolved. Use it directly.
                    int currentK = K;
                    for (int i = 0; i < currentK; ++i) {
                        size_t src = idxBuf[i];
                        valPtr[b * currentK * inner + i * inner + j] = inPtr[offset + src * inner];
                        idxPtr[b * currentK * inner + i * inner + j] = static_cast<int64_t>(src);
                    }
                    cmp.minusOffset(offset);
                }
            }
        };
        parallel_for_(Range(0, static_cast<int>(outer)), worker);
    }

public:
    void forward(InputArrayOfArrays inputs_arr,
                 OutputArrayOfArrays outputs_arr,
                 OutputArrayOfArrays internals_arr) CV_OVERRIDE
    {
        CV_TRACE_FUNCTION();
        std::vector<Mat> inputs, outputs;
        inputs_arr.getMatVector(inputs);
        outputs_arr.getMatVector(outputs);

Review comment (bug):
this is very suspicious place. Subsequent offset values inside loop form arithmetic progression (for fixed b). You add this offset to data inside cmp.addOffset(), so data addresses form geometric progression. It's wrong, they should form arithmetic progression as well. That's why I suggest you to remove addOffset and minusOffset and pass the proper data pointer to comparator constructor, which you need to place inside the inner loop.

Code after:
            for (size_t b = r.start; b < r.end; ++b) {
                for (size_t j = 0; j < inner; ++j) {
                    size_t offset = b * dimAxis * inner + j;
                    for (uint32_t u = 0; u < (uint32_t)dimAxis; ++u) {
                        sortbuf[u].first  = u;
                        sortbuf[u].second = WT(inPtr[offset + u * inner]);
                    }
                    if (largest){
                        ComparatorGreater<uint32_t,WT> cmp;
                        std::partial_sort(sortbuf.begin(), sortbuf.begin() + kVal, sortbuf.end(), cmp);
                    }
                    else{
                        ComparatorLess<uint32_t,WT> cmp;
                        std::partial_sort(sortbuf.begin(), sortbuf.begin() + kVal, sortbuf.end(), cmp);
                    }
                    for (int i = 0; i < kVal; ++i) {
                        auto &p = sortbuf[i];
                        valPtr[b * kVal * inner + i * inner + j] = T(p.second);
                        idxPtr[b * kVal * inner + i * inner + j] = p.first;
                    }
                }
            }
        };
        parallel_for_(Range(0, static_cast<int>(outer)), worker);
    }

public:
    void forward(InputArrayOfArrays inputs_arr,
                 OutputArrayOfArrays outputs_arr,
                 OutputArrayOfArrays internals_arr) CV_OVERRIDE
    {
        CV_TRACE_FUNCTION();
        std::vector<Mat> inputs, outputs;
        inputs_arr.getMatVector(inputs);
        outputs_arr.getMatVector(outputs);

        const auto &input = inputs.front();
        Mat output_value, output_index;

        // Normalize axis to handle negative values
        int normalized_axis = normalize_axis(axis, input.dims);

        int kVal = K;
        if (dynamicK) {
            CV_Assert(inputs.size() == 2);
            CV_Assert(inputs[1].type() == CV_64S);
            CV_Assert(inputs[1].total() == 1);
            int64_t kTemp = inputs[1].at<int64_t>(0);
            CV_CheckGT(kTemp, 0, "TopK2: dynamic K must be > 0");
            CV_CheckLT(kTemp, input.size[normalized_axis] + 1, "TopK2: dynamic K is out of range");
            kVal = static_cast<int>(kTemp);