From 2893f2fc7c67a26d3e550ed09a0bb0cbe10ce329 Mon Sep 17 00:00:00 2001 From: Alexandra Sidorova Date: Fri, 7 Jun 2024 15:13:26 +0300 Subject: [PATCH] [Snippets] Added dynamism support to MHA Tokenization pass (#24897) ### Details: - *Updated `MHATokenization` pass to support dynamic nodes tokenization* - *Added MHATokenization unit tests* - *Disabled dynamic MHA tokenization in CPU Plugin* ### Tickets: - *123329* --- .../snippets/src/pass/mha_tokenization.cpp | 85 +++++--- .../tests/src/pass/mha_tokenization.cpp | 70 ++++++- .../transformation_pipeline.cpp | 2 +- .../include/subgraph_mha.hpp | 18 +- .../ov_snippets_models/src/subgraph_mha.cpp | 196 +++++++++--------- 5 files changed, 227 insertions(+), 144 deletions(-) diff --git a/src/common/snippets/src/pass/mha_tokenization.cpp b/src/common/snippets/src/pass/mha_tokenization.cpp index 08deb95b12e..d928cdd1d33 100644 --- a/src/common/snippets/src/pass/mha_tokenization.cpp +++ b/src/common/snippets/src/pass/mha_tokenization.cpp @@ -16,7 +16,7 @@ namespace { bool is_supported_tensor(const ov::descriptor::Tensor& t) { - return t.get_partial_shape().is_static() && ov::snippets::utils::one_of(t.get_shape().size(), 3lu, 4lu); + return t.get_partial_shape().rank().is_static() && ov::snippets::utils::one_of(t.get_partial_shape().size(), 3lu, 4lu); } bool is_supported_intermediate_op(const std::shared_ptr& node) { @@ -68,6 +68,10 @@ void tokenize_broadcast(const std::shared_ptr& interm_op, ov::NodeVect // TODO: Can we reuse AppropriateForSubgraph here? Seems like it's huge check for Broadcast if (broadcast && broadcast->get_broadcast_spec().m_type == ov::op::AutoBroadcastType::NUMPY && broadcast->get_output_target_inputs(0).size() == 1) { + // TODO: Add support of Broadcast with ShapeOf subgraph on second input + if (!ov::is_type(broadcast->input_value(1).get_node_shared_ptr())) + continue; + broadcast_nodes.push_back(broadcast); const auto pshape = broadcast->get_input_partial_shape(0); @@ -96,10 +100,17 @@ void tokenize_broadcast(const std::shared_ptr& interm_op, ov::NodeVect bool tokenize_reshape_around_softmax(std::shared_ptr& interm_op, std::shared_ptr& reshape, ov::NodeVector& ordered_ops) { reshape = ov::as_type_ptr(interm_op); if (reshape) { - const auto in_shape = reshape->get_input_shape(0); - const auto out_shape = reshape->get_output_shape(0); - if (in_shape.back() != out_shape.back() || reshape->get_output_target_inputs(0).size() != 1) + // TODO: Add support of Reshape with ShapeOf subgraph on second input + if (!ov::is_type(reshape->input_value(1).get_node_shared_ptr())) return false; + + const auto in_shape = reshape->get_input_partial_shape(0); + const auto out_shape = reshape->get_output_partial_shape(0); + const auto in_last_dim = *in_shape.crbegin(); + const auto out_last_dim = *out_shape.crbegin(); + if (in_last_dim.is_dynamic() || out_last_dim.is_dynamic() || in_last_dim != out_last_dim || reshape->get_output_target_inputs(0).size() != 1) + return false; + ordered_ops.push_back(reshape); interm_op = reshape->get_output_target_inputs(0).begin()->get_node()->shared_from_this(); } @@ -204,8 +215,7 @@ bool ov::snippets::pass::TokenizeMHASnippets::is_matmul0_supported(const std::sh ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsTokenization::Config& config) { MATCHER_SCOPE(TokenizeMHASnippets); - auto m_matmul0 = std::make_shared(ov::pass::pattern::any_input(ov::pass::pattern::has_static_shape()), - ov::pass::pattern::any_input(ov::pass::pattern::has_static_shape())); + auto m_matmul0 = std::make_shared(ov::pass::pattern::any_input(), ov::pass::pattern::any_input()); register_matcher(std::make_shared(m_matmul0, matcher_name), [OV_CAPTURE_CPY_AND_THIS](ov::pass::pattern::Matcher &m) { @@ -224,20 +234,14 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken // Example: // Buffer - i32 [32, 128] -> ~ Loop ~ -> Buffer - i8 [32, 128] // After each Loop iteration we should increment pointers of Buffers: accordingly on 4 byte and 1 byte for scalar case. - // It means that these Buffers cannot be inplace => Each Buffer should have the own register + // It means that these increments are not proportional => Each Buffer should have the own register // For that we can just check the following "branches": // - Between MatMul0 and MatMul1 - Softmax is sync point. The operations between MatMul0 -> Softmax and Softmax -> MatMul1 // will be fused into one loop after conversion to snippet dialect (Because it's just FQ, Eltwise nodes) - // - Between MatMul0 and Transpose1 - At the moment operations after Transpose1 cannot be fused in Transpose Loop (to avoid performance regressions). + // - Between MatMul0 and Transpose1 - At the moment operations after Transpose1 cannot be fused in inner Transpose Loop + // (to avoid performance regressions due to scalar calculations). // But operations after Transpose1 and before MatMul0 will be fused into one loop as well (look at first point) - // Note: If the pass is updated, need to check the new possible branches for potential non-inplace Buffers! - // Default value is 2 because - // - Firstly, Softmax always needs Buffers - // - Secondly, Softmax needs 2 Buffers but they can be inplace - One virtual port is enough for Softmax => buffer_count = 1 - // - Thirdly, MatMul requires unique Buffers on inputs and outputs because blocking implementation increments input/output pointers during computations - // However, all of the Buffers are usually reused by the next MatMul and Softmax. - // So on sufficiently large subgraphs we use only one additional unique buffer => buffer_count increments by 1 - size_t buffer_count = 2; + size_t uniqie_buffer_reg_group_count = 1; // After MatMul0 there is always one Buffer std::string fused_names; ov::NodeVector ordered_ops; @@ -260,24 +264,20 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken if (!is_matmul0_supported(matmul0)) return false; - const auto matmul0_prc = op::Brgemm::get_output_type(matmul0->get_input_element_type(0), matmul0->get_input_element_type(1)); - // Between MatMul0 and Softmax will be the one Loop because of LoopFusing optimization. - // The Loop will have one Buffer with the same shape both on input and output. - // Need to check for precision to get if we need one more register for Buffer - if (matmul0_prc.size() != ov::element::f32.size()) { - if (buffer_count < 2) - buffer_count++; - } - ordered_ops.push_back(matmul0); const auto pattern_rank = matmul0->get_output_partial_shape(0).size(); + const auto ops_count_before_softmax = ordered_ops.size(); auto interm_op = matmul0->get_output_target_inputs(0).begin()->get_node()->shared_from_this(); // Add supported operations which are between MatMul0 and Softmax to ordered_ops if (!update_intermediate_supported_ops(interm_op, ordered_ops, hidden_virtual_ports_count, potential_body_params_count)) return false; + // If before Softmax there is Eltwise ops, there will be one more Buffer + if (ops_count_before_softmax != ordered_ops.size() && interm_op->get_output_partial_shape(0).rbegin()->is_dynamic()) + uniqie_buffer_reg_group_count++; + std::shared_ptr reshape0 = nullptr; if (!tokenize_reshape_around_softmax(interm_op, reshape0, ordered_ops)) return false; @@ -294,6 +294,11 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken if (axis != rank.get_length() - 1 || interm_op->get_output_target_inputs(0).size() != 1) return false; + + // Softmax need one buffer at least + if (interm_op->get_output_partial_shape(0).rbegin()->is_dynamic()) + uniqie_buffer_reg_group_count++; + ordered_ops.push_back(interm_op); interm_op = interm_op->get_output_target_inputs(0).begin()->get_node()->shared_from_this(); @@ -302,7 +307,7 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken return false; if (((reshape0 == nullptr) != (reshape1 == nullptr)) || - (reshape0 && reshape1 && (reshape0->get_input_shape(0) != reshape1->get_output_shape(0)))) + (reshape0 && reshape1 && (reshape0->get_input_partial_shape(0) != reshape1->get_output_partial_shape(0)))) return false; // Add supported operations which are between Softmax and MatMul1 to ordered_ops @@ -310,8 +315,7 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken return false; const auto matmul1 = ov::as_type_ptr(interm_op); - if (!matmul1 || matmul1->get_output_target_inputs(0).size() != 1 || - matmul1->get_transpose_a() || matmul1->get_transpose_b()) + if (!matmul1 || matmul1->get_transpose_a() || matmul1->get_transpose_b()) return false; const auto matmul1_out_type = op::Brgemm::get_output_type(matmul1->get_input_element_type(0), @@ -328,8 +332,9 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken // Between Softmax and MatMul1 will be the one Loop because of LoopFusing optimization. // The Loop will have one Buffer with the same shape both on input and output. // Need to check for precision to get if we need one more register for Buffer - if (matmul1->get_input_element_type(0).size() != ov::element::f32.size()) { - buffer_count++; + const auto matmul0_prc = op::Brgemm::get_output_type(matmul0->get_input_element_type(0), matmul0->get_input_element_type(1)); + if (matmul1->get_input_element_type(0).size() != matmul0_prc.size() || matmul1->get_input_partial_shape(0).is_dynamic()) { + uniqie_buffer_reg_group_count++; } /***********************/ @@ -358,6 +363,7 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken // There is transformation ExplicitTransposeMatMulInputs that set supported order and transposed_b(false). // We can allow to call this pass only if ops have scalar shapes to avoid shape mismatching const auto is_transposed_b_0 = matmul0->get_transpose_b(); + bool has_matmul0_has_ops_on_input = false; while (is_supported_intermediate_op(parent)) { // All supported ops have only one output port if (parent->get_output_target_inputs(0).size() != 1) @@ -379,6 +385,11 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken ordered_ops.insert(ordered_ops.begin(), parent); // [107731] To go always through 0-th port - is it safe? parent = parent->get_input_node_shared_ptr(0); + has_matmul0_has_ops_on_input = true; + } + // If there are ops on second input of MatMul0 -> there always will be unique Buffer + if (has_matmul0_has_ops_on_input) { + uniqie_buffer_reg_group_count++; } auto tokenize_transpose = [&](const std::shared_ptr& transpose, @@ -412,7 +423,9 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken bool are_ops_after_matmul1 = false; auto child = matmul1->get_output_target_inputs(0).begin()->get_node()->shared_from_this(); - while (is_supported_intermediate_op(child)) { + const auto can_be_ops_after_matmul1_tokenized = matmul1->get_output_target_inputs(0).size() == 1; + bool has_matmul1_has_ops_on_output = false; + while (can_be_ops_after_matmul1_tokenized && is_supported_intermediate_op(child)) { are_ops_after_matmul1 = true; // All supported ops have only one output port if (child->get_output_target_inputs(0).size() != 1) @@ -427,19 +440,23 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken // TODO [75567]: move this plugin-specific constraint to the plugin callback // We cannot collapse op to Subgraph if count of potential Parameter and Result count is higher 12 - if (potential_body_params_count + child->get_output_target_inputs(0).size() + hidden_virtual_ports_count + buffer_count > 12) { + if (potential_body_params_count + child->get_output_target_inputs(0).size() + hidden_virtual_ports_count + uniqie_buffer_reg_group_count > 12) { break; } ordered_ops.push_back(child); child = child->get_output_target_inputs(0).begin()->get_node()->shared_from_this(); + has_matmul1_has_ops_on_output = true; + } + if (has_matmul1_has_ops_on_output) { + uniqie_buffer_reg_group_count++; } // At the moment Snippets don't support nodes between MatMul1 and Transpose3 due to Loop and strided calculations limitations // MatMul1 // // Transpose3 - if (!are_ops_after_matmul1) { + if (can_be_ops_after_matmul1_tokenized && !are_ops_after_matmul1) { auto transpose3 = config.get_mha_token_enable_transpose_on_output() ? ov::as_type_ptr(child) : nullptr; if (is_valid_transpose(transpose3, config.get_mha_supported_transpose_ranks(), get_fusion_transpose_order(pattern_rank)) && transpose3->get_input_element_type(0) == matmul1_out_type) { // To avoid Convert between MatMul1 and Transpose3 @@ -455,7 +472,7 @@ ov::snippets::pass::TokenizeMHASnippets::TokenizeMHASnippets(const SnippetsToken // TODO [75567]: move this plugin-specific constraint to the plugin callback const auto last_node = ordered_ops.back(); - if (potential_body_params_count + last_node->get_output_size() + hidden_virtual_ports_count + buffer_count > 11) { + if (potential_body_params_count + last_node->get_output_size() + hidden_virtual_ports_count + uniqie_buffer_reg_group_count > 11) { return false; } diff --git a/src/common/snippets/tests/src/pass/mha_tokenization.cpp b/src/common/snippets/tests/src/pass/mha_tokenization.cpp index 6438dff516c..b411aace066 100644 --- a/src/common/snippets/tests/src/pass/mha_tokenization.cpp +++ b/src/common/snippets/tests/src/pass/mha_tokenization.cpp @@ -39,6 +39,30 @@ TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_4D) { run(); } +TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_4D_Dynamic) { + const auto &f = MHAFunction(std::vector{{-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}}, + std::vector({ov::element::f32, ov::element::f32, ov::element::f32, ov::element::f32}), true, false); + model = f.getOriginal(); + model_ref = f.getReference(); + run(); +} + +TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_4D_Dynamic_M) { + const auto &f = MHAFunction(std::vector{{1, -1, 12, 64}, {1, 128, 12, 64}, {1, 12, -1, 128}, {1, 128, 12, 64}}, + std::vector({ov::element::f32, ov::element::f32, ov::element::f32, ov::element::f32}), true, false); + model = f.getOriginal(); + model_ref = f.getReference(); + run(); +} + +TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_4D_Dynamic_K) { + const auto &f = MHAFunction(std::vector{{1, 128, 12, -1}, {1, 128, 12, -1}, {1, 12, 128, 128}, {1, 128, 12, 64}}, + std::vector({ov::element::f32, ov::element::f32, ov::element::f32, ov::element::f32}), true, false); + model = f.getOriginal(); + model_ref = f.getReference(); + run(); +} + TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_3D) { const auto &f = MHAFunction(std::vector{{128, 12, 64}, {128, 12, 64}, {12, 128, 128}, {128, 12, 64}}, std::vector({ov::element::f32, ov::element::f32, ov::element::f32, ov::element::f32})); @@ -47,8 +71,15 @@ TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_3D) { run(); } -TEST_F(SKIP_TokenizeMHASnippetsTests /* CVS-114607 */, smoke_Snippets_MHA_with_MatMul0_Transpose) { - GTEST_SKIP(); +TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_3D_Dynamic) { + const auto &f = MHAFunction(std::vector{{-1, -1, -1}, {-1, -1, -1}, {-1, -1, -1}, {-1, -1, -1}}, + std::vector({ov::element::f32, ov::element::f32, ov::element::f32, ov::element::f32}), true, false); + model = f.getOriginal(); + model_ref = f.getReference(); + run(); +} + +TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_with_MatMul0_Transpose) { const auto &f = MHAMatMul0TransposeFunction(std::vector{{1, 128, 12, 64}, {1, 128, 12, 64}, {1, 12, 128, 128}, {1, 128, 12, 64}}, std::vector({ov::element::f32, ov::element::f32, ov::element::f32, ov::element::f32})); model = f.getOriginal(); @@ -56,6 +87,16 @@ TEST_F(SKIP_TokenizeMHASnippetsTests /* CVS-114607 */, smoke_Snippets_MHA_with_M run(); } +TEST_F(SKIP_TokenizeMHASnippetsTests /* CVS-142098 */, smoke_Snippets_MHA_with_MatMul0_Transpose_Dynamic) { + GTEST_SKIP(); + const auto &f = MHAMatMul0TransposeFunction(std::vector{{-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}}, + std::vector({ov::element::f32, ov::element::f32, ov::element::f32, ov::element::f32}), + false); + model = f.getOriginal(); + model_ref = f.getReference(); + run(); +} + TEST_F(SKIP_TokenizeMHASnippetsTests /* CVS-114607 */, smoke_Snippets_MHA_with_int_Matmuls) { GTEST_SKIP(); const auto &f = MHAINT8MatMulTypeRelaxedFunction(std::vector{{1, 128, 12, 64}, {1, 128, 12, 64}, {1, 12, 128, 128}, {1, 128, 12, 64}}); @@ -71,6 +112,14 @@ TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_Transpose_extraction) { run(); } +TEST_F(SKIP_TokenizeMHASnippetsTests /* CVS-142098 */, smoke_Snippets_MHA_Dynamic_Transpose_extraction) { + GTEST_SKIP(); + const auto& f = MHATransposedInputFunction(std::vector{{-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}}, true); + model = f.getOriginal(); + model_ref = f.getReference(); + run(); +} + TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_Transpose_extraction_and_unsupported_existing_transpose) { const auto& f = MHATransposedInputFunction(std::vector{{1, 128, 12, 64}, {1, 12, 64, 128}, {1, 128, 12, 64}}, true, std::vector{0, 3, 1, 2}); @@ -79,6 +128,15 @@ TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_Transpose_extraction_and_uns run(); } +TEST_F(SKIP_TokenizeMHASnippetsTests /* CVS-142098 */, smoke_Snippets_MHA_Dynamic_Transpose_extraction_and_unsupported_existing_transpose) { + GTEST_SKIP(); + const auto& f = MHATransposedInputFunction(std::vector{{-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}}, true, + std::vector{0, 3, 1, 2}); + model = f.getOriginal(); + model_ref = f.getReference(); + run(); +} + TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_Transpose_fusion) { const auto& f = MHATransposedInputFunction(std::vector{{1, 128, 12, 64}, {1, 64, 128, 12}, {1, 128, 12, 64}}, false, std::vector{0, 2, 1, 3}); @@ -87,6 +145,14 @@ TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_Transpose_fusion) { run(); } +TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA_Dyanmic_Transpose_fusion) { + const auto& f = MHATransposedInputFunction(std::vector{{-1, -1, -1, -1}, {-1, -1, -1, -1}, {-1, -1, -1, -1}}, false, + std::vector{0, 2, 1, 3}); + model = f.getOriginal(); + model_ref = f.getReference(); + run(); +} + TEST_F(TokenizeMHASnippetsTests, smoke_Snippets_MHA3D_SplitM) { const auto& f = MHASplitMFunction(std::vector{{128, 12, 64}, {128, 12, 64}, {12, 128, 128}, {128, 12, 64}}, std::vector({ov::element::f32, ov::element::f32, ov::element::f32, ov::element::f32}), diff --git a/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp b/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp index 3d13cab76db..006935a85e8 100644 --- a/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp +++ b/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp @@ -867,7 +867,7 @@ void Transformations::MainSnippets(void) { #if defined(OPENVINO_ARCH_X86_64) auto is_supported_matmul = [this](const std::shared_ptr& n) { const auto matmul = ov::as_type_ptr(n); - if (!matmul) + if (!matmul || matmul->is_dynamic()) return false; const auto in_type0 = matmul->get_input_element_type(0); const auto in_type1 = matmul->get_input_element_type(1); diff --git a/src/tests/ov_helpers/ov_snippets_models/include/subgraph_mha.hpp b/src/tests/ov_helpers/ov_snippets_models/include/subgraph_mha.hpp index 3e3cb995b70..9ec265c9322 100644 --- a/src/tests/ov_helpers/ov_snippets_models/include/subgraph_mha.hpp +++ b/src/tests/ov_helpers/ov_snippets_models/include/subgraph_mha.hpp @@ -42,8 +42,9 @@ namespace snippets { */ class MHAFunction : public SnippetsFunctionBase { public: - explicit MHAFunction(const std::vector& inputShapes, const std::vector& precisions, bool with_mul = true) - : SnippetsFunctionBase(inputShapes), with_mul(with_mul), precisions(precisions) { + explicit MHAFunction(const std::vector& inputShapes, const std::vector& precisions, + bool with_mul = true, bool with_reshape = true) + : SnippetsFunctionBase(inputShapes), with_mul(with_mul), with_reshape(with_reshape), precisions(precisions) { OPENVINO_ASSERT(input_shapes.size() == 4, "Got invalid number of input shapes"); OPENVINO_ASSERT(precisions.size() == 4, "Got invalid number of input precisions"); } @@ -51,8 +52,9 @@ protected: std::shared_ptr initOriginal() const override; std::shared_ptr initReference() const override; - bool with_mul = true; - std::vector precisions; + const bool with_mul = true; + const bool with_reshape = true; + const std::vector precisions; }; class MHASplitMFunction : public MHAFunction { @@ -85,8 +87,9 @@ protected: */ class MHAMatMul0TransposeFunction : public SnippetsFunctionBase { public: - explicit MHAMatMul0TransposeFunction(const std::vector& inputShapes, const std::vector& precisions) - : SnippetsFunctionBase(inputShapes), precisions(precisions) { + explicit MHAMatMul0TransposeFunction(const std::vector& inputShapes, const std::vector& precisions, + bool with_reshape = true) + : SnippetsFunctionBase(inputShapes), with_reshape(with_reshape), precisions(precisions) { OPENVINO_ASSERT(input_shapes.size() == 4, "Got invalid number of input shapes"); OPENVINO_ASSERT(precisions.size() == 4, "Got invalid number of input precisions"); } @@ -94,7 +97,8 @@ protected: std::shared_ptr initOriginal() const override; std::shared_ptr initReference() const override; - std::vector precisions; + const bool with_reshape = true; + const std::vector precisions; }; /* Graph: diff --git a/src/tests/ov_helpers/ov_snippets_models/src/subgraph_mha.cpp b/src/tests/ov_helpers/ov_snippets_models/src/subgraph_mha.cpp index 3157c53fbb3..f923a9a3aa1 100644 --- a/src/tests/ov_helpers/ov_snippets_models/src/subgraph_mha.cpp +++ b/src/tests/ov_helpers/ov_snippets_models/src/subgraph_mha.cpp @@ -70,26 +70,35 @@ std::shared_ptr MHAFunction::initOriginal() const { std::shared_ptr matmul_parent1 = transpose1; if (with_mul) { ov::Shape shape(rank, 1); - shape[rank - 3] = transpose1->get_output_shape(0)[rank - 3]; - std::vector mulConstData(ov::shape_size(shape)); + if (transpose1->get_output_partial_shape(0).is_static()) { + shape[rank - 3] = transpose1->get_output_shape(0)[rank - 3]; + } const auto mulConst = ov::test::utils::make_constant(precisions[1], shape); matmul_parent1 = std::make_shared(transpose1, mulConst); } const auto matMul0 = std::make_shared(transpose0, matmul_parent1); const auto add = std::make_shared(matMul0, addParam); - const auto interm_shape = add->get_output_shape(0); - const auto batch = std::accumulate(interm_shape.cbegin(), interm_shape.cbegin() + rank - 1, 1, std::multiplies()); - const auto reshape0ConstData = std::vector{ batch, -1 }; - const auto reshape1ConstData = interm_shape; - const auto reshape0Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{reshape0ConstData.size()}, reshape0ConstData); - const auto reshape1Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{reshape1ConstData.size()}, reshape1ConstData); + auto softmax_out = add->output(0); + if (with_reshape) { + const auto interm_shape = add->get_output_shape(0); + const auto batch = std::accumulate(interm_shape.cbegin(), interm_shape.cbegin() + rank - 1, 1, std::multiplies()); + const auto reshape0ConstData = std::vector{ batch, -1 }; + const auto reshape1ConstData = interm_shape; + const auto reshape0Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{reshape0ConstData.size()}, reshape0ConstData); + const auto reshape1Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{reshape1ConstData.size()}, reshape1ConstData); + + const auto reshape0 = std::make_shared(add, reshape0Const, true); + const auto softMax = std::make_shared(reshape0, 1); + const auto reshape1 = std::make_shared(softMax, reshape1Const, true); + softmax_out = reshape1->output(0); + } else { + const auto softMax = std::make_shared(add, rank - 1); + softmax_out = softMax->output(0); + } - const auto reshape0 = std::make_shared(add, reshape0Const, true); - const auto softMax = std::make_shared(reshape0, 1); - const auto reshape1 = std::make_shared(softMax, reshape1Const, true); const auto transpose2 = std::make_shared(transpose2Param, transpose2Const); - const auto matMul1 = std::make_shared(reshape1, transpose2); + const auto matMul1 = std::make_shared(softmax_out, transpose2); const auto transpose3 = std::make_shared(matMul1, transpose3Const); ov::ResultVector results{std::make_shared(transpose3)}; @@ -124,13 +133,19 @@ std::shared_ptr MHAFunction::initReference() const { std::shared_ptr matmul_parent1 = transpose1; if (with_mul) { ov::Shape shape(rank, 1); - shape[rank - 3] = transpose1->get_output_shape(0)[rank - 3]; - std::vector mulConstData(ov::shape_size(shape)); + if (transpose1->get_output_partial_shape(0).is_static()) { + shape[rank - 3] = transpose1->get_output_shape(0)[rank - 3]; + } const auto mulConst = ov::test::utils::make_constant(precisions[1], shape); - const auto mulParam = std::make_shared(precisions[1], mulConst->get_shape()); - matmul_parent1 = std::make_shared(transpose1, mulParam); - subgraph_params = {transpose0Param, transpose1Param, mulParam, addParam, transpose2Param}; - subgraph_inputs = {data0, data1, mulConst, data2, data3}; + + if (ov::shape_size(shape) > 1) { + const auto mulParam = std::make_shared(precisions[1], mulConst->get_shape()); + matmul_parent1 = std::make_shared(transpose1, mulParam); + subgraph_params = {transpose0Param, transpose1Param, mulParam, addParam, transpose2Param}; + subgraph_inputs = {data0, data1, mulConst, data2, data3}; + } else { + matmul_parent1 = std::make_shared(transpose1, mulConst); + } } const auto matMul0 = std::make_shared(transpose0, matmul_parent1); @@ -182,16 +197,22 @@ std::shared_ptr MHASplitMFunction::initReference() const { std::shared_ptr matmul_parent1 = transpose1; if (with_mul) { ov::Shape shape(rank - 1, 1); - shape[rank - 4] = transpose1->get_output_shape(0)[rank - 4]; - ov::Shape reshape_shape = shape; - reshape_shape.insert(reshape_shape.cbegin() + rank - 3, 1); - std::vector mulConstData(ov::shape_size(shape)); + if (transpose1->get_output_partial_shape(0).is_static()) { + shape[rank - 4] = transpose1->get_output_shape(0)[rank - 4]; + } const auto mulConst = ov::test::utils::make_constant(precisions[1], shape); - const auto reshape_mul = make_reshape(mulConst, reshape_shape); - const auto mulParam = std::make_shared(precisions[1], reshape_mul->get_shape()); - matmul_parent1 = std::make_shared(transpose1, mulParam); - subgraph_params = {transpose0Param, transpose1Param, mulParam, addParam, transpose2Param}; - subgraph_inputs = {reshape0, reshape1, reshape_mul, reshape2, reshape3}; + + if (ov::shape_size(shape) > 1) { + ov::Shape reshape_shape = shape; + reshape_shape.insert(reshape_shape.cbegin() + rank - 3, 1); + const auto mulReshape = make_reshape(mulConst, reshape_shape); + const auto mulParam = std::make_shared(precisions[1], mulReshape->get_shape()); + matmul_parent1 = std::make_shared(transpose1, mulParam); + subgraph_params = {transpose0Param, transpose1Param, mulParam, addParam, transpose2Param}; + subgraph_inputs = {reshape0, reshape1, mulReshape, reshape2, reshape3}; + } else { + matmul_parent1 = std::make_shared(transpose1, mulConst); + } } const auto matMul0 = std::make_shared(transpose0, matmul_parent1); @@ -217,47 +238,42 @@ std::shared_ptr MHAMatMul0TransposeFunction::initOriginal() const { auto transpose2Param = std::make_shared(precisions[3], input_shapes[3]); ov::ParameterVector ngraphParam = {transpose0Param, transpose1Param, addParam, transpose2Param}; - std::vector constantShapes; - constantShapes.push_back(ov::Shape({input_shapes[0].get_shape().size()})); - constantShapes.push_back(ov::Shape({input_shapes[0].get_shape().size()})); - constantShapes.push_back(ov::Shape({1, input_shapes[1].get_shape()[2], 1, 1})); - constantShapes.push_back(ov::Shape({2})); - constantShapes.push_back(ov::Shape({4})); - constantShapes.push_back(ov::Shape({input_shapes[0].get_shape().size()})); - constantShapes.push_back(ov::Shape({input_shapes[0].get_shape().size()})); + const auto rank = input_shapes[0].size(); + const auto fusion_order = get_fusion_order(rank); - const auto order = std::vector{0, 2, 1, 3}; - auto transpose0Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[0], order); - auto transpose1Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[1], order); - auto transpose2Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[5], order); - auto transpose3Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[6], order); + const auto transpose0Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{rank}, fusion_order); + const auto transpose1Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{rank}, fusion_order); + const auto transpose2Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{rank}, fusion_order); + const auto transpose3Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{rank}, fusion_order); - std::vector mulConstData(1); - auto mulConst = ov::test::utils::make_constant(precisions[1], ov::Shape{1}); - - std::vector reshape0ConstData = {static_cast(input_shapes[0].get_shape()[0] * - input_shapes[0].get_shape()[1] * input_shapes[0].get_shape()[2]), - -1}; - auto reshape0Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[3], reshape0ConstData); - - std::vector reshape1ConstData = {static_cast(input_shapes[0].get_shape()[0]), - static_cast(input_shapes[0].get_shape()[2]), - static_cast(input_shapes[0].get_shape()[1]), - static_cast(input_shapes[0].get_shape()[1])}; - auto reshape1Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[4], reshape1ConstData); - - float transA = false; - float transB = false; const auto transpose0 = std::make_shared(transpose0Param, transpose0Const); const auto transpose1 = std::make_shared(transpose1Param, transpose1Const); + + const auto mulConst = ov::test::utils::make_constant(precisions[1], ov::Shape{1}); const auto mul = std::make_shared(transpose1, mulConst); - const auto matMul0 = std::make_shared(transpose0, mul, transA, true); + const auto matMul0 = std::make_shared(transpose0, mul, false, true); const auto add = std::make_shared(matMul0, addParam); - const auto reshape0 = std::make_shared(add, reshape0Const, true); - const auto softMax = std::make_shared(reshape0, 1); - const auto reshape1 = std::make_shared(softMax, reshape1Const, true); + + auto softmax_out = add->output(0); + if (with_reshape) { + const auto interm_shape = add->get_output_shape(0); + const auto batch = std::accumulate(interm_shape.cbegin(), interm_shape.cbegin() + rank - 1, 1, std::multiplies()); + const auto reshape0ConstData = std::vector{ batch, -1 }; + const auto reshape1ConstData = interm_shape; + const auto reshape0Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{reshape0ConstData.size()}, reshape0ConstData); + const auto reshape1Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{reshape1ConstData.size()}, reshape1ConstData); + + const auto reshape0 = std::make_shared(add, reshape0Const, true); + const auto softMax = std::make_shared(reshape0, 1); + const auto reshape1 = std::make_shared(softMax, reshape1Const, true); + softmax_out = reshape1->output(0); + } else { + const auto softMax = std::make_shared(add, rank - 1); + softmax_out = softMax->output(0); + } + const auto transpose2 = std::make_shared(transpose2Param, transpose2Const); - const auto matMul1 = std::make_shared(reshape1, transpose2, transA, transB); + const auto matMul1 = std::make_shared(softmax_out, transpose2); const auto transpose3 = std::make_shared(matMul1, transpose3Const); ov::ResultVector results{std::make_shared(transpose3)}; @@ -269,58 +285,38 @@ std::shared_ptr MHAMatMul0TransposeFunction::initReference() const { auto data2 = std::make_shared(precisions[2], input_shapes[2]); auto data3 = std::make_shared(precisions[3], input_shapes[3]); ov::ParameterVector ngraphParams = {data0, data1, data2, data3}; + NodeVector subgraph_inputs = {data0, data1, data2, data3}; auto transpose0Param = std::make_shared(precisions[0], input_shapes[0]); auto transpose1Param = std::make_shared(precisions[1], input_shapes[1]); auto addParam = std::make_shared(precisions[2], input_shapes[2]); auto transpose2Param = std::make_shared(precisions[3], input_shapes[3]); - std::vector constantShapes; - constantShapes.push_back(ov::Shape({input_shapes[0].get_shape().size()})); - constantShapes.push_back(ov::Shape({input_shapes[0].get_shape().size()})); - constantShapes.push_back(ov::Shape({1, input_shapes[1].get_shape()[2], 1, 1})); - constantShapes.push_back(ov::Shape({2})); - constantShapes.push_back(ov::Shape({4})); - constantShapes.push_back(ov::Shape({input_shapes[0].get_shape().size()})); - constantShapes.push_back(ov::Shape({input_shapes[0].get_shape().size()})); + ov::ParameterVector subgraph_params = {transpose0Param, transpose1Param, addParam, transpose2Param}; - auto transpose0Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[0], std::vector{0, 2, 1, 3}); - auto transpose1Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[1], std::vector{0, 2, 3, 1}); - auto transpose2Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[5], std::vector{0, 2, 1, 3}); - auto transpose3Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[6], std::vector{0, 2, 1, 3}); + const auto rank = input_shapes[0].size(); + const auto fusion_order = get_fusion_order(rank); + const auto decomposed_order = get_decomposed_order(rank); - std::vector mulConstData(1); - auto mulConst = ov::test::utils::make_constant(precisions[1], ov::Shape{1}); - ov::ParameterVector subgraphParams = {transpose0Param, transpose1Param, addParam, transpose2Param}; + const auto transpose0Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{rank}, fusion_order); + const auto transpose1Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{rank}, decomposed_order); + const auto transpose2Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{rank}, fusion_order); + const auto transpose3Const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{rank}, fusion_order); - std::vector reshape0ConstData = {static_cast(input_shapes[0].get_shape()[0] * - input_shapes[0].get_shape()[1] * input_shapes[0].get_shape()[2]), - -1}; - auto reshape0Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[3], reshape0ConstData); - - std::vector reshape1ConstData = {static_cast(input_shapes[0].get_shape()[0]), - static_cast(input_shapes[0].get_shape()[2]), - static_cast(input_shapes[0].get_shape()[1]), - static_cast(input_shapes[0].get_shape()[1])}; - auto reshape1Const = ov::op::v0::Constant::create(ov::element::i64, constantShapes[4], reshape1ConstData); - - float transA = false; - float transB = false; const auto transpose0 = std::make_shared(transpose0Param, transpose0Const); const auto transpose1 = std::make_shared(transpose1Param, transpose1Const); + + const auto mulConst = ov::test::utils::make_constant(precisions[1], ov::Shape{1}); const auto mul = std::make_shared(transpose1, mulConst); - const auto matMul0 = std::make_shared(transpose0, mul, transA, transB); + const auto matMul0 = std::make_shared(transpose0, mul); const auto add = std::make_shared(matMul0, addParam); - const auto reshape0 = std::make_shared(add, reshape0Const, true); - const auto softMax = std::make_shared(reshape0, 1); - const auto reshape1 = std::make_shared(softMax, reshape1Const, true); + const auto softMax = std::make_shared(add, rank - 1); const auto transpose2 = std::make_shared(transpose2Param, transpose2Const); - const auto matMul1 = std::make_shared(reshape1, transpose2, transA, transB); + const auto matMul1 = std::make_shared(softMax, transpose2); const auto transpose3 = std::make_shared(matMul1, transpose3Const); - auto subgraph = std::make_shared( - NodeVector{data0, data1, data2, data3}, - std::make_shared(NodeVector{transpose3}, subgraphParams)); + auto subgraph = std::make_shared(subgraph_inputs, + std::make_shared(NodeVector{transpose3}, subgraph_params)); return std::make_shared(NodeVector{subgraph}, ngraphParams); } @@ -982,9 +978,9 @@ std::shared_ptr MHATransposedInputFunction::initReference() const { } } - const auto param0 = std::make_shared(precision, data0->get_shape()); - const auto param1 = std::make_shared(precision, in1->get_shape()); - const auto param2 = std::make_shared(precision, data2->get_shape()); + const auto param0 = std::make_shared(precision, data0->get_output_partial_shape(0)); + const auto param1 = std::make_shared(precision, in1->get_output_partial_shape(0)); + const auto param2 = std::make_shared(precision, data2->get_output_partial_shape(0)); std::shared_ptr matmul0_in1 = param1; if (!m_order.empty() && is_supported) {