Function topkTest

test/topk.cpp:75–160

template<typename T>
void topkTest(const int ndims, const dim_t* dims, const unsigned k,
              const int dim, const af_topk_function order) {
    SUPPORTED_TYPE_CHECK(T);
    af_dtype dtype = (af_dtype)dtype_traits<T>::af_type;

    af_array input, output, outindex;

    size_t ielems = 1;
    size_t oelems = 1;

    for (int i = 0; i < ndims; i++) {
        ielems *= dims[i];
        oelems *= (i == dim ? k : dims[i]);
    }

    size_t bCount = ielems / dims[dim];
    size_t bSize  = dims[dim];

    vector<T> inData(ielems);
    T val{std::numeric_limits<T>::lowest()};
    generate(begin(inData), end(inData), [&]() {
        increment_next(val);
        return val;
    });

    random_device rnd_device;
    mt19937 g(rnd_device());
    shuffle(begin(inData), end(inData), g);

    vector<T> outData(oelems);
    vector<uint> outIdxs(oelems);

    for (size_t b = 0; b < bCount; b++) {
        using KeyValuePair = pair<T, uint>;

        vector<KeyValuePair> kvPairs;
        kvPairs.reserve(((b + 1) * bSize));

        for (size_t i = b * bSize; i < ((b + 1) * bSize); ++i)
            kvPairs.push_back(make_pair(inData[i], (i - b * bSize)));

        if (order & AF_TOPK_MIN) {
            stable_sort(kvPairs.begin(), kvPairs.end(),
                        [](const KeyValuePair& lhs, const KeyValuePair& rhs) {
                            return lhs.first < rhs.first;
                        });
        } else {
            stable_sort(kvPairs.begin(), kvPairs.end(),
                        [](const KeyValuePair& lhs, const KeyValuePair& rhs) {
                            return lhs.first > rhs.first;
                        });
        }

        auto it = kvPairs.begin();
        for (size_t i = 0; i < k; ++it, ++i) {
            outData[i + b * k] = it->first;
            outIdxs[i + b * k] = it->second;
        }