Function run_program

test/ref/topk.cpp:33–82

#include <test.hpp>

static auto run_program(const migraphx::value& op, bool custom_idx = false, bool fill1 = false)
{
    migraphx::program p;
    auto* mm = p.get_main_module();
    migraphx::shape s{migraphx::shape::float_type, {3, 5}};
    auto data                                          = mm->add_parameter("data", s);
    std::vector<migraphx::instruction_ref> topk_inputs = {data};
    if(custom_idx)
    {
        migraphx::shape is{migraphx::shape::uint16_type, {3, 5}};
        std::vector<int16_t> indices(is.elements());
        std::iota(indices.rbegin(), indices.rend(), 1);
        topk_inputs.push_back(mm->add_literal(migraphx::literal{is, indices}));
    }
    auto r  = mm->add_instruction(migraphx::make_op("topk", op), topk_inputs);
    auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r);
    auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r);
    mm->add_return({r0, r1});

    p.compile(migraphx::make_target("ref"));

    std::vector<float> input_data = {
        2.1,
        2.3,
        2.0,
        2.5,
        1.9,
        3.3,
        0.2,
        4.5,
        0.1,
        0.8,
        1.0,
        4.5,
        2.1,
        0.8,
        1.5,
    };
    if(fill1)
        input_data.assign(input_data.size(), 1);
    migraphx::parameter_map pp;
    pp["data"] = migraphx::argument(s, input_data.data());
    auto rets  = p.eval(pp);
    std::vector<float> ret_val;
    rets.front().visit([&](auto v) { ret_val.assign(v.begin(), v.end()); });
    std::vector<int64_t> ret_ind;
    rets.back().visit([&](auto v) { ret_ind.assign(v.begin(), v.end()); });

    return std::make_pair(ret_val, ret_ind);
}

TEST_CASE(topk_largest0)
{