[GNA] Fixed handling of transposes around MatMul (#8600)

* [GNA] Fixed handling of transposes around MatMul

* [GNA] Fixed swap matmul inputs tests

* [GNA] Comments applying

inference-engine/src/gna_plugin/backend/gna_limitations.hpp

@@ -31,10 +31,14 @@ constexpr uint32_t maxPoolMaxWindowSize = 6;
 constexpr uint32_t copyMaxGrouping = 8;
 constexpr uint32_t transposeMaxSize = 65528;
 
+inline bool IsTranspose2d(const std::vector<size_t>& shape) {
+    return std::count_if(std::begin(shape), std::end(shape), [](size_t dim) { return dim != 1; }) == 2;
+}
+
 inline bool IsTransposeSupported(const std::vector<size_t>& shape) {
+    if (!IsTranspose2d(shape)) return false;
     auto shape_no_1 = shape;
     shape_no_1.erase(std::remove(shape_no_1.begin(), shape_no_1.end(), 1), shape_no_1.end());
-    if (shape_no_1.size() != 2) return false;
     size_t min, max;
     std::tie(min, max) = std::minmax(shape_no_1[0], shape_no_1[1]);
     return min <= 8 && max % 8 == 0 && max >= 8 && max <= transposeMaxSize;

inference-engine/src/gna_plugin/gna_plugin.cpp

@@ -711,6 +711,8 @@ void GNAPlugin::LoadNetwork(CNNNetwork & _network) {
         manager.register_pass<InsertReshapeAroundMatmulWithFq>();
         manager.register_pass<InsertReshapeAroundMatmulWithAdd>();
         manager.register_pass<InsertReshapeAroundMatmul>();
+        manager.register_pass<SwapInputMatMulWithTrailingTranspose>();
+        manager.register_pass<SwapInputMatMulWithAct>();
         manager.register_pass<SwapInputMatMulWithFq>();
         manager.register_pass<SwapInputMatMulWithBias>();
         manager.register_pass<SwapInputMatMul>();

inference-engine/src/gna_plugin/transformations/handle_transposes_around_matmul.cpp

@@ -34,7 +34,15 @@ void ReplaceTransposeWithReshape(std::shared_ptr<ngraph::Node> transpose_node) {
     transpose_node->output(0).replace(reshape_node->output(0));
 }
 
-void InsertTranspose(std::shared_ptr<ngraph::Node> prev_node, const std::string& base_name) {
+void InsertTranspose(std::shared_ptr<ngraph::Node> prev_node, const std::string& base_name, bool before_matmul) {
+    auto create_reshape = [](const ngraph::Shape& shape, std::shared_ptr<ngraph::Node> input_node, const std::string& name) {
+        auto reshape_const = std::make_shared<ngraph::opset8::Constant>(ngraph::element::Type_t::i64,
+            ngraph::Shape{shape.size()}, shape);
+        auto node = std::make_shared<ngraph::opset8::Reshape>(input_node, reshape_const, false);
+        node->set_friendly_name(name);
+        return node;
+    };
+
     auto consumers = prev_node->output(0).get_target_inputs();
     const auto orig_shape = prev_node->get_output_shape(0);
     std::vector<size_t> transpose_ids;
@@ -48,18 +56,29 @@ void InsertTranspose(std::shared_ptr<ngraph::Node> prev_node, const std::string&
     std::iota(std::begin(permute_order), std::end(permute_order), 0);
     std::swap(permute_order[transpose_ids[0]], permute_order[transpose_ids[1]]);
 
-    auto transpose_order = ngraph::opset8::Constant::create(ngraph::element::i64, ngraph::Shape{permute_order.size()}, permute_order);
-    auto transpose = std::make_shared<ngraph::opset8::Transpose>(prev_node, transpose_order);
-    transpose->set_friendly_name(base_name + "/in_transpose");
+    ngraph::NodeVector new_ops;
+    std::shared_ptr<ngraph::Node> node = prev_node;
+    if (!before_matmul) {
+        auto shape = prev_node->get_output_shape(0);
+        std::swap(shape[0], shape[1]);
+        node = create_reshape(shape, node, base_name + "/reshape_before_transpose");
+        new_ops.push_back(node);
+    }
 
-    auto reshapeConstAfter = std::make_shared<ngraph::opset8::Constant>(ngraph::element::Type_t::i64,
-        ngraph::Shape{orig_shape.size()}, orig_shape);
-    auto reshapeAfter = std::make_shared<ngraph::opset8::Reshape>(transpose, reshapeConstAfter, false);
-    reshapeAfter->set_friendly_name(base_name + "/reshape_after_transpose");
-    ngraph::copy_runtime_info(prev_node, ngraph::NodeVector{transpose, reshapeAfter});
+    auto transpose_order = ngraph::opset8::Constant::create(ngraph::element::i64, ngraph::Shape{permute_order.size()}, permute_order);
+    node = std::make_shared<ngraph::opset8::Transpose>(node, transpose_order);
+    node->set_friendly_name(base_name + "/in_transpose");
+    new_ops.push_back(node);
+
+    if (before_matmul) {
+        node = create_reshape(orig_shape, node, base_name + "/reshape_after_transpose");
+        new_ops.push_back(node);
+    }
+
+    ngraph::copy_runtime_info(prev_node, new_ops);
 
     for (auto input : consumers) {
-        input.replace_source_output(reshapeAfter);
+        input.replace_source_output(node);
     }
 }
 
@@ -94,24 +113,25 @@ HandleTransposeBeforeMatMul::HandleTransposeBeforeMatMul() {
             return false;
         }
 
+        auto matmul_node = matmul_iter->second.get_node_shared_ptr();
         auto transpose_reshape_it = pattern_map.find(transpose);
         if (transpose_reshape_it != std::end(pattern_map)) {
             ReplaceTransposeWithReshape(transpose_reshape_it->second.get_node_shared_ptr());
         } else if ((transpose_reshape_it = pattern_map.find(reshape)) != std::end(pattern_map)) {
             auto reshape_node = pattern_map.at(reshape).get_node_shared_ptr();
             if (GNALimitations::IsTransposeSupported(reshape_node->get_output_shape(0))) {
-                auto matmul_node = matmul_iter->second.get_node_shared_ptr();
-                InsertTranspose(reshape_node, matmul_node->get_friendly_name());
+                InsertTranspose(reshape_node, matmul_node->get_friendly_name(), true);
             }
         }
 
+        // Transpose the constant input if it's the first input
         auto iter = pattern_map.find(fq);
         if (iter != pattern_map.end() ||
             (iter = pattern_map.find(constant)) != pattern_map.end()) {
             auto prev_node = iter->second.get_node_shared_ptr();
-            if (!GNALimitations::IsTransposeSupported(prev_node->get_output_shape(0))) return false;
-            auto matmul_node = iter->second.get_node_shared_ptr();
-            InsertTranspose(prev_node, matmul_node->get_friendly_name());
+            if (GNALimitations::IsTranspose2d(prev_node->get_output_shape(0))) {
+                InsertTranspose(prev_node, prev_node->get_friendly_name(), true);
+            }
         }
         return true;
     };
@@ -129,7 +149,11 @@ HandleTransposeAfterMatMul::HandleTransposeAfterMatMul() {
     auto fq_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{matmul, add_left, add_right});
     auto fq = ngraph::pattern::wrap_type<ngraph::opset8::FakeQuantize>({fq_input, ngraph::pattern::any_input(),
         ngraph::pattern::any_input(), ngraph::pattern::any_input(), ngraph::pattern::any_input()});
-    auto transpose_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{fq_input, fq});
+    auto act_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{fq_input, fq});
+    auto act = ngraph::pattern::wrap_type<ngraph::opset8::Relu, ngraph::opset8::Sigmoid,
+            ngraph::opset8::Tanh, ngraph::opset8::Abs, ngraph::opset8::Log, ngraph::opset8::Exp,
+            ngraph::opset8::Sign, ngraph::opset8::Clamp>({act_input});
+    auto transpose_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{act_input, act});
     auto transpose = ngraph::pattern::wrap_type<ngraph::opset8::Transpose>({transpose_input, ngraph::pattern::any_input()});
     auto reshape_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{transpose_input, transpose});
     auto reshape = ngraph::pattern::wrap_type<ngraph::opset8::Reshape>(
@@ -142,16 +166,17 @@ HandleTransposeAfterMatMul::HandleTransposeAfterMatMul() {
             ReplaceTransposeWithReshape(transpose_it->second.get_node_shared_ptr());
         } else {
             auto reshape_node = pattern_map.at(reshape).get_node_shared_ptr();
-            if (!GNALimitations::IsTransposeSupported(reshape_node->get_output_shape(0))) return false;
-            auto iter = pattern_map.find(fq);
+            if (!GNALimitations::IsTransposeSupported(reshape_node->get_input_shape(0))) return false;
+            auto iter = pattern_map.find(act);
             if (iter == pattern_map.end() &&
+                (iter = pattern_map.find(fq)) == pattern_map.end() &&
                 (iter = pattern_map.find(add_left)) == pattern_map.end() &&
                 (iter = pattern_map.find(add_right)) == pattern_map.end() &&
                 (iter = pattern_map.find(matmul)) == pattern_map.end()) {
                 return false;
             }
             auto node = iter->second.get_node_shared_ptr();
-            InsertTranspose(node, node->get_friendly_name());
+            InsertTranspose(node, node->get_friendly_name(), false);
         }
         return true;
     };

