Method apply

src/split_single_dyn_dim.cpp:118–171

* Makes all the shapes in the dynamic_dimension range. Probably won't work for `if` * and `loop` instructions, depending on how the submodules for those * work. Inserts select_module instruction to the top. Replaces return, bypassing other * instructions. Skips if the dynamic parameter outputs to a select_module operator. */

 * instructions. Skips if the dynamic parameter outputs to a select_module operator.
 */
void split_single_dyn_dim::apply(module_pass_manager& mpm) const
{
    module_ref mm                               = &mpm.get_module();
    auto param_names                            = mm->get_parameter_names();
    auto param_shapes                           = mm->get_parameter_shapes();
    optional<std::vector<dynamic_dimensions_check>> dd_check_vec =
        has_one_unique_dyn_dim(param_shapes);
    if(dd_check_vec.has_value() and not any_sm_next(mm, dd_check_vec.value()))
    {
        // all dynamic dimension objects should be the same for all parameters in dd_check_vec
        auto dyn_dim = dd_check_vec->at(0).dd;
        // create submodules for each dimension size
        std::vector<module_ref> submodules;
        for(size_t dim_size : migraphx::range(dyn_dim.min, dyn_dim.max + 1))
        {
            auto* submod = mpm.create_module("dim_" + std::to_string(dim_size));
            // instruction map for new static shaped submodule parameters
            std::unordered_map<instruction_ref, instruction_ref> map_ins;
            for(const auto& dd_check : dd_check_vec.value())
            {
                // create static shape using dim_size
                const auto& dyn_param = mm->get_parameter(dd_check.dyn_param_str);
                auto dyn_param_shape  = mm->get_parameter_shape(dd_check.dyn_param_str);
                auto static_shape     = dyn_param_shape.to_static(dim_size);
                map_ins[dyn_param]    = submod->add_parameter(dd_check.dyn_param_str, static_shape);
            }
            auto outputs = submod->add_instructions(mm, &map_ins);
            submod->add_return({outputs});
            submodules.push_back(submod);
        }
        // sort parameters by name for consistency (vs. parameter order attr)
        std::sort(param_names.begin(), param_names.end());
        // redirect to select_module operator and return
        std::vector<instruction_ref> sm_inputs;
        std::transform(param_names.cbegin(),
                       param_names.cend(),
                       std::back_inserter(sm_inputs),
                       [&](auto pn) { return mm->get_parameter(std::move(pn)); });
        auto output_shapes       = mm->get_output_shapes();
        migraphx::shape out_attr = migraphx::shape{output_shapes};
        auto sm_ins              = mm->add_instruction(
            migraphx::make_op("select_module",
                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
            sm_inputs,
            submodules);
        std::vector<instruction_ref> outputs(output_shapes.size());
        for(size_t i = 0; i < output_shapes.size(); ++i)
        {
            outputs.at(i) =
                mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", i}}), sm_ins);
        }
        mm->replace_return(outputs);
    }
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx