| 133 | }; |
| 134 | |
| 135 | migraphx::argument |
| 136 | compute(migraphx::context& ctx, |
| 137 | const migraphx::shape& out_shape, |
| 138 | const std::vector<migraphx::argument>& args, |
| 139 | const std::vector<migraphx::module_ref>& mods, |
| 140 | const std::function<std::vector<migraphx::argument>( |
| 141 | migraphx::module_ref&, const std::unordered_map<std::string, migraphx::argument>&)>& |
| 142 | run) const |
| 143 | { |
| 144 | // wrap up the arguments vector, so ref and gpu impl are the same |
| 145 | auto cpy_args = args; |
| 146 | bool in_cond = args.at(1).at<bool>(); |
| 147 | bool cond = in_cond; |
| 148 | int64_t iter = 0; |
| 149 | // insert iter and cond used in the loop |
| 150 | auto s_cond = args.at(1).get_shape(); |
| 151 | auto s_iter = args.at(0).get_shape(); |
| 152 | cpy_args.push_back({s_iter, &iter}); |
| 153 | cpy_args.push_back({s_cond, &cond}); |
| 154 | cpy_args.insert(cpy_args.end(), args.begin() + 2, args.end()); |
| 155 | // add cond and mod outputs to the argument list |
| 156 | cpy_args.push_back(migraphx::argument(s_cond)); |
| 157 | cpy_args.push_back(migraphx::argument(out_shape)); |
| 158 | // run loop |
| 159 | return run_loop(test_loop{max_iterations}, {}, ctx, cpy_args, mods, run); |
| 160 | } |
| 161 | }; |
| 162 | |
| 163 | static auto create_program(int64_t max_loop_iterations = 10) |