inference-engine/src/gna_plugin/transformations/handle_transposes_around_matmul.hpp

@@ -11,15 +11,18 @@ namespace GNAPluginNS {
 /**
  * @brief Inserts Transpose before MatMul or removes it (if it exists) if there is Reshape
  * before MatMul which changes the batch size:
- *    [1, A*B]                 [1, A*B]
+ *    [1, A*B]                [1, A*B]
  *       |                       |
  *    Reshape                 Reshape
  *       |                       |
- * [1, A, 1, B]            [1, A, 1, B]
+ *    [A, B]                  [A, B]
  *       |                       |
  *       |                   Transpose
  *       |           ->          |
- *       |           <-     [1, B, 1, A]
+ *       |           <-       [B, A]
+ *       |                       |
+ *       |                    Reshape
+ *       |                    [A, B]
  *       |                       |
  *    MatMul                   MatMul
  */
@@ -33,12 +36,20 @@ public:
  * @brief Inserts Transpose after MatMul or removes it (if it exists) if there is Reshape
  * after MatMul which changes the batch size:
  *    MatMul                  MatMul
+ *    [A, B]                  [A, B]
  *       |                       |
- * [1, A, 1, B]            [1, A, 1, B]
+ *     [Add]                   [Add]
+ *       |                       |
+ *  [FakeQuantize]        [FakeQuantize]
+ *       |                       |
+ *   [Activation]          [Activation]
+ *       |                       |
+ *       |                    Reshape
+ *       |                    [B, A]
  *       |                       |
  *       |                   Transpose
  *       |           ->          |
- *       |           <-     [1, B, 1, A]
+ *       |           <-        [A, B]
  *       |                       |
  *    Reshape                 Reshape
  *       |                       |

inference-engine/src/gna_plugin/transformations/swap_input_matmul_gna.cpp

@@ -23,13 +23,17 @@ namespace GNAPluginNS {
 NGRAPH_RTTI_DEFINITION(SwapInputMatMul, "SwapInputMatMul", 0);
 NGRAPH_RTTI_DEFINITION(SwapInputMatMulWithBias, "SwapInputMatMulWithBias", 0);
 NGRAPH_RTTI_DEFINITION(SwapInputMatMulWithFq, "SwapInputMatMulWithFq", 0);
+NGRAPH_RTTI_DEFINITION(SwapInputMatMulWithAct, "SwapInputMatMulWithAct", 0);
+NGRAPH_RTTI_DEFINITION(SwapInputMatMulWithTrailingTranspose, "SwapInputMatMulWithTrailingTranspose", 0);
 
 static void SwapAndTransposeInputs(
     std::shared_ptr<ngraph::opset8::MatMul> matmul_node,
-    std::shared_ptr<ngraph::Node> add,
-    std::shared_ptr<ngraph::Node> bias,
-    std::shared_ptr<ngraph::Node> fq,
-    const std::string& last_layer_name) {
+    const std::string& last_layer_name,
+    std::shared_ptr<ngraph::Node> add = nullptr,
+    std::shared_ptr<ngraph::Node> bias = nullptr,
+    std::shared_ptr<ngraph::Node> fq = nullptr,
+    std::shared_ptr<ngraph::Node> act = nullptr,
+    std::shared_ptr<ngraph::Node> transpose = nullptr) {
     auto create_transpose =
         [](ngraph::Output<ngraph::Node> node, const std::string& transpose_name) -> std::shared_ptr<ngraph::Node> {
         ngraph::Shape output_shape = node.get_node_shared_ptr()->get_shape();
@@ -46,12 +50,37 @@ static void SwapAndTransposeInputs(
 
     ngraph::NodeVector new_ops;
 
+    auto transpose_matmul_input = [matmul_node, &new_ops, create_transpose](size_t ix) {
+        std::shared_ptr<ngraph::Node> matmul_input = matmul_node->input_value(ix).get_node_shared_ptr();
+        auto input_transpose = std::dynamic_pointer_cast<ngraph::opset8::Transpose>(matmul_input);
+        if (input_transpose != nullptr) {
+            matmul_input = matmul_input->input_value(0).get_node_shared_ptr();
+            ngraph::replace_output_update_name(input_transpose->output(0), input_transpose->input_value(0));
+        } else {
+            matmul_input = create_transpose(matmul_node->input_value(ix), matmul_node->get_friendly_name() + "/input_transpose");
+            new_ops.push_back(matmul_input);
+        }
+        return matmul_input;
+    };
+
     gnalog() << "Swap and transpose inputs for " << matmul_node->get_friendly_name() << "\n";
-    std::shared_ptr<ngraph::Node> new_matmul = std::make_shared<ngraph::opset8::MatMul>(
-        matmul_node->input(1).get_source_output(), matmul_node->input(0).get_source_output(),
-        !matmul_node->get_transpose_b(), !matmul_node->get_transpose_a());
-    new_matmul->set_friendly_name(matmul_node->get_friendly_name() + "/swap_inputs");
-    new_ops.push_back(new_matmul);
+
+    bool first_input_const = false;
+    auto first_input = matmul_node->input_value(0).get_node_shared_ptr();
+    if (std::dynamic_pointer_cast<ngraph::opset8::FakeQuantize>(first_input)) {
+        first_input = first_input->input_value(0).get_node_shared_ptr();
+    }
+    if (std::dynamic_pointer_cast<ngraph::opset8::Constant>(first_input)) {
+        first_input_const = true;
+    }
+
+    auto input1 = first_input_const ? transpose_matmul_input(1) : matmul_node->input_value(1);
+    auto input2 = first_input_const ? matmul_node->input_value(0) : transpose_matmul_input(0);
+    bool transpose_1 = first_input_const ? matmul_node->get_transpose_b() : !matmul_node->get_transpose_b();
+    bool transpose_2 = first_input_const ? !matmul_node->get_transpose_a() : matmul_node->get_transpose_a();
+    std::shared_ptr<ngraph::Node> new_node = std::make_shared<ngraph::opset8::MatMul>(input1, input2, transpose_1, transpose_2);
+    new_node->set_friendly_name(matmul_node->get_friendly_name() + "/swap_inputs");
+    new_ops.push_back(new_node);
 
     std::shared_ptr<ngraph::Node> old_root_node = matmul_node;
     if (bias != nullptr) {
@@ -74,23 +103,34 @@ static void SwapAndTransposeInputs(
             }
         }
 
-        new_matmul = std::make_shared<ngraph::opset8::Add>(new_matmul, bias);
+        new_node = std::make_shared<ngraph::opset8::Add>(new_node, bias);
         old_root_node = add;
-        new_ops.push_back(new_matmul);
+        new_ops.push_back(new_node);
     }
 
     if (fq != nullptr) {
-        new_matmul = fq->clone_with_new_inputs({new_matmul, fq->input_value(1), fq->input_value(2),
+        new_node = fq->clone_with_new_inputs({new_node, fq->input_value(1), fq->input_value(2),
             fq->input_value(3), fq->input_value(4)});
         old_root_node = fq;
-        new_ops.push_back(new_matmul);
+        new_ops.push_back(new_node);
     }
 
-    auto output = create_transpose(new_matmul, last_layer_name);
-    new_ops.push_back(output);
+    if (act != nullptr) {
+        new_node = act->clone_with_new_inputs({new_node});
+        old_root_node = act;
+        new_ops.push_back(new_node);
+    }
+
+    if (transpose == nullptr) {
+        new_node = create_transpose(new_node, last_layer_name);
+        new_ops.push_back(new_node);
+    } else {
+        ngraph::replace_output_update_name(transpose->output(0), transpose->input_value(0));
+        new_node->set_friendly_name(last_layer_name);
+    }
 
     ngraph::copy_runtime_info(matmul_node, new_ops);
-    ngraph::replace_node(old_root_node, output);
+    ngraph::replace_node(old_root_node, new_node);
 }
 
 static std::shared_ptr<ngraph::Node> CreateMatmul(
@@ -103,13 +143,16 @@ static std::shared_ptr<ngraph::Node> CreateMatmul(
         ngraph::pattern::wrap_type<ngraph::opset8::Constant>(),
         ngraph::pattern::wrap_type<ngraph::opset8::Constant>(),
         ngraph::pattern::wrap_type<ngraph::opset8::Constant>()});
-    auto matmul_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{constant, fake_quantize});
+    auto matmul_const_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{constant, fake_quantize});
+    auto transpose = ngraph::pattern::wrap_type<ngraph::opset8::Transpose>();
+    auto matmul_non_const_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{transpose,
+        ngraph::pattern::any_input()});
     if (is_first_constant) {
         return ngraph::pattern::wrap_type<ngraph::opset8::MatMul>(
-            {matmul_input, ngraph::pattern::any_input()}, matmul_predicate);
+            {matmul_const_input, matmul_non_const_input}, matmul_predicate);
     }
     return ngraph::pattern::wrap_type<ngraph::opset8::MatMul>(
-        {ngraph::pattern::any_input(), matmul_input}, matmul_predicate);
+        {ngraph::pattern::any_input(), matmul_non_const_input}, matmul_predicate);
 }
 
 static std::shared_ptr<ngraph::Node> CreateMatmuls(
@@ -151,14 +194,13 @@ SwapInputMatMul::SwapInputMatMul() {
     auto callback = [=](ngraph::pattern::Matcher& m) {
         const auto& pattern_map = m.get_pattern_value_map();
         auto iter = pattern_map.find(matmul1);
-        if (iter == pattern_map.end() &&
-            (iter = pattern_map.find(matmul2)) == pattern_map.end()) {
+        if (iter == pattern_map.end() && (iter = pattern_map.find(matmul2)) == pattern_map.end()) {
             return false;
         }
 
         auto matmul_node = std::dynamic_pointer_cast<ngraph::opset8::MatMul>(iter->second.get_node_shared_ptr());
         IE_ASSERT(matmul_node != nullptr);
-        SwapAndTransposeInputs(matmul_node, nullptr, nullptr, nullptr, "");
+        SwapAndTransposeInputs(matmul_node, matmul_node->get_friendly_name());
         return true;
     };
 
@@ -176,8 +218,7 @@ SwapInputMatMulWithBias::SwapInputMatMulWithBias() {
     auto callback = [=](ngraph::pattern::Matcher& m) {
         const auto& pattern_map = m.get_pattern_value_map();
         auto iter = pattern_map.find(matmul1);
-        if (iter == pattern_map.end() &&
-            (iter = pattern_map.find(matmul2)) == pattern_map.end()) {
+        if (iter == pattern_map.end() && (iter = pattern_map.find(matmul2)) == pattern_map.end()) {
             return false;
         }
 
@@ -185,10 +226,9 @@ SwapInputMatMulWithBias::SwapInputMatMulWithBias() {
         IE_ASSERT(matmul_node != nullptr);
         SwapAndTransposeInputs(
             matmul_node,
+            pattern_map.at(add).get_node_shared_ptr()->get_friendly_name(),
             pattern_map.at(add).get_node_shared_ptr(),
-            pattern_map.at(bias).get_node_shared_ptr(),
-            nullptr,
-            pattern_map.at(add).get_node_shared_ptr()->get_friendly_name());
+            pattern_map.at(bias).get_node_shared_ptr());
         return true;
     };
 
@@ -212,8 +252,7 @@ SwapInputMatMulWithFq::SwapInputMatMulWithFq() {
     auto callback = [=](ngraph::pattern::Matcher& m) {
         const auto& pattern_map = m.get_pattern_value_map();
         auto iter = pattern_map.find(matmul1);
-        if (iter == pattern_map.end() &&
-            (iter = pattern_map.find(matmul2)) == pattern_map.end()) {
+        if (iter == pattern_map.end() && (iter = pattern_map.find(matmul2)) == pattern_map.end()) {
             return false;
         }
 
@@ -223,14 +262,104 @@ SwapInputMatMulWithFq::SwapInputMatMulWithFq() {
         IE_ASSERT(matmul_node != nullptr);
         SwapAndTransposeInputs(
             matmul_node,
+            pattern_map.at(fq).get_node_shared_ptr()->get_friendly_name(),
             iter_add != pattern_map.end() ? iter_add->second.get_node_shared_ptr() : nullptr,
             iter_bias != pattern_map.end() ? iter_bias->second.get_node_shared_ptr() : nullptr,
-            pattern_map.at(fq).get_node_shared_ptr(),
-            pattern_map.at(fq).get_node_shared_ptr()->get_friendly_name());
+            pattern_map.at(fq).get_node_shared_ptr());
         return true;
     };
 
     auto matcher = std::make_shared<ngraph::pattern::Matcher>(fq, "SwapInputMatMulWithFq");
     this->register_matcher(matcher, callback);
 }
+
+SwapInputMatMulWithAct::SwapInputMatMulWithAct() {
+    MATCHER_SCOPE(SwapInputMatMulWithAct);
+    std::shared_ptr<ngraph::Node> matmul1;
+    std::shared_ptr<ngraph::Node> matmul2;
+    auto matmul = CreateMatmuls(matmul1, matmul2);
+    auto bias = ngraph::pattern::wrap_type<ngraph::opset8::Constant>();
+    auto add = ngraph::pattern::wrap_type<ngraph::opset8::Add>({matmul, bias});
+    auto fq_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{add, matmul});
+    auto fq = ngraph::pattern::wrap_type<ngraph::opset8::FakeQuantize>({fq_input,
+        ngraph::pattern::wrap_type<ngraph::opset8::Constant>(),
+        ngraph::pattern::wrap_type<ngraph::opset8::Constant>(),
+        ngraph::pattern::wrap_type<ngraph::opset8::Constant>(),
+        ngraph::pattern::wrap_type<ngraph::opset8::Constant>()});
+    auto act_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{fq_input, fq});
+    auto act = ngraph::pattern::wrap_type<ngraph::opset8::Relu, ngraph::opset8::Sigmoid,
+        ngraph::opset8::Tanh, ngraph::opset8::Abs, ngraph::opset8::Log, ngraph::opset8::Exp,
+        ngraph::opset8::Sign, ngraph::opset8::Clamp>({act_input});
+    auto callback = [=](ngraph::pattern::Matcher& m) {
+        const auto& pattern_map = m.get_pattern_value_map();
+        auto iter = pattern_map.find(matmul1);
+        if (iter == pattern_map.end() && (iter = pattern_map.find(matmul2)) == pattern_map.end()) {
+            return false;
+        }
+
+        auto iter_add = pattern_map.find(add);
+        auto iter_bias = pattern_map.find(bias);
+        auto iter_fq = pattern_map.find(fq);
+        auto matmul_node = std::dynamic_pointer_cast<ngraph::opset8::MatMul>(iter->second.get_node_shared_ptr());
+        IE_ASSERT(matmul_node != nullptr);
+        SwapAndTransposeInputs(
+            matmul_node,
+            pattern_map.at(act).get_node_shared_ptr()->get_friendly_name(),
+            iter_add != pattern_map.end() ? iter_add->second.get_node_shared_ptr() : nullptr,
+            iter_bias != pattern_map.end() ? iter_bias->second.get_node_shared_ptr() : nullptr,
+            iter_fq != pattern_map.end() ? iter_fq->second.get_node_shared_ptr() : nullptr,
+            pattern_map.at(act).get_node_shared_ptr());
+        return true;
+    };
+
+    auto matcher = std::make_shared<ngraph::pattern::Matcher>(act, "SwapInputMatMulWithAct");
+    this->register_matcher(matcher, callback);
+}
+
+SwapInputMatMulWithTrailingTranspose::SwapInputMatMulWithTrailingTranspose() {
+    MATCHER_SCOPE(SwapInputMatMulWithTrailingTranspose);
+    std::shared_ptr<ngraph::Node> matmul1;
+    std::shared_ptr<ngraph::Node> matmul2;
+    auto matmul = CreateMatmuls(matmul1, matmul2);
+    auto bias = ngraph::pattern::wrap_type<ngraph::opset8::Constant>();
+    auto add = ngraph::pattern::wrap_type<ngraph::opset8::Add>({matmul, bias});
+    auto fq_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{add, matmul});
+    auto fq = ngraph::pattern::wrap_type<ngraph::opset8::FakeQuantize>({fq_input,
+        ngraph::pattern::wrap_type<ngraph::opset8::Constant>(),
+        ngraph::pattern::wrap_type<ngraph::opset8::Constant>(),
+        ngraph::pattern::wrap_type<ngraph::opset8::Constant>(),
+        ngraph::pattern::wrap_type<ngraph::opset8::Constant>()});
+    auto act_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{fq_input, fq});
+    auto act = ngraph::pattern::wrap_type<ngraph::opset8::Relu, ngraph::opset8::Sigmoid,
+        ngraph::opset8::Tanh, ngraph::opset8::Abs, ngraph::opset8::Log, ngraph::opset8::Exp,
+        ngraph::opset8::Sign, ngraph::opset8::Clamp>({act_input});
+    auto transpose_input = std::make_shared<ngraph::pattern::op::Or>(ngraph::OutputVector{act_input, act});
+    auto transpose = ngraph::pattern::wrap_type<ngraph::opset8::Transpose>({transpose_input, ngraph::pattern::any_input()});
+    auto callback = [=](ngraph::pattern::Matcher& m) {
+        const auto& pattern_map = m.get_pattern_value_map();
+        auto iter = pattern_map.find(matmul1);
+        if (iter == pattern_map.end() && (iter = pattern_map.find(matmul2)) == pattern_map.end()) {
+            return false;
+        }
+
+        auto iter_add = pattern_map.find(add);
+        auto iter_bias = pattern_map.find(bias);
+        auto iter_fq = pattern_map.find(fq);
+        auto iter_act = pattern_map.find(act);
+        auto matmul_node = std::dynamic_pointer_cast<ngraph::opset8::MatMul>(iter->second.get_node_shared_ptr());
+        IE_ASSERT(matmul_node != nullptr);
+        SwapAndTransposeInputs(
+            matmul_node,
+            pattern_map.at(transpose).get_node_shared_ptr()->get_friendly_name(),
+            iter_add != pattern_map.end() ? iter_add->second.get_node_shared_ptr() : nullptr,
+            iter_bias != pattern_map.end() ? iter_bias->second.get_node_shared_ptr() : nullptr,
+            iter_fq != pattern_map.end() ? iter_fq->second.get_node_shared_ptr() : nullptr,
+            iter_act != pattern_map.end() ? iter_act->second.get_node_shared_ptr() : nullptr,
+            pattern_map.at(transpose).get_node_shared_ptr());
+        return true;
+    };
+
+    auto matcher = std::make_shared<ngraph::pattern::Matcher>(transpose, "SwapInputMatMulWithTrailingTranspose");
+    this->register_matcher(matcher, callback);
+}
 } // namespace GNAPluginNS

inference-engine/src/gna_plugin/transformations/swap_input_matmul_gna.hpp

@@ -8,9 +8,28 @@
 #include <ngraph/pass/graph_rewrite.hpp>
 
 namespace GNAPluginNS {
-// @brief Swaps and transposes inputs of MatMul if
-// 1. its first input is const and its batch size isn't supported by GNA
-// 2. its first input is non-const and its batch size isn't supported by GNA
+/* @brief Swaps and transposes inputs of MatMul and transposes its output if
+ * 1. its first input is const and its batch size isn't supported by GNA
+ * 2. its first input is non-const and its batch size isn't supported by GNA
+ * The following pattern is detected:
+ *    Constant           Any input
+ *       |                   |
+ *   [FakeQuantize]      [Transpose]
+ *        \                 /
+ *               MatMul
+ *                 |
+ *               [Add]
+ *                 |
+ *           [FakeQuantize]
+ *                 |
+ *            [Activation]
+ *                 |
+ *            [Transpose]
+ * 
+ * The existed Transposes will be removed instead of inserting new ones during transposition.
+ * Separate matchers are required for different last nodes. They should be registered in the order from the longest
+ * to the shortest pattern.
+ **/
 class SwapInputMatMul: public ngraph::pass::MatcherPass {
 public:
     NGRAPH_RTTI_DECLARATION;
@@ -28,6 +47,18 @@ public:
     NGRAPH_RTTI_DECLARATION;
     SwapInputMatMulWithFq();
 };
+
+class SwapInputMatMulWithAct: public ngraph::pass::MatcherPass {
+public:
+    NGRAPH_RTTI_DECLARATION;
+    SwapInputMatMulWithAct();
+};
+
+class SwapInputMatMulWithTrailingTranspose: public ngraph::pass::MatcherPass {
+public:
+    NGRAPH_RTTI_DECLARATION;
+    SwapInputMatMulWithTrailingTranspose();
+};
 } // namespace GNAPluginNS
 
 #endif // SWAP_INPUT_MATMUL_GNA_HPP

inference-engine/tests/functional/plugin/gna/pass_tests/convert_matmul_to_fullyconnected.cpp

@@ -75,10 +75,77 @@ protected:
     }
 };
 
