diff --git a/src/common/transformations/src/transformations/common_optimizations/simplify_shape_of_sub_graph.cpp b/src/common/transformations/src/transformations/common_optimizations/simplify_shape_of_sub_graph.cpp index 7facf950ee7..429d2c0ccd5 100644 --- a/src/common/transformations/src/transformations/common_optimizations/simplify_shape_of_sub_graph.cpp +++ b/src/common/transformations/src/transformations/common_optimizations/simplify_shape_of_sub_graph.cpp @@ -36,7 +36,7 @@ pass::GroupedGatherElimination::GroupedGatherElimination() { matcher_pass_callback callback = [=](Matcher& m) { auto concat = m.get_match_root(); OutputVector inputs = concat->input_values(); - NodeVector new_ops; + NodeRegistry new_ops; size_t i = 0, original_inputs_size = inputs.size(); while (inputs.size() > i + 1) { auto curr = inputs[i].get_node_shared_ptr(), next = inputs[i + 1].get_node_shared_ptr(); @@ -47,16 +47,36 @@ pass::GroupedGatherElimination::GroupedGatherElimination() { continue; } - // Scalar inputs are not supported by Concat and we don't want to throw an exception here. + auto curr_indices = curr->input_value(1); + auto next_indices = next->input_value(1); + // Scalar or not same type inputs are not supported by Concat and we don't want to throw an exception here. // The transformation should not be applied instead. - if (curr->input_value(1).get_partial_shape().same_scheme(Shape{}) || - next->input_value(1).get_partial_shape().same_scheme(Shape{})) { + if (curr_indices.get_partial_shape().same_scheme(Shape{}) || + next_indices.get_partial_shape().same_scheme(Shape{})) { return false; } + const auto& curr_et = curr_indices.get_element_type(); + const auto& next_et = next_indices.get_element_type(); + + if (!curr_et.compatible(next_et)) { + if (curr_et == element::u64 || next_et == element::u64) { + // Not apply transformation as conversion can change values interpretation + return false; + } else { + constexpr auto common_type = element::i64; + if (curr_et != common_type) { + curr_indices = new_ops.make(curr_indices, common_type); + } + if (next_et != common_type) { + next_indices = new_ops.make(next_indices, common_type); + } + } + } + // curr and next are the same type of gather which takes data from the same source auto joint_indices = - op::util::make_try_fold(OutputVector{curr->input_value(1), next->input_value(1)}, 0); + new_ops.add(op::util::make_try_fold(OutputVector{curr_indices, next_indices}, 0)); std::shared_ptr new_gather; if (is_type(curr)) { new_gather = register_new_node(curr->input_value(0), @@ -73,12 +93,12 @@ pass::GroupedGatherElimination::GroupedGatherElimination() { } else { OPENVINO_THROW("Unexpected Gather version"); } - new_ops.push_back(joint_indices); - new_ops.push_back(new_gather); + + new_ops.add(new_gather); inputs.erase(inputs.begin() + i); inputs[i] = new_gather->output(0); } - ov::copy_runtime_info(concat, new_ops); + ov::copy_runtime_info(concat, new_ops.get()); if (inputs.size() == 1) // we can optimize out concat return replace_output_update_name(concat->output(0), inputs[0]); if (original_inputs_size > inputs.size()) { diff --git a/src/common/transformations/tests/common_optimizations/simplify_shape_of_sub_graph.cpp b/src/common/transformations/tests/common_optimizations/simplify_shape_of_sub_graph.cpp index 7c09340b375..fa7abfab9ca 100644 --- a/src/common/transformations/tests/common_optimizations/simplify_shape_of_sub_graph.cpp +++ b/src/common/transformations/tests/common_optimizations/simplify_shape_of_sub_graph.cpp @@ -196,3 +196,179 @@ TEST_F(TransformationTestsF, GroupedGatherEliminationNegative) { manager.register_pass(); } } + +TEST_F(TransformationTestsF, GroupedGatherEliminationNotCompatibleIndiciesSameSign) { + PartialShape data_shape{2, 128}; + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + + auto indices_1 = opset7::Constant::create(element::i32, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::i64, {1}, {1}); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + auto gather_1 = std::make_shared(shape_op, indices_1, axis); + auto gather_2 = std::make_shared(shape_op, indices_2, axis); + + auto concat = std::make_shared(OutputVector{gather_1, gather_2}, 0); + + model = std::make_shared(NodeVector{concat}, ParameterVector{data}); + manager.register_pass(); + } + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + + auto indices_1 = opset7::Constant::create(element::i32, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::i64, {1}, {1}); + auto convert_1 = std::make_shared(indices_1, element::i64); + auto joint_indices = std::make_shared(OutputVector{convert_1, indices_2}, 0); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + + auto gather = std::make_shared(shape_op, joint_indices, axis); + + model_ref = std::make_shared(NodeVector{gather}, ParameterVector{data}); + } +} + +TEST_F(TransformationTestsF, GroupedGatherEliminationNotCompatibleIndiciesCanConvert_1) { + PartialShape data_shape{2, 128}; + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + + auto indices_1 = opset7::Constant::create(element::u32, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::i64, {1}, {1}); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + auto gather_1 = std::make_shared(shape_op, indices_1, axis); + auto gather_2 = std::make_shared(shape_op, indices_2, axis); + + auto concat = std::make_shared(OutputVector{gather_1, gather_2}, 0); + + model = std::make_shared(NodeVector{concat}, ParameterVector{data}); + manager.register_pass(); + } + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + + auto indices_1 = opset7::Constant::create(element::u32, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::i64, {1}, {1}); + auto convert_1 = std::make_shared(indices_1, element::i64); + auto joint_indices = std::make_shared(OutputVector{convert_1, indices_2}, 0); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + + auto gather = std::make_shared(shape_op, joint_indices, axis); + + model_ref = std::make_shared(NodeVector{gather}, ParameterVector{data}); + } +} + +TEST_F(TransformationTestsF, GroupedGatherEliminationNotCompatibleIndiciesCanConvert_2) { + PartialShape data_shape{2, 128}; + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + + auto indices_1 = opset7::Constant::create(element::u32, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::i32, {1}, {1}); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + auto gather_1 = std::make_shared(shape_op, indices_1, axis); + auto gather_2 = std::make_shared(shape_op, indices_2, axis); + + auto concat = std::make_shared(OutputVector{gather_1, gather_2}, 0); + + model = std::make_shared(NodeVector{concat}, ParameterVector{data}); + manager.register_pass(); + } + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + auto indices_1 = opset7::Constant::create(element::u32, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::i32, {1}, {1}); + auto convert_1 = std::make_shared(indices_1, element::i64); + auto convert_2 = std::make_shared(indices_2, element::i64); + auto joint_indices = std::make_shared(OutputVector{convert_1, convert_2}, 0); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + + auto gather = std::make_shared(shape_op, joint_indices, axis); + + model_ref = std::make_shared(NodeVector{gather}, ParameterVector{data}); + } +} + +TEST_F(TransformationTestsF, GroupedGatherEliminationCompatibleIndicies_1) { + PartialShape data_shape{2, 128}; + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + + auto indices_1 = opset7::Constant::create(element::i64, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::i64, {1}, {1}); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + auto gather_1 = std::make_shared(shape_op, indices_1, axis); + auto gather_2 = std::make_shared(shape_op, indices_2, axis); + + auto concat = std::make_shared(OutputVector{gather_1, gather_2}, 0); + + model = std::make_shared(NodeVector{concat}, ParameterVector{data}); + manager.register_pass(); + } + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + auto joint_indices = opset7::Constant::create(element::i64, {2}, {1, 1}); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + + auto gather = std::make_shared(shape_op, joint_indices, axis); + + model_ref = std::make_shared(NodeVector{gather}, ParameterVector{data}); + } +} + +TEST_F(TransformationTestsF, GroupedGatherEliminationCompatibleIndicies_2) { + PartialShape data_shape{2, 128}; + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + + auto indices_1 = opset7::Constant::create(element::u64, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::u64, {1}, {1}); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + auto gather_1 = std::make_shared(shape_op, indices_1, axis); + auto gather_2 = std::make_shared(shape_op, indices_2, axis); + + auto concat = std::make_shared(OutputVector{gather_1, gather_2}, 0); + + model = std::make_shared(NodeVector{concat}, ParameterVector{data}); + manager.register_pass(); + } + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + auto joint_indices = opset7::Constant::create(element::u64, {2}, {1, 1}); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + + auto gather = std::make_shared(shape_op, joint_indices, axis); + + model_ref = std::make_shared(NodeVector{gather}, ParameterVector{data}); + } +} + +TEST_F(TransformationTestsF, GroupedGatherEliminationNotCompatibleIndiciesCannotConvert) { + PartialShape data_shape{2, 128}; + { + auto data = std::make_shared(element::f32, data_shape); + auto shape_op = std::make_shared(data); + + auto indices_1 = opset7::Constant::create(element::u64, {1}, {1}); + auto indices_2 = opset7::Constant::create(element::i64, {1}, {1}); + auto axis = opset7::Constant::create(element::i64, {}, {0}); + auto gather_1 = std::make_shared(shape_op, indices_1, axis); + auto gather_2 = std::make_shared(shape_op, indices_2, axis); + + auto concat = std::make_shared(OutputVector{gather_1, gather_2}, 0); + + model = std::make_shared(NodeVector{concat}, ParameterVector{data}); + manager.register_pass(); + } +}