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

Code before:
                CV_Assert(outs.size() == 2);
                outs[0].fit(outShape, input.type());
                outs[1].fit(outShape, CV_64S);
                output_value = outs[0];
                output_index = outs[1];
            } else if (kind == _InputArray::STD_VECTOR_UMAT) {
                std::vector<UMat>& uouts = outputs_arr.getUMatVecRef();
                CV_Assert(uouts.size() == 2);
                uouts[0].fit(outShape, input.type());
                uouts[1].fit(outShape, CV_64S);
                output_value = uouts[0].getMat(ACCESS_WRITE);
                output_index = uouts[1].getMat(ACCESS_WRITE);
            } else {
                // Fallback – recreate directly on provided Mat refs.
                output_value.create(outShape, input.type());
                output_index.create(outShape, CV_64S);
            }
        }

        if (largest) {
            switch (input.depth()) {
                case CV_8U: FindTopK<ComparatorGreater<uint8_t>, uint8_t>(input, output_value, output_index, normalized_axis); break;
                case CV_8S: FindTopK<ComparatorGreater<int8_t>, int8_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16U: FindTopK<ComparatorGreater<uint16_t>, uint16_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16S: FindTopK<ComparatorGreater<int16_t>, int16_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16F: FindTopK<ComparatorGreater<hfloat>, hfloat>(input, output_value, output_index, normalized_axis); break;
                case CV_32U: FindTopK<ComparatorGreater<unsigned>, unsigned>(input, output_value, output_index, normalized_axis); break;
                case CV_32S: FindTopK<ComparatorGreater<int>, int>(input, output_value, output_index, normalized_axis); break;
                case CV_32F: FindTopK<ComparatorGreater<float>, float>(input, output_value, output_index, normalized_axis); break;
                case CV_64U: FindTopK<ComparatorGreater<uint64_t>, uint64_t>(input, output_value, output_index, normalized_axis); break;
                case CV_64S: FindTopK<ComparatorGreater<int64_t>, int64_t>(input, output_value, output_index, normalized_axis); break;
                case CV_64F: FindTopK<ComparatorGreater<double>, double>(input, output_value, output_index, normalized_axis); break;
                default: CV_Error(Error::BadDepth, "Unsupported input data type");
            }
        } else {
            switch (input.depth()) {
                case CV_8U: FindTopK<ComparatorLess<uint8_t>, uint8_t>(input, output_value, output_index, normalized_axis); break;
                case CV_8S: FindTopK<ComparatorLess<int8_t>, int8_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16U: FindTopK<ComparatorLess<uint16_t>, uint16_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16S: FindTopK<ComparatorLess<int16_t>, int16_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16F: FindTopK<ComparatorLess<hfloat>, hfloat>(input, output_value, output_index, normalized_axis); break;
                case CV_32U: FindTopK<ComparatorLess<unsigned>, unsigned>(input, output_value, output_index, normalized_axis); break;
                case CV_32S: FindTopK<ComparatorLess<int>, int>(input, output_value, output_index, normalized_axis); break;
                case CV_32F: FindTopK<ComparatorLess<float>, float>(input, output_value, output_index, normalized_axis); break;
                case CV_64U: FindTopK<ComparatorLess<uint64_t>, uint64_t>(input, output_value, output_index, normalized_axis); break;
                case CV_64S: FindTopK<ComparatorLess<int64_t>, int64_t>(input, output_value, output_index, normalized_axis); break;
                case CV_64F: FindTopK<ComparatorLess<double>, double>(input, output_value, output_index, normalized_axis); break;
                default: CV_Error(Error::BadDepth, "Unsupported input data type");
            }
        }
    }

Review comment (performance):
1) add CV_16BF branch as well
2) sorting hfloat's on x86 and sorting bfloat's on any platform will be slow because there is no support in hardware for scalar arithmetics on such types. Add yet another parameter, WT (work type) to the template method, set it to the type itself for all types except for hfloat and bfloat. For hfloat and bfloat set it to float.
3) in the comment above I recommended you to cache sorted values into a temporary buffer. Use WT for elements of that buffer.

Code after:
            CV_Error(cv::Error::StsBadArg, cv::format("Unsupported output array kind: %d", kind));
        }

        switch (input.depth()) {
            case CV_8U:   FindTopK<uint8_t          >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_8S:   FindTopK<int8_t           >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_16U:  FindTopK<uint16_t         >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_16S:  FindTopK<int16_t          >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_16F:  FindTopK<hfloat, float    >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_16BF: FindTopK<bfloat, float    >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_32U:  FindTopK<uint32_t         >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_32S:  FindTopK<int32_t          >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_32F:  FindTopK<float            >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_64U:  FindTopK<uint64_t         >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_64S:  FindTopK<int64_t          >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_64F:  FindTopK<double           >(input, output_value, output_index, normalized_axis, kVal); break;
            default: CV_Error(Error::BadDepth, "Unsupported input data type");
        }
    }
};

Ptr<TopK2Layer> TopK2Layer::create(const LayerParams& params)
{
    return makePtr<TopK2LayerImpl>(params);
}

}} // namespace cv::dnn