diff --git a/src/common/transformations/src/transformations/transpose_sinking/ts_gather.cpp b/src/common/transformations/src/transformations/transpose_sinking/ts_gather.cpp index d7609c935f1..42fa6f85cc6 100644 --- a/src/common/transformations/src/transformations/transpose_sinking/ts_gather.cpp +++ b/src/common/transformations/src/transformations/transpose_sinking/ts_gather.cpp @@ -57,16 +57,17 @@ TSGatherForward::TSGatherForward() { } } - size_t axis; + size_t order_axis; if (axes[0] < 0) { auto data_rank = main_node->get_input_partial_shape(0).rank(); if (data_rank.is_dynamic()) { return false; } - axis = static_cast(axes[0] + data_rank.get_length()); + order_axis = static_cast(axes[0] + data_rank.get_length()); } else { - axis = static_cast(axes[0]); + order_axis = static_cast(axes[0]); } + const size_t axis = order_val[order_axis]; /* https://docs.openvino.ai/2023.0/openvino_docs_ops_movement_Gather_8.html The Gather output shape has the same shape as the input, @@ -136,8 +137,7 @@ TSGatherForward::TSGatherForward() { if (!success) { return false; } - auto new_axis = - ov::op::v0::Constant::create(gather_axis->get_element_type(), gather_axis->get_shape(), {order_val[axis]}); + auto new_axis = ov::op::v0::Constant::create(gather_axis->get_element_type(), gather_axis->get_shape(), {axis}); main_node->input(2).replace_source_output(new_axis); copy_runtime_info(gather_axis, new_axis); diff --git a/src/common/transformations/tests/transpose_sinking/ts_gather_test.cpp b/src/common/transformations/tests/transpose_sinking/ts_gather_test.cpp index 34586bcdd81..2a410fb0181 100644 --- a/src/common/transformations/tests/transpose_sinking/ts_gather_test.cpp +++ b/src/common/transformations/tests/transpose_sinking/ts_gather_test.cpp @@ -63,6 +63,7 @@ auto wrapper = [](const TestCase& test_case) { struct GatherForwardArguments { OutputVector inputs_to_main; + std::function&, const OutputVector&)> create_input_transpose_to_main; Output new_Gather_first_input; AxisVector new_transpose_order; }; @@ -76,7 +77,7 @@ auto test_forward_gather = [](const GatherForwardArguments& test_arguments) { test_case.inputs_to_main = test_arguments.inputs_to_main; // Test model description: - test_case.model.preprocess_inputs_to_main = {{set_transpose_for}, {{0}}}; + test_case.model.preprocess_inputs_to_main = {{test_arguments.create_input_transpose_to_main}, {{0}}}; test_case.model.main_op = {CREATE_GATHER_FACTORY(Gather)}; test_case.model.model_template = create_model; @@ -104,24 +105,44 @@ auto test_forward_gather = [](const GatherForwardArguments& test_arguments) { return wrapper(test_case); }; +class SetTransposeWithOrder { +public: + SetTransposeWithOrder(const AxisVector& order) : _order(order) {} + OutputVector operator()(const vector& idxs, const OutputVector& out_vec) const { + return set_transpose_with_order(idxs, out_vec, _order); + } + +private: + const AxisVector _order; +}; + vector tests_arguments_fw{ {{{parameter(f32, {3, 4, 5, 6}), constant(i32, {2}, {0, 2}), constant(i32, {1}, {2})}}, + set_transpose_for, constant(i32, {1}, {1}), AxisVector{3, 2, 1, 0}}, {{parameter(f32, {2, 4}), constant(i32, {}, {0}), constant(i32, {1}, {1})}, + set_transpose_for, constant(i32, {1}, {0}), AxisVector{0}}, {{parameter(f32, {2, 4}), constant(i32, {1}, {0}), constant(i32, {1}, {1})}, + set_transpose_for, constant(i32, {1}, {0}), AxisVector{1, 0}}, {{parameter(f32, {2, 3, 4}), constant(i32, {2, 3}, {0, 1, 0, 1, 0, 1}), constant(i32, {1}, {1})}, + set_transpose_for, constant(i32, {1}, {1}), - AxisVector{3, 1, 2, 0}}}; + AxisVector{3, 1, 2, 0}}, + {{parameter(f32, {64, 49, 3, 3, 32}), constant(i32, {}, {1}), constant(i32, {}, {0})}, + SetTransposeWithOrder(AxisVector{2, 0, 3, 1, 4}), + constant(i32, {}, {2}), + AxisVector{0, 2, 1, 3}}}; INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_0, TSTestFixture, test_forward_gather(tests_arguments_fw[0])); INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_1, TSTestFixture, test_forward_gather(tests_arguments_fw[1])); INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_2, TSTestFixture, test_forward_gather(tests_arguments_fw[2])); INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_3, TSTestFixture, test_forward_gather(tests_arguments_fw[3])); +INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_4, TSTestFixture, test_forward_gather(tests_arguments_fw[4])); struct GatherBackwardArguments { OutputVector inputs_to_main; diff --git a/src/common/transformations/tests/transpose_sinking/ts_test_utils.cpp b/src/common/transformations/tests/transpose_sinking/ts_test_utils.cpp index d0a81c5dd40..31be099f03a 100644 --- a/src/common/transformations/tests/transpose_sinking/ts_test_utils.cpp +++ b/src/common/transformations/tests/transpose_sinking/ts_test_utils.cpp @@ -52,6 +52,19 @@ OutputVector set_transpose_for(const vector& idxs, const OutputVector& o return result; } +OutputVector set_transpose_with_order(const vector& idxs, + const OutputVector& out_vec, + const vector& transpose_order_axes) { + OutputVector result = out_vec; + for (const auto& idx : idxs) { + const auto& out = out_vec[idx]; + auto order = make_shared(element::i32, Shape{transpose_order_axes.size()}, transpose_order_axes); + auto transpose = make_shared(out, order); + result[idx] = transpose; + } + return result; +} + OutputVector set_gather_for(const vector& idxs, const OutputVector& out_vec) { OutputVector result = out_vec; for (const auto& idx : idxs) { diff --git a/src/common/transformations/tests/transpose_sinking/ts_test_utils.hpp b/src/common/transformations/tests/transpose_sinking/ts_test_utils.hpp index 826d75e0b1b..cc8685c1d7e 100644 --- a/src/common/transformations/tests/transpose_sinking/ts_test_utils.hpp +++ b/src/common/transformations/tests/transpose_sinking/ts_test_utils.hpp @@ -21,6 +21,9 @@ std::string to_string(const ov::Shape& shape); ov::ParameterVector filter_parameters(const ov::OutputVector& out_vec); ov::OutputVector set_transpose_for(const std::vector& idxs, const ov::OutputVector& out_vec); +ov::OutputVector set_transpose_with_order(const std::vector& idxs, + const ov::OutputVector& out_vec, + const std::vector& transpose_order_axes); ov::OutputVector set_gather_for(const std::vector& idxs, const ov::OutputVector& out_vec); std::shared_ptr create_main_node(const ov::OutputVector& inputs, size_t num_ops, const FactoryPtr& creator);