Function TEST_P

test/topk.cpp:367–409

}

TEST_P(TopKParams, CPP) {
    topk_params params = GetParam();
    int d0             = params.d0;
    int d1             = params.d1;
    int k              = params.k;
    int dim            = params.dim;
    topkFunction order = params.order;

    array in = iota(dim4(d0, d1));

    // reverse the array if the order is ascending
    if (order == AF_TOPK_MIN) { in = -in + (d0 * d1 - 1); }
    array val, idx;
    topk(val, idx, in, k, dim, order);

    vector<float> hval(k * d1);
    vector<unsigned> hidx(k * d1);
    val.host(&hval[0]);
    idx.host(&hidx[0]);

    if (order == AF_TOPK_MIN) {
        for (int j = d1 - 1, i = 0; j > 0; j--) {
            for (int kidx = 0, goldidx = d0 - 1; kidx < k;
                 i++, kidx++, goldidx--) {
                float gold = static_cast<float>(j * d0 + kidx);
                ASSERT_FLOAT_EQ(gold, hval[i])
                    << print_context(i, kidx, hval, hidx);
                ASSERT_EQ(goldidx, hidx[i])
                    << print_context(i, kidx, hval, hidx);
            }
        }
    } else {
        for (int ii = 0, i = 0; ii < d1; ii++) {
            for (int j = d0 - 1; j >= d0 - k; --j, i++) {
                float gold  = static_cast<float>(ii * d0 + j);
                int goldidx = j;
                ASSERT_FLOAT_EQ(gold, hval[i])
                    << print_context(i, j, hval, hidx);
                ASSERT_EQ(goldidx, hidx[i]) << print_context(i, j, hval, hidx);
            }
        }
    }
}

TEST(TopK, KGreaterThan256) {
    af::array a = af::randu(500);