[Snippets] Improve Subgraph tokenization (#31900)

### Details:
 - *Error-prone constructor is removed from Subgraph node*
- *`tokenize_ordered_nodes` helper: avoid Parameters duplication in case
of shared external input*

### Tickets:
 - *N\A*

Author: v-Golubev

src/common/snippets/include/snippets/op/subgraph.hpp

@@ -92,8 +92,6 @@ class Subgraph : public ov::op::util::SubGraphOp {
 
     explicit Subgraph(const OutputVector& args, const std::shared_ptr<ov::Model>& body);
 
-    explicit Subgraph(const NodeVector& args, const std::shared_ptr<ov::Model>& body);
-
     bool visit_attributes(AttributeVisitor& visitor) override;
 
     void validate_and_infer_types() override;

src/common/snippets/include/snippets/utils/tokenization_utils.hpp

@@ -31,7 +31,13 @@ bool tokenize_node(const std::shared_ptr<ov::Node>& node,
  *        1. The user is responsible for valid count of parameters, results and hidden virtual ports (constants)
  *        2. The list of nodes cannot contain Subgraph ops
  * @param ordered_ops node list which should be tokenized
+ * @param are_shared_internal_params_allowed if true, allows sharing internal parameters.
+ * Note: Snippets support only internal parameters which are used by all the consumers as is.
+ * This means that e.g. if the shared parameter is used by 2 MatMuls on B input,
+ * both matmuls must have equal transpose_b parameter.
+ * This is a user responsibility to ensure that the shared internal parameters can be used.
  * @return tokenized subgraph
  */
-std::shared_ptr<ov::snippets::op::Subgraph> tokenize_ordered_nodes(const ov::NodeVector& ordered_ops);
+std::shared_ptr<ov::snippets::op::Subgraph> tokenize_ordered_nodes(const ov::NodeVector& ordered_ops,
+                                                                   bool are_shared_internal_params_allowed = false);
 }  // namespace ov::snippets::utils

src/common/snippets/src/op/subgraph.cpp

@@ -219,9 +219,6 @@ Subgraph::Subgraph(const OutputVector& args, const std::shared_ptr<ov::Model>& b
     m_shape_infer = std::make_shared<OVShapeInfer>(body);
 }
 
-Subgraph::Subgraph(const NodeVector& args, const std::shared_ptr<ov::Model>& body)
-    : Subgraph(as_output_vector(args), body) {}
-
 std::shared_ptr<Node> Subgraph::clone_with_new_inputs(const OutputVector& inputs) const {
     INTERNAL_OP_SCOPE(Subgraph);
     return make_shared<Subgraph>(inputs, body().clone());

src/common/snippets/src/pass/gated_mlp_tokenization.cpp

@@ -6,6 +6,7 @@
 
 #include <memory>
 
+#include "openvino/core/except.hpp"
 #include "openvino/core/node.hpp"
 #include "openvino/core/node_output.hpp"
 #include "openvino/core/type.hpp"
@@ -116,7 +117,13 @@ TokenizeGatedMLPSnippets::TokenizeGatedMLPSnippets(const SnippetsTokenization::C
         }
 
         const auto ordered_ops = ov::NodeVector{fc_gate, fc_up, act, mul, fc_down};
-        const auto subgraph = ov::snippets::utils::tokenize_ordered_nodes(ordered_ops);
+        const bool allow_shared_params = [&]() {
+            const auto mm_gate = ov::as_type_ptr<ov::op::v0::MatMul>(fc_gate);
+            const auto mm_up = ov::as_type_ptr<ov::op::v0::MatMul>(fc_up);
+            OPENVINO_ASSERT(mm_gate && mm_up, "fc_gate and fc_up must have MatMul type");
+            return mm_gate->get_transpose_a() == mm_up->get_transpose_a();
+        }();
+        const auto subgraph = ov::snippets::utils::tokenize_ordered_nodes(ordered_ops, allow_shared_params);
 
         // mark the Subgraph as Completed to not allow Snippets to include any nodes into this Subgraph in common
         // Tokenization

src/common/snippets/src/utils/tokenization_utils.cpp

@@ -9,6 +9,7 @@
 #include <climits>
 #include <cstddef>
 #include <cstdint>
+#include <iterator>
 #include <map>
 #include <memory>
 #include <numeric>
@@ -442,56 +443,49 @@ bool tokenize_node(const std::shared_ptr<ov::Node>& node, const SnippetsTokeniza
     return true;
 }
 
-std::shared_ptr<ov::snippets::op::Subgraph> tokenize_ordered_nodes(const ov::NodeVector& ordered_ops) {
+std::shared_ptr<ov::snippets::op::Subgraph> tokenize_ordered_nodes(const ov::NodeVector& ordered_ops,
+                                                                   bool are_shared_internal_params_allowed) {
     OPENVINO_ASSERT(!ordered_ops.empty(), "Nothing to be tokenized!");
 
-    ov::OutputVector body_inputs, subgraph_inputs;
+    ov::OutputVector subgraph_inputs;
     ov::ParameterVector body_parameters;
-    ov::ResultVector body_results;
-    std::vector<std::set<Input<Node>>> subgraph_result_inputs;
-
     auto create_body_inputs = [&](const std::shared_ptr<ov::Node>& node) -> void {
         for (size_t i = 0; i < node->get_input_size(); ++i) {
             const auto input = node->input(i);
             const auto parent = input.get_source_output().get_node_shared_ptr();
             const auto constant = ov::as_type_ptr<ov::op::v0::Constant>(parent);
             if (constant && (ov::shape_size(input.get_shape()) == 1 || ov::is_type<ov::op::v0::FakeQuantize>(node) ||
                              op::Subgraph::constant_input_should_be_inside_body(node))) {
-                // If Constant has one consumer - target node, we add Constant to body_inputs
-                // If Constant has several consumers, we should check that all these consumers are inside Subgraph body
-                // and if all of them are inside body, we can explicitly add Constant to the body_inputs, otherwise we
-                // should make a copy and add copy of Constant to body_inputs For example, this case is especially valid
-                // for Transposes nodes
-                //     (several Transposes have the same order so there can be the common Constant with this order)
-                if (constant->get_output_target_inputs(0).size() == 1) {
-                    body_inputs.push_back(input.get_source_output());
-                } else {
+                // If not all Constant consumers are inside Subgraph body,
+                // we should make a copy of this Constant for Subgraph body.
+                if (constant->get_output_target_inputs(0).size() > 1) {
                     const auto constant_consumers = constant->get_output_target_inputs(0);
-                    bool all_consumers_are_inside =
-                        std::all_of(constant_consumers.begin(),
+                    bool has_external_consumers =
+                        std::any_of(constant_consumers.begin(),
                                     constant_consumers.end(),
                                     [&ordered_ops](const ov::Input<ov::Node>& input) {
                                         return std::find(ordered_ops.begin(),
                                                          ordered_ops.end(),
-                                                         input.get_node()->shared_from_this()) != ordered_ops.end();
+                                                         input.get_node()->shared_from_this()) == ordered_ops.end();
                                     });
-                    if (all_consumers_are_inside) {
-                        body_inputs.push_back(input.get_source_output());
-                    } else {
+                    if (has_external_consumers) {
                         const auto constant_copy = constant->clone_with_new_inputs({});
                         node->set_argument(input.get_index(), constant_copy);
-                        body_inputs.emplace_back(constant_copy);
                     }
                 }
             } else if (std::find(ordered_ops.begin(), ordered_ops.end(), parent) == ordered_ops.end()) {
-                auto parameter =
-                    std::make_shared<ov::opset1::Parameter>(input.get_element_type(), input.get_partial_shape());
-                body_parameters.push_back(parameter);
-                body_parameters.back()->set_friendly_name(input.get_node()->get_friendly_name());
-                body_inputs.push_back(parameter->output(0));
-
-                subgraph_inputs.push_back(input.get_source_output());
-
+                const auto& parent_output = input.get_source_output();
+                auto it = std::find(subgraph_inputs.begin(), subgraph_inputs.end(), parent_output);
+                if (!are_shared_internal_params_allowed || it == subgraph_inputs.end()) {
+                    auto new_param =
+                        std::make_shared<ov::op::v0::Parameter>(input.get_element_type(), input.get_partial_shape());
+                    new_param->set_friendly_name(input.get_node()->get_friendly_name());
+                    subgraph_inputs.push_back(parent_output);
+                    body_parameters.push_back(new_param);
+                    it = subgraph_inputs.end() - 1;
+                }
+                const auto param_index = static_cast<size_t>(std::distance(subgraph_inputs.begin(), it));
+                const auto& parameter = body_parameters[param_index];
                 node->input(i).replace_source_output(parameter);
             }
         }
@@ -511,20 +505,17 @@ std::shared_ptr<ov::snippets::op::Subgraph> tokenize_ordered_nodes(const ov::Nod
         }
     }
 
+    ov::ResultVector body_results;
+    std::vector<std::set<Input<Node>>> subgraph_result_inputs;
     for (const auto& output : last_node->outputs()) {
+        // Note: since we need to save only original consumers,
+        // subgraph_result_inputs must be taken before result creation
         subgraph_result_inputs.push_back(output.get_target_inputs());
-    }
-    for (const auto& output : last_node->outputs()) {
         body_results.push_back(std::make_shared<ov::opset1::Result>(last_node->output(output.get_index())));
     }
 
-    if (body_results.size() != subgraph_result_inputs.size()) {
-        OPENVINO_THROW("body results and node results size mismatch during subgraph collapse");
-    }
-
     auto body = op::create_body(last_node->get_friendly_name(), body_results, body_parameters);
     auto subgraph = std::make_shared<op::Subgraph>(subgraph_inputs, body);
-    // Copy runtime info from last node to subgraph - to copy topological order
     copy_runtime_info(last_node, subgraph);
     subgraph->set_friendly_name(last_node->get_friendly_name());
 

src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/x64/subgraph_serialize.cpp

@@ -35,7 +35,7 @@ TEST_F(SubgraphSnippetSerializationTest, smoke_SerializeSubgraph) {
         auto add = std::make_shared<Add>(ininput0, ininput1);
         auto subgraph_body =
             std::make_shared<ov::Model>(ov::OutputVector{add}, ov::ParameterVector{ininput0, ininput1});
-        auto subgraph = std::make_shared<ov::snippets::op::Subgraph>(ov::NodeVector{input0, input1}, subgraph_body.get()->clone());
+        auto subgraph = std::make_shared<ov::snippets::op::Subgraph>(ov::OutputVector{input0, input1}, subgraph_body.get()->clone());
         return std::make_shared<ov::Model>(ov::OutputVector{subgraph}, ov::ParameterVector{input0, input1});
     })();
     ov::Core core;
@@ -84,7 +84,7 @@ TEST_F(SubgraphSnippetSerializationTest, smoke_SerializeSubgraphWithScalarConst)
         auto internal_add = std::make_shared<Add>(internal_input, internal_constant);
         auto subgraph_body =
             std::make_shared<ov::Model>(ov::OutputVector{internal_add}, ov::ParameterVector{internal_input});
-        auto subgraph = std::make_shared<ov::snippets::op::Subgraph>(ov::NodeVector{add}, subgraph_body.get()->clone());
+        auto subgraph = std::make_shared<ov::snippets::op::Subgraph>(ov::OutputVector{add}, subgraph_body.get()->clone());
         return std::make_shared<ov::Model>(ov::OutputVector{subgraph}, ov::ParameterVector{input});
     })();
     ov::Core core;

src/plugins/intel_cpu/tests/functional/shared_tests_instances/snippets/mha_transposed_b.cpp

@@ -50,6 +50,23 @@ INSTANTIATE_TEST_SUITE_P(
                        ::testing::Values(CPUTestUtils::empty_plugin_config)),
     MHA::getTestCaseName);
 
+std::vector<std::vector<ov::test::InputShape>> shared_kv_shapes = {{
+    {PartialShape{-1, -1, -1, -1}, {{1, 3, 64, 128}}},
+    {PartialShape{-1, -1, -1, -1}, {{1, 3, 64, 128}}},
+}};
+
+INSTANTIATE_TEST_SUITE_P(smoke_Snippets_MHA_SharedKV,
+                         MHASharedKV,
+                         ::testing::Combine(::testing::ValuesIn(shared_kv_shapes),
+                                            ::testing::ValuesIn(precision_f32(2)),
+                                            ::testing::Values(ov::element::f32),
+                                            ::testing::Values(false),
+                                            ::testing::Values(MHA::default_thread_count),
+                                            ::testing::Values(expected_num_nodes),
+                                            ::testing::Values(1),
+                                            ::testing::Values(ov::test::utils::DEVICE_CPU),
+                                            ::testing::Values(CPUTestUtils::empty_plugin_config)),
+                         MHA::getTestCaseName);
 }  // namespace
 }  // namespace snippets
 }  // namespace test

