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, ...@@ -69,22 +69,30 @@ static bool try_compute_shape(instruction_ref ins,
return try_compute_shape(ins, inputs, mods); 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)) for(auto ins : iterator_for(m))
{ {
// return instruction should have inputs with standard shape // return instruction should have inputs with standard shape
if(ins->name() == "@return") if(ins->name() == "@return")
continue; 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 // Make a copy so we can modify it while we iterate
auto args = ins->inputs(); auto args = ins->inputs();
auto new_args = args; auto new_args = args;
auto mod_args = ins->module_inputs(); auto mod_args = ins->module_inputs();
for(auto arg : ins->inputs()) for(auto arg : ins->inputs())
{ {
if(arg->name() == op_name) if(arg->name() != op_name)
{ continue;
auto prev = arg->inputs().front(); auto prev = arg->inputs().front();
replace(new_args, arg, prev); replace(new_args, arg, prev);
if(try_compute_shape(ins, new_args, mod_args)) if(try_compute_shape(ins, new_args, mod_args))
...@@ -101,7 +109,17 @@ void eliminate_contiguous::apply(module& m) const ...@@ -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 } // 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