+class ConvertMatmulToFcWithTransposesPass : public testing::WithParamInterface<ConvertMatmulToFcPassParams >,
+                              public LayerTestsUtils::LayerTestsCommon {
+public:
+    static std::string getTestCaseName(testing::TestParamInfo<ConvertMatmulToFcPassParams> obj) {
+        InferenceEngine::Precision netPrecision;
+        std::string targetDevice;
+        std::map<std::string, std::string> configuration;
+        std::vector<std::vector<size_t>> inputShape;
+        std::tie(inputShape, netPrecision, targetDevice, configuration) = obj.param;
+
+        std::ostringstream result;
+        result << "netPRC=" << netPrecision.name() << "_";
+        result << "targetDevice=" << targetDevice << "_";
+        for (auto const& configItem : configuration) {
+            result << "_configItem=" << configItem.first << "_" << configItem.second;
+        }
+        result << "_IS=" << CommonTestUtils::vec2str(inputShape[1]) << "_";
+        result << "_CS=" << CommonTestUtils::vec2str(inputShape[0]) << "_";
+        return result.str();
+    }
+
+InferenceEngine::Blob::Ptr GenerateInput(const InferenceEngine::InputInfo& info) const override {
+        InferenceEngine::Blob::Ptr blob = make_blob_with_precision(info.getTensorDesc());
+        blob->allocate();
+
+        auto* rawBlobDataPtr = blob->buffer().as<float*>();
+        std::vector<float> values = CommonTestUtils::generate_float_numbers(blob->size(), -0.1f, 0.1f);
+        for (size_t i = 0; i < blob->size(); i++) {
+            rawBlobDataPtr[i] = values[i];
+        }
+        return blob;
+    }
+
+protected:
+    void SetUp() override {
+        InferenceEngine::Precision netPrecision;
+        std::vector<std::vector<size_t>> inputShape;
+        std::tie(inputShape, netPrecision, targetDevice, configuration) = this->GetParam();
+        auto ngPrc = FuncTestUtils::PrecisionUtils::convertIE2nGraphPrc(netPrecision);
+
+        auto params = ngraph::builder::makeParams(ngPrc, { {1, inputShape[1][0] * inputShape[1][1]} });
+
+        auto reshape1 = std::make_shared<ngraph::opset1::Reshape>(params[0],
+            ngraph::builder::makeConstant(ngraph::element::i64, { inputShape[1].size() }, inputShape[1]), false);
+        auto transpose1 = std::make_shared<ngraph::opset1::Transpose>(reshape1,
+            ngraph::opset1::Constant::create(ngraph::element::i64, ngraph::Shape{ 2 }, std::vector<int64_t>{1, 0}));
+
+        std::vector<float> weights = CommonTestUtils::generate_float_numbers(inputShape[0][0] * inputShape[0][1], -0.1f, 0.1f);
+        auto const_mult2 = ngraph::builder::makeConstant<float>(ngPrc, inputShape[0], weights);
+        auto matmul = std::make_shared<ngraph::opset1::MatMul>(const_mult2, transpose1, false, false);
+        auto relu = std::make_shared<ngraph::opset1::Relu>(matmul);
+
+        auto transpose2 = std::make_shared<ngraph::opset1::Transpose>(relu,
+            ngraph::opset1::Constant::create(ngraph::element::i64, ngraph::Shape{ 2 }, std::vector<int64_t>{1, 0}));
+        auto transpose_output_shape = transpose2->get_output_shape(0);
+        ngraph::Shape output_shape = {1, transpose_output_shape[0] * transpose_output_shape[1]};
+        auto reshape2 = std::make_shared<ngraph::opset1::Reshape>(transpose2,
+            ngraph::builder::makeConstant(ngraph::element::i64, { output_shape.size() }, output_shape), false);
+
+        function = std::make_shared<ngraph::Function>(reshape2, params, "ConvertMatmulToFCWithTransposes");
+    }
+};
+
 TEST_P(ConvertMatmulToFcPass, CompareWithRefImpl) {
     Run();
 }
 