src/tests/functional/plugin/shared/include/snippets/mha.hpp

@@ -134,6 +134,11 @@ class MHAWithDynamicMul : public testing::WithParamInterface<ov::test::snippets:
     void init_params(std::vector<InputShape>& input_shapes, ov::element::Type& prc, ov::AnyMap& additional_config) override;
 };
 
+class MHASharedKV : public MHA {
+protected:
+    std::shared_ptr<SnippetsFunctionBase> get_subgraph() const override;
+};
+
 }  // namespace snippets
 }  // namespace test
 }  // namespace ov

src/tests/functional/plugin/shared/src/snippets/mha.cpp

@@ -246,6 +246,10 @@ std::shared_ptr<SnippetsFunctionBase> MHAWithDynamicMul::get_subgraph() const {
     return std::make_shared<ov::test::snippets::MHAWithDynamicMulFunction>(inputDynamicShapes, m_input_types);
 }
 
+std::shared_ptr<SnippetsFunctionBase> MHASharedKV::get_subgraph() const {
+    return std::make_shared<ov::test::snippets::MHASharedKVFunction>(inputDynamicShapes, m_input_types);
+}
+
 TEST_P(MHA, CompareWithRefImpl) {
     SKIP_IF_CURRENT_TEST_IS_DISABLED()
     run();
@@ -331,6 +335,12 @@ TEST_P(MHAWithDynamicMul, CompareWithRefImpl) {
     validateNumSubgraphs();
 }
 
+TEST_P(MHASharedKV, CompareWithRefImpl) {
+    SKIP_IF_CURRENT_TEST_IS_DISABLED()
+    run();
+    validateNumSubgraphs();
+}
+
 }  // namespace snippets
 }  // namespace test
 }  // namespace ov

