Commit 3fb1b7f3 authored by Paul's avatar Paul
Browse files

Fix contiguous after splits

parent 14ac1aed
......@@ -69,22 +69,30 @@ static bool try_compute_shape(instruction_ref ins,
return try_compute_shape(ins, inputs, mods);
}
void eliminate_contiguous::apply(module& m) const
template<class F>
static void remove_contiguous(const std::string& op_name, module& m, F f)
{
auto last = std::prev(m.end());
for(auto ins : iterator_for(m))
{
// return instruction should have inputs with standard shape
if(ins->name() == "@return")
continue;
if (ins != last and ins->outputs().empty())
continue;
if (not f(ins))
continue;
// Make a copy so we can modify it while we iterate
auto args = ins->inputs();
auto new_args = args;
auto mod_args = ins->module_inputs();
for(auto arg : ins->inputs())
{
if(arg->name() == op_name)
{
if(arg->name() != op_name)
continue;
auto prev = arg->inputs().front();
replace(new_args, arg, prev);
if(try_compute_shape(ins, new_args, mod_args))
......@@ -101,7 +109,17 @@ void eliminate_contiguous::apply(module& m) const
}
}
}
}
}
void eliminate_contiguous::apply(module& m) const
{
// Skip contiguous from splits first
remove_contiguous(op_name, m, [](auto ins) {
if (ins->name() != "slice")
return true;
return (ins->inputs().front()->outputs().size() == 1);
});
remove_contiguous(op_name, m, [](auto) { return true; });
}
} // namespace MIGRAPHX_INLINE_NS
......
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