+TEST_P(ConvertMatmulToFcWithTransposesPass, CompareWithRefImpl) {
+    Run();
+}
+
 const std::vector<InferenceEngine::Precision> netPrecisions = {
         InferenceEngine::Precision::FP32,
         InferenceEngine::Precision::FP16,
@@ -103,6 +170,13 @@ const std::vector<std::vector<std::vector<size_t>>> input_shapes = {
         {{6, 16}, {16, 8}}
 };
 
+const std::vector<std::vector<std::vector<size_t>>> input_shapes_transposes = {
+        {{1, 8}, {1, 8}},
+        {{128, 8}, {1, 8}},
+        {{8, 8}, {8, 8}},
+        {{1, 16}, {8, 16}},
+        {{6, 16}, {1, 16}}
+};
 
 INSTANTIATE_TEST_SUITE_P(smoke_convert_matmul_to_fc, ConvertMatmulToFcPass,
                         ::testing::Combine(
@@ -111,4 +185,12 @@ INSTANTIATE_TEST_SUITE_P(smoke_convert_matmul_to_fc, ConvertMatmulToFcPass,
                                 ::testing::Values(CommonTestUtils::DEVICE_GNA),
                                 ::testing::ValuesIn(configs)),
                         ConvertMatmulToFcPass::getTestCaseName);
+
+INSTANTIATE_TEST_SUITE_P(smoke_convert_matmul_to_fc, ConvertMatmulToFcWithTransposesPass,
+                        ::testing::Combine(
+                                ::testing::ValuesIn(input_shapes_transposes),
+                                ::testing::ValuesIn(netPrecisions),
+                                ::testing::Values(CommonTestUtils::DEVICE_GNA),
+                                ::testing::ValuesIn(configs)),
+                        ConvertMatmulToFcWithTransposesPass::getTestCaseName);
 } // namespace LayerTestsDefinitions

inference-engine/tests/unit/gna/ngraph/transformations/gna_swap_input_matmul.cpp

@@ -20,9 +20,17 @@ static std::shared_ptr<ngraph::Function> CreateMatMulFunction(const ngraph::Shap
                                                               bool withBias,
                                                               bool withWeightsFq,
                                                               bool withOutFq,
+                                                              bool withAct,
                                                               bool swappedInputs,
-                                                              bool needTranspose) {
+                                                              bool needTranspose,
+                                                              bool expected = false) {
     auto input_params = std::make_shared<ngraph::opset8::Parameter>(ngraph::element::i64, input2_shape);
+    std::shared_ptr<ngraph::Node> input = input_params;
+    if (input->get_output_shape(0).size() == 2 && needTranspose) {
+        auto transpose_order = ngraph::opset8::Constant::create(ngraph::element::i64, ngraph::Shape{2},
+                                                                std::vector<size_t>{1, 0});
+        input = std::make_shared<ngraph::opset8::Transpose>(input, transpose_order);
+    }
 
     auto constant = ngraph::opset8::Constant::create(ngraph::element::i64, input1_shape, {1});
     std::shared_ptr<ngraph::Node> const_input = constant;
@@ -34,14 +42,18 @@ static std::shared_ptr<ngraph::Function> CreateMatMulFunction(const ngraph::Shap
         const_input = std::make_shared<ngraph::opset8::FakeQuantize>(const_input, input_low, input_high,
                                                                      output_low, output_high, 11);
     }
-    auto matmul = swappedInputs ? std::make_shared<ngraph::opset8::MatMul>(input_params, const_input, needTranspose, needTranspose) :
-        std::make_shared<ngraph::opset8::MatMul>(const_input, input_params, needTranspose, needTranspose);
+    auto matmul = swappedInputs ? std::make_shared<ngraph::opset8::MatMul>(input, const_input, false, needTranspose) :
+        std::make_shared<ngraph::opset8::MatMul>(const_input, input, needTranspose, false);
 
     std::shared_ptr<ngraph::Node> final_node = matmul;
     if (withBias) {
-        auto bias = ngraph::opset8::Constant::create(ngraph::element::i64, bias_shape, {1});
+        auto shape = bias_shape;
+        if ((needTranspose && !expected || !needTranspose && expected) && bias_shape.size() > 1) {
+            std::swap(shape[0], shape[1]);
+        }
+        auto bias = ngraph::opset8::Constant::create(ngraph::element::i64, shape, {1});
         std::shared_ptr<ngraph::Node> bias_node = bias;
-        if (needTranspose && bias_shape.size() > 1) {
+        if (expected && bias_shape.size() > 1) {
             auto transpose_order = ngraph::opset8::Constant::create(ngraph::element::i64, ngraph::Shape{2},
                                                                     std::vector<size_t>{1, 0});
             bias_node = std::make_shared<ngraph::opset8::Transpose>(bias_node, transpose_order);
@@ -58,7 +70,11 @@ static std::shared_ptr<ngraph::Function> CreateMatMulFunction(const ngraph::Shap
                                                                     output_low, output_high, 11);
     }
 
-    if (needTranspose) {
+    if (withAct) {
+        final_node = std::make_shared<ngraph::opset8::Relu>(final_node);
+    }
+
+    if (final_node->get_output_shape(0).size() == 2 && needTranspose) {
         auto transpose_order = ngraph::opset8::Constant::create(ngraph::element::i64, ngraph::Shape{2},
                                                                 std::vector<size_t>{1, 0});
         final_node = std::make_shared<ngraph::opset8::Transpose>(final_node, transpose_order);
@@ -72,6 +88,8 @@ static std::shared_ptr<ngraph::Function> CreateMatMulFunction(const ngraph::Shap
 static void Execute(std::shared_ptr<ngraph::Function> function, std::shared_ptr<ngraph::Function> reference_function) {
     ngraph::pass::Manager m;
     m.register_pass<ngraph::pass::InitNodeInfo>();
+    m.register_pass<GNAPluginNS::SwapInputMatMulWithTrailingTranspose>();
+    m.register_pass<GNAPluginNS::SwapInputMatMulWithAct>();
     m.register_pass<GNAPluginNS::SwapInputMatMulWithFq>();
     m.register_pass<GNAPluginNS::SwapInputMatMulWithBias>();
     m.register_pass<GNAPluginNS::SwapInputMatMul>();
@@ -87,13 +105,15 @@ typedef std::tuple<
         std::vector<ngraph::Shape>,         // constant input shape, non-const input shape, bias shape
         bool,                               // with bias
         bool,                               // with weights FakeQuantize
-        bool                                // with output FakeQuantize
+        bool,                               // with output FakeQuantize
+        bool,                               // with activation after MatMul
+        bool                                // with transposes
 > SwapInputMatmulParams;
 
 static std::string getTestCaseName(testing::TestParamInfo<SwapInputMatmulParams> obj) {
     std::vector<ngraph::Shape> shapes;
-    bool withBias, withWeightsFq, withOutFq;
-    std::tie(shapes, withBias, withWeightsFq, withOutFq) = obj.param;
+    bool withBias, withWeightsFq, withOutFq, withAct, withTransposes;
+    std::tie(shapes, withBias, withWeightsFq, withOutFq, withAct, withTransposes) = obj.param;
 
     std::ostringstream result;
     result << "IS1=" << shapes[0] << "_";
@@ -101,7 +121,9 @@ static std::string getTestCaseName(testing::TestParamInfo<SwapInputMatmulParams>
     result << "BS=" << shapes[2] << "_";
     result << "bias=" << withBias << "_";
     result << "wFQ=" << withWeightsFq << "_";
-    result << "oFQ=" << withOutFq;
+    result << "oFQ=" << withOutFq << "_";
+    result << "act=" << withAct << "_";
+    result << "transposes=" << withTransposes;
     return result.str();
 }
 
@@ -110,14 +132,26 @@ enum class MatmulInputType {
     SecondInputConstant
 }; // enum class MatmulInputType
 
+static void transposeInputShapes(std::vector<ngraph::Shape> &shapes) {
+    if (shapes[0].size() > 1) {
+        std::swap(shapes[0][0], shapes[0][1]);
+    }
+    if (shapes[1].size() > 1) {
+        std::swap(shapes[1][0], shapes[1][1]);
+    }
+    if (shapes[2].size() > 1) {
+        std::swap(shapes[2][0], shapes[2][1]);
+    }
+}
+
 template<MatmulInputType E>
 class SwapInputMatmul : public CommonTestUtils::TestsCommon,
                         public ::testing::WithParamInterface<SwapInputMatmulParams> {
 public:
     void SetUp() override {
         std::vector<ngraph::Shape> shapes;
-        bool withBias, withWeightsFq, withOutFq;
-        std::tie(shapes, withBias, withWeightsFq, withOutFq) = this->GetParam();
+        bool withBias, withWeightsFq, withOutFq, withAct, withTransposes;
+        std::tie(shapes, withBias, withWeightsFq, withOutFq, withAct, withTransposes) = this->GetParam();
 
         bool swap_inputs = false;
         switch (E) {
@@ -128,9 +162,13 @@ public:
             break;
         }
 
-        function = CreateMatMulFunction(shapes[0], shapes[1], shapes[2], withBias, withWeightsFq, withOutFq, swap_inputs, false);
+        if (withTransposes) {
+            transposeInputShapes(shapes);
+        }
+        function = CreateMatMulFunction(shapes[0], shapes[1], shapes[2], withBias, withWeightsFq, withOutFq,
+                                        withAct, swap_inputs, withTransposes);
         reference_function = CreateMatMulFunction(shapes[0], shapes[1], shapes[2], withBias, withWeightsFq,
-                                                  withOutFq, !swap_inputs, true);
+                                                  withOutFq, withAct, !swap_inputs, !withTransposes, true);
     }
 public:
     std::shared_ptr<ngraph::Function> function, reference_function;
@@ -142,8 +180,8 @@ class SwapInputMatmulNotApplied : public CommonTestUtils::TestsCommon,
 public:
     void SetUp() override {
         std::vector<ngraph::Shape> shapes;
-        bool withBias, withWeightsFq, withOutFq;
-        std::tie(shapes, withBias, withWeightsFq, withOutFq) = this->GetParam();
+        bool withBias, withWeightsFq, withOutFq, withAct, withTransposes;
+        std::tie(shapes, withBias, withWeightsFq, withOutFq, withAct, withTransposes) = this->GetParam();
 
         bool swap_inputs = false;
         switch (E) {
@@ -154,7 +192,11 @@ public:
             break;
         }
 
-        function = CreateMatMulFunction(shapes[0], shapes[1], shapes[2], withBias, withWeightsFq, withOutFq, swap_inputs, false);
+        if (withTransposes) {
+            transposeInputShapes(shapes);
+        }
+        function = CreateMatMulFunction(shapes[0], shapes[1], shapes[2], withBias, withWeightsFq, withOutFq,
+            withAct, swap_inputs, withTransposes);
         reference_function = ngraph::clone_function(*function);
     }
 public:
@@ -209,6 +251,8 @@ INSTANTIATE_TEST_SUITE_P(smoke_swap_input_matmul, SwapInputMatmulWithFirstInputC
         ::testing::ValuesIn(input_shapes_for_matmul_with_first_constant_applied),
         ::testing::ValuesIn(std::vector<bool>{false, true}),
         ::testing::ValuesIn(std::vector<bool>{false, true}),
+        ::testing::ValuesIn(std::vector<bool>{false, true}),
+        ::testing::ValuesIn(std::vector<bool>{false, true}),
         ::testing::ValuesIn(std::vector<bool>{false, true})),
     getTestCaseName);
 
@@ -217,6 +261,8 @@ INSTANTIATE_TEST_SUITE_P(smoke_swap_input_matmul, SwapInputMatmulWithFirstInputC
         ::testing::ValuesIn(input_shapes_for_matmul_with_first_constant_not_applied),
         ::testing::ValuesIn(std::vector<bool>{false, true}),
         ::testing::ValuesIn(std::vector<bool>{false, true}),
+        ::testing::ValuesIn(std::vector<bool>{false, true}),
+        ::testing::ValuesIn(std::vector<bool>{false, true}),
         ::testing::ValuesIn(std::vector<bool>{false, true})),
     getTestCaseName);
 
@@ -225,6 +271,8 @@ INSTANTIATE_TEST_SUITE_P(smoke_swap_input_matmul, SwapInputMatmulWithSecondInput
         ::testing::ValuesIn(input_shapes_for_matmul_with_second_constant_applied),
         ::testing::ValuesIn(std::vector<bool>{false, true}),
         ::testing::ValuesIn(std::vector<bool>{false, true}),
+        ::testing::ValuesIn(std::vector<bool>{false, true}),
+        ::testing::ValuesIn(std::vector<bool>{false, true}),
         ::testing::ValuesIn(std::vector<bool>{false, true})),
     getTestCaseName);
 
@@ -233,6 +281,8 @@ INSTANTIATE_TEST_SUITE_P(smoke_swap_input_matmul, SwapInputMatmulWithSecondInput
         ::testing::ValuesIn(input_shapes_for_matmul_with_second_constant_not_applied),
         ::testing::ValuesIn(std::vector<bool>{false, true}),
         ::testing::ValuesIn(std::vector<bool>{false, true}),
+        ::testing::ValuesIn(std::vector<bool>{false, true}),
+        ::testing::ValuesIn(std::vector<bool>{false, true}),
         ::testing::ValuesIn(std::vector<bool>{false, true})),
     getTestCaseName);
 

inference-engine/tests/unit/gna/ngraph/transformations/handle_transposes_around_matmul.cpp

@@ -74,7 +74,7 @@ std::shared_ptr<ngraph::Function> CreateMatmulTransposeFunction(
     const ngraph::Shape& input_shape,
     const ngraph::Shape& matmul_shape,
     const ngraph::Shape& reshape_shape,
-    bool create_reshape_after_transpose,
+    bool create_reshape_before_transpose,
     bool enable_last_reshape,
     bool enable_add,
     bool matmul_on_left_side,
@@ -105,29 +105,23 @@ std::shared_ptr<ngraph::Function> CreateMatmulTransposeFunction(
             255);
     }
 
-    auto transpose_order = ngraph::opset7::Constant::create(ngraph::element::i64, ngraph::Shape{2}, {1, 0});
-    auto transpose = std::make_shared<ngraph::opset7::Transpose>(node, transpose_order);
-    const auto transpose_output_shape = transpose->get_output_shape(0);
-
-    std::shared_ptr<ngraph::Node> reshape;
-    auto shape_const = ngraph::opset7::Constant::create(ngraph::element::i64, ngraph::Shape{reshape_shape.size()}, reshape_shape);
-    if (create_reshape_after_transpose) {
-        const auto matmul_output_shape = node->get_output_shape(0);
-        auto reshape_after_transpose_const = ngraph::opset7::Constant::create(ngraph::element::i64,
+    if (create_reshape_before_transpose) {
+        auto matmul_output_shape = node->get_output_shape(0);
+        std::swap(matmul_output_shape[0], matmul_output_shape[1]);
+        auto reshape_before_transpose_const = ngraph::opset7::Constant::create(ngraph::element::i64,
             ngraph::Shape{matmul_output_shape.size()}, matmul_output_shape);
-        auto reshape_after_transpose = std::make_shared<ngraph::opset7::Reshape>(transpose, reshape_after_transpose_const, false);
-        reshape = reshape_after_transpose;
-        if (enable_last_reshape) {
-            reshape = std::make_shared<ngraph::opset7::Reshape>(reshape_after_transpose, shape_const, false);
-        }
-    } else {
-        reshape = transpose;
-        if (enable_last_reshape) {
-            reshape = std::make_shared<ngraph::opset7::Reshape>(transpose, shape_const, false);
-        }
+        node = std::make_shared<ngraph::opset7::Reshape>(node, reshape_before_transpose_const, false);
     }
 
-    auto result = std::make_shared<ngraph::opset7::Result>(reshape);
+    auto transpose_order = ngraph::opset7::Constant::create(ngraph::element::i64, ngraph::Shape{2}, {1, 0});
+    node = std::make_shared<ngraph::opset7::Transpose>(node, transpose_order);
+
+    if (enable_last_reshape) {
+        auto shape_const = ngraph::opset7::Constant::create(ngraph::element::i64, ngraph::Shape{reshape_shape.size()}, reshape_shape);
+        node = std::make_shared<ngraph::opset7::Reshape>(node, shape_const, false);
+    }
+
+    auto result = std::make_shared<ngraph::opset7::Result>(node);
     return std::make_shared<ngraph::Function>(ngraph::ResultVector{result}, ngraph::ParameterVector{input_params});
 }