src/tests/ov_helpers/ov_snippets_models/include/subgraph_mha.hpp

@@ -494,6 +494,32 @@ class MHARankUpgradeToReductionFunction : public SnippetsFunctionBase {
     std::shared_ptr<ov::Model> initReference() const override;
 };
 
+/* Graph:
+ * input0   input1
+ *    \     /  \
+ *    MatMul0   \
+ *       |       \
+ *    Softmax     \
+ *       \       /
+ *       MatMul1
+ * Note: This is a MHA pattern with shared K and V inputs, duplicating one of the python TF tests
+ */
+class MHASharedKVFunction : public SnippetsFunctionBase {
+public:
+    explicit MHASharedKVFunction(const std::vector<PartialShape>& inputShapes,
+                                 const std::vector<ov::element::Type>& precisions)
+        : SnippetsFunctionBase(inputShapes),
+          precisions(precisions) {
+        OPENVINO_ASSERT(input_shapes.size() == 2, "Got invalid number of input shapes");
+        OPENVINO_ASSERT(precisions.size() == 2, "Got invalid number of input precisions");
+    }
+
+protected:
+    std::shared_ptr<ov::Model> initOriginal() const override;
+
+    const std::vector<ov::element::Type> precisions;
+};
+
 }  // namespace snippets
 }  // namespace test
 }  // namespace ov