diff --git a/src/eliminate_concat.cpp b/src/eliminate_concat.cpp index 625629d4c5a..f98d4daac6a 100644 --- a/src/eliminate_concat.cpp +++ b/src/eliminate_concat.cpp @@ -29,16 +29,25 @@ #include #include #include +#include #include #include #include +#include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { namespace { +bool reshape_lazy_accepts_input_shape(instruction_ref ins, const shape& input_shape) +{ + if(ins->name() != "reshape_lazy") + return true; + return reshape_dims(input_shape, ins->get_shape().lens(), {.lazy = true}).has_value(); +} + struct concat_optimizer { module* m = nullptr; @@ -66,14 +75,60 @@ struct concat_optimizer return ins->name() == "allocate" or ins->name() == am.name(); } - bool need_copy(instruction_ref ins) const { return not is_allocation(get_output_alias(ins)); } + bool need_copy(instruction_ref ins, const shape& replacement_shape) const + { + auto alias = get_output_alias(ins); + return not is_allocation(alias) or + not replacement_supports_reshape_lazy(alias, replacement_shape); + } + + bool replacement_supports_reshape_lazy(instruction_ref replaced, + const shape& replacement_shape) const + { + std::unordered_map shapes = {{replaced, replacement_shape}}; + auto downstream_range = range(std::next(replaced), m->end()); + auto downstream_refs = iterator_for(downstream_range); + return std::all_of( + downstream_refs.begin(), downstream_refs.end(), [&](instruction_ref ins) { + bool changed = false; + auto input_shapes = to_shapes(ins->inputs()); + std::transform(ins->inputs().begin(), + ins->inputs().end(), + input_shapes.begin(), + [&](instruction_ref arg) { + auto it = shapes.find(arg); + if(it == shapes.end()) + return arg->get_shape(); + changed = true; + return it->second; + }); + if(not changed) + return true; + if(not reshape_lazy_accepts_input_shape(ins, input_shapes.front())) + return false; + shape output_shape; + try + { + output_shape = + ins->get_operator().compute_shape(input_shapes, ins->module_inputs()); + } + catch(const std::exception&) + { + return false; + } + if(output_shape != ins->get_shape()) + shapes.emplace(ins, std::move(output_shape)); + return true; + }); + } instruction_ref insert_copy(const operation& op, instruction_ref input, instruction_ref super) const { auto slice = m->insert_instruction(std::next(super), op, super); // If its packed then replace the allocation with the slice instead - if(not need_copy(input) and slice->get_shape().packed() and input->outputs().size() == 1) + if(not need_copy(input, slice->get_shape()) and slice->get_shape().packed() and + input->outputs().size() == 1) { m->replace_instruction(get_output_alias(input), slice); return input; @@ -142,9 +197,23 @@ void eliminate_concat::apply(module& m) const continue; auto lens = ins->inputs().front()->get_shape().lens(); std::size_t axis = tune_axis(lens.size(), concat_op->axis, concat_op->name()); - auto ncopies = std::count_if( + if(std::any_of( + ins->inputs().begin(), std::prev(ins->inputs().end()), [&](instruction_ref input) { + auto slice_shape = shape{input->get_shape().type(), + input->get_shape().lens(), + ins->get_shape().strides()}; + auto needs_copy = co.need_copy(input, slice_shape) or not slice_shape.packed() or + input->outputs().size() != 1; + return needs_copy and + not co.replacement_supports_reshape_lazy(input, slice_shape); + })) + continue; + auto ncopies = std::count_if( ins->inputs().begin(), std::prev(ins->inputs().end()), [&](instruction_ref input) { - if(co.need_copy(input)) + auto slice_shape = shape{input->get_shape().type(), + input->get_shape().lens(), + ins->get_shape().strides()}; + if(co.need_copy(input, slice_shape)) return true; if(is_packed(input, axis)) return false;