Function TEST_CASE

test/ref/if.cpp:33–88

#include <test.hpp>

TEST_CASE(if_literal_test)
{
    auto create_program = [] {
        migraphx::program p;
        auto* mm = p.get_main_module();
        migraphx::shape cond_s{migraphx::shape::bool_type};
        auto cond = mm->add_parameter("cond", cond_s);

        migraphx::shape s{migraphx::shape::float_type, {5}};

        auto* then_mod           = p.create_module("If_0_if");
        std::vector<float> data1 = {1, 2, 3, 4, 5};
        auto l1                  = then_mod->add_literal(migraphx::literal(s, data1));
        then_mod->add_return({l1});

        auto* else_mod           = p.create_module("If_0_else");
        std::vector<float> data2 = {5, 4, 3, 2, 1};
        auto l2                  = else_mod->add_literal(migraphx::literal(s, data2));
        else_mod->add_return({l2});

        auto ret = mm->add_instruction(migraphx::make_op("if"), {cond}, {then_mod, else_mod});
        auto r   = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), ret);
        mm->add_return({r});

        return p;
    };

    auto run_prog = [&](bool cond) {
        auto p = create_program();
        p.compile(migraphx::make_target("ref"));
        std::vector<char> c_data = {static_cast<char>(cond)};
        migraphx::shape cs{migraphx::shape::bool_type};
        migraphx::parameter_map m;
        m["cond"] = migraphx::argument(cs, c_data.data());

        auto res = p.eval(m).back();
        std::vector<float> ret;
        res.visit([&](auto v) { ret.assign(v.begin(), v.end()); });

        return ret;
    };

    // then branch
    {
        std::vector<float> gold_ret = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
        auto ret                    = run_prog(true);
        EXPECT(gold_ret == ret);
    }

    // else branch
    {
        std::vector<float> gold_ret = {5.0f, 4.0f, 3.0f, 2.0f, 1.0f};
        auto ret                    = run_prog(false);
        EXPECT(gold_ret == ret);
    }
}

TEST_CASE(if_param_test)