Function create_pointwise_module

src/targets/gpu/jit/mlir.cpp:45–71

MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_DUMP);

static module create_pointwise_module(module_ref in_mod)
{
    module pw_mod;
    std::unordered_map<instruction_ref, instruction_ref> map_ins;
    for(auto param : in_mod->get_parameters())
    {
        map_ins[param] =
            pw_mod.add_parameter(any_cast<builtin::param>(param->get_operator()).parameter,
                                 shape{param->get_shape().type()});
    }
    auto return_args =
        pw_mod.add_instructions(in_mod,
                                &map_ins,
                                [](module& m,
                                   instruction_ref ins,
                                   const operation& op,
                                   const std::vector<instruction_ref>& inputs,
                                   const std::vector<module_ref>& mod_args) -> instruction_ref {
                                    auto out_aliases = op.output_alias(to_shapes(inputs));
                                    if(out_aliases.size() == 1)
                                        return inputs.at(out_aliases[0]);
                                    else
                                        return m.insert_instruction(ins, op, inputs, mod_args);
                                });
    pw_mod.add_return(return_args);
    return pw_mod;
}

static code_object_op
compile_pointwise_module(context& ctx, const std::vector<shape>& inputs, module_ref mod)