Commit 43c7d307 authored by Paul's avatar Paul
Browse files

Add fuse mlir

parent cea63d59
#include <migraphx/gpu/fuse_mlir.hpp>
#include <migraphx/gpu/mlir.hpp>
#include <migraphx/matcher.hpp>
#include <migraphx/pass_manager.hpp>
#include <migraphx/make_op.hpp>
namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
struct module;
namespace gpu {
struct find_conv_pointwise
{
// Find a convolution followed by a pointwise operation.
auto matcher() const
{
auto convolution = match::skip(match::name("contiguous"))(match::name("convolution").bind("convolution"));
return match::name("pointwise")(match::any_of[match::inputs()](convolution.bind("x")));
}
void apply(module& m, match::matcher_result r) const
{
auto ins = r.result;
auto conv_ins = r.instructions["convolution"];
auto x_ins = r.instructions["x"]; // input after contiguous
auto pm = ins->module_inputs().front();
auto names = pm->get_parameter_names();
std::sort(names.begin(), names.end());
module mm{};
std::unordered_map<instruction_ref, instruction_ref> param_map;
auto x = mm.add_parameter("x"+std::to_string(names.size()), conv_ins->inputs().at(0)->get_shape());
auto w = mm.add_parameter("x"+std::to_string(names.size()+1), conv_ins->inputs().at(1)->get_shape());
auto conv = mm.add_instruction(conv_ins->get_operator(), {x, w});
std::transform(names.begin(), names.end(), ins->inputs().begin(), std::inserter(param_map, param_map.end()), [&](auto name, auto input) {
if (input == x_ins)
return std::make_pair(pm->get_parameter(name), conv);
return std::make_pair(pm->get_parameter(name), mm.add_parameter(name, input->get_shape()));
});
mm.add_return(mm.insert_module_instructions(mm.end(), pm, param_map));
auto inputs = ins->inputs();
inputs.insert(inputs.end(), conv_ins->inputs().begin(), conv_ins->inputs().end());
inputs.push_back(m.insert_instruction(ins, make_op("hip::allocate", {{"shape", to_value(ins->get_shape())}})));
auto mlir = insert_mlir(m, ins, mm, inputs);
m.replace_instruction(ins, mlir);
}
};
void fuse_mlir::apply(module& m) const
{
match::find_matches(m, find_conv_pointwise{});
}
}// namespace gpu
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment