| 31 | #include "test.hpp" |
| 32 | |
| 33 | static auto run_prog(int64_t iter_num, bool cond, int64_t ini_val) |
| 34 | { |
| 35 | migraphx::shape si{migraphx::shape::int64_type}; |
| 36 | migraphx::shape s{migraphx::shape::int64_type, {1}}; |
| 37 | migraphx::shape sc{migraphx::shape::bool_type}; |
| 38 | |
| 39 | auto create_program = [&]() { |
| 40 | migraphx::program p; |
| 41 | auto* mm = p.get_main_module(); |
| 42 | |
| 43 | auto in_iter = mm->add_parameter("iter_num", si); |
| 44 | auto in_cond = mm->add_parameter("ccond", sc); |
| 45 | auto in_val = mm->add_parameter("val", s); |
| 46 | |
| 47 | auto* body = p.create_module("loop_module"); |
| 48 | auto iter = body->add_parameter("#loop_module_in_0", si); |
| 49 | body->add_parameter("#loop_module_in_1", sc); |
| 50 | auto in_v = body->add_parameter("#loop_module_in_2", s); |
| 51 | std::vector<int64_t> vd = {3}; |
| 52 | auto l = body->add_literal(migraphx::literal(si, vd)); |
| 53 | auto ad = body->add_instruction(migraphx::make_op("add"), iter, l); |
| 54 | auto val = body->add_instruction(migraphx::make_op("add"), in_v, ad); |
| 55 | auto eq = body->add_instruction(migraphx::make_op("equal"), iter, l); |
| 56 | auto beq = body->add_instruction( |
| 57 | migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), eq); |
| 58 | auto neq = body->add_instruction(migraphx::make_op("not"), beq); |
| 59 | body->add_return({neq, val, val}); |
| 60 | |
| 61 | auto rl = mm->add_instruction(migraphx::make_op("loop", {{"max_iterations", 10}}), |
| 62 | {in_iter, in_cond, in_val}, |
| 63 | {body}); |
| 64 | auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), rl); |
| 65 | auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), rl); |
| 66 | mm->add_return({r0, r1}); |
| 67 | |
| 68 | return p; |
| 69 | }; |
| 70 | |
| 71 | auto p = create_program(); |
| 72 | p.compile(migraphx::make_target("ref")); |
| 73 | migraphx::parameter_map pp; |
| 74 | pp["iter_num"] = migraphx::argument(si, &iter_num); |
| 75 | pp["ccond"] = migraphx::argument(sc, &cond); |
| 76 | pp["val"] = migraphx::argument(s, &ini_val); |
| 77 | auto rets = p.eval(pp); |
| 78 | |
| 79 | std::vector<std::vector<int64_t>> res; |
| 80 | for(auto& arg : rets) |
| 81 | { |
| 82 | std::vector<int64_t> vec; |
| 83 | arg.visit([&](auto v) { vec.assign(v.begin(), v.end()); }); |
| 84 | res.push_back(vec); |
| 85 | } |
| 86 | |
| 87 | return res; |
| 88 | } |
| 89 | |
| 90 | TEST_CASE(loop_test1) |