[Core][CPU]markup rope's sin/cos generation with f32 (openvinotoolkit#25662)

### Details:
- *Sin/Cos table generation must run in f32 otherwise it has accuracy
issue*
 -  Reference : huggingface/transformers#29285

### Tickets:
 - *CVS-146672*

Author: zhangYiIntel

src/common/transformations/include/transformations/common_optimizations/mark_rope_input_to_keep_in_mixed_precision.hpp

@@ -0,0 +1,38 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/graph_rewrite.hpp"
+#include "transformations_visibility.hpp"
+
+namespace ov {
+namespace pass {
+
+/**
+ * @ingroup ov_transformation_common_api
+ * @brief This transformation markups the 2nd/3rd inputs of Rope with FP32 to mantian accuracy.
+ * +-------+    +-------+    +-------+
+ * |intput1|    |input2 |    |input3 |
+ * |(orig) |    |(fp32) |    |(fp32) |
+ * +---|---+    +---|---+    +---|---+
+ *     |            |            |
+ *     |            |            |
+ *  +--+------------|------------+--+
+ *  |                               |
+ *  |             ROPE              |
+ *  +-------------------------------+
+ */
+
+class TRANSFORMATIONS_API MarkRopeInputsToKeepInMixedPrecision : public ov::pass::MatcherPass {
+public:
+    OPENVINO_RTTI("MarkRopeInputsToKeepInMixedPrecision", "0");
+    MarkRopeInputsToKeepInMixedPrecision();
+
+private:
+    std::unordered_set<ov::Node*> visited;
+};
+
+}  // namespace pass
+}  // namespace ov

src/common/transformations/include/transformations/utils/utils.hpp

@@ -215,6 +215,19 @@ TRANSFORMATIONS_API std::shared_ptr<Node> clone_try_fold(const std::shared_ptr<N
 TRANSFORMATIONS_API bool shapes_equal_except_dynamic_expected_batch(const PartialShape& expected,
                                                                     const PartialShape& actual);
 
+/**
+ * \brief Traverses path starting from `node`, and calls "func" for each ov::Node.
+ *
+ * \param node  The node from which path is started.
+ * \param visited  Set of nodes which were visited.
+ * \param func  The function which is called for each visited node.
+ * \param skip_node_predicate  predicte to skip nodes.
+ */
+TRANSFORMATIONS_API void visit_path(ov::Node* node,
+                                    std::unordered_set<ov::Node*>& visited,
+                                    std::function<void(ov::Node*)> func,
+                                    std::function<bool(ov::Node*)> skip_node_predicate);
+
 /**
  * \brief Traverses a shapeOf subgraph starting from the node and not including the ShapeOf nodes,
  * and calls "func" for each ov::Node.

src/common/transformations/src/transformations/common_optimizations/mark_rope_input_to_keep_in_mixed_precision.cpp

@@ -0,0 +1,49 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "transformations/common_optimizations/mark_rope_input_to_keep_in_mixed_precision.hpp"
+
+#include <unordered_set>
+
+#include "itt.hpp"
+#include "openvino/core/rt_info.hpp"
+#include "openvino/pass/constant_folding.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+#include "ov_ops/rotary_positional_embeddings.hpp"
+#include "transformations/rt_info/disable_fp16_compression.hpp"
+#include "transformations/utils/gen_pattern.hpp"
+#include "transformations/utils/utils.hpp"
+
+ov::pass::MarkRopeInputsToKeepInMixedPrecision::MarkRopeInputsToKeepInMixedPrecision() {
+    MATCHER_SCOPE(MarkRopeInputsToKeepInMixedPrecision);
+    using namespace ov::pass::pattern;
+    using namespace ov::gen_pattern;
+    auto cos_tab = any_input();
+    auto sin_tab = any_input();
+    auto rope = makePattern<ov::op::internal::RoPE>({any_input(), cos_tab, sin_tab});
+
+    ov::matcher_pass_callback callback = [=](ov::pass::pattern::Matcher& m) {
+        const auto& pattern_map = m.get_pattern_value_map();
+        auto cos_input_node = pattern_map.at(cos_tab).get_node();
+        auto sin_input_node = pattern_map.at(sin_tab).get_node();
+        // mark the node as disable_fp16_compression
+        auto visit_func = [](ov::Node* node) {
+            ov::disable_fp16_compression(node->shared_from_this());
+        };
+        // skip constant and parameter node
+        auto skip_node_predicate = [](ov::Node* node) -> bool {
+            return ov::is_type<ov::op::v0::Constant>(node) || ov::is_type<ov::op::v0::Parameter>(node);
+        };
+        if (!visited.count(cos_input_node)) {
+            ov::op::util::visit_path(cos_input_node, visited, visit_func, skip_node_predicate);
+        }
+        if (!visited.count(sin_input_node)) {
+            ov::op::util::visit_path(sin_input_node, visited, visit_func, skip_node_predicate);
+        }
+        return false;
+    };
+
+    auto m = std::make_shared<ov::pass::pattern::Matcher>(rope, matcher_name);
+    this->register_matcher(m, callback);
+}

src/common/transformations/src/transformations/utils/utils.cpp

@@ -30,11 +30,10 @@ namespace ov {
 namespace op {
 namespace util {
 
-namespace {
-void visit_path_impl(ov::Node* node,
-                     std::unordered_set<ov::Node*>& visited,
-                     std::function<void(ov::Node*)> func,
-                     std::function<bool(ov::Node*)> skip_node_predicate) {
+void visit_path(ov::Node* node,
+                std::unordered_set<ov::Node*>& visited,
+                std::function<void(ov::Node*)> func,
+                std::function<bool(ov::Node*)> skip_node_predicate) {
     if (!node)
         return;
     visited.insert(node);
@@ -56,7 +55,6 @@ void visit_path_impl(ov::Node* node,
         }
     }
 }
-}  // namespace
 
 bool get_single_value(const std::shared_ptr<op::v0::Constant>& const_node, float& value, bool check_value_range) {
     switch (const_node->get_element_type()) {
@@ -287,7 +285,7 @@ void visit_shape_path(Node* node, std::unordered_set<ov::Node*>& visited, std::f
     auto is_shapeof = [](ov::Node* node) {
         return ov::is_type<ov::op::v0::ShapeOf>(node) || ov::is_type<ov::op::v3::ShapeOf>(node);
     };
-    visit_path_impl(node, visited, func, is_shapeof);
+    visit_path(node, visited, func, is_shapeof);
 }
 
 void visit_constant_path(ov::Node* node, std::unordered_set<ov::Node*>& visited, std::function<void(ov::Node*)> func) {
@@ -296,7 +294,7 @@ void visit_constant_path(ov::Node* node, std::unordered_set<ov::Node*>& visited,
                         "visit_constant_path is called for non-constant path.");
         return false;
     };
-    visit_path_impl(node, visited, func, check_parameter);
+    visit_path(node, visited, func, check_parameter);
 }
 
 bool is_dequantization_subgraph(const Output<Node>& node) {

src/common/transformations/tests/common_optimizations/mark_rope_input_to_keep_in_mixed_precision_test.cpp

@@ -0,0 +1,75 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+#include "transformations/common_optimizations/mark_rope_input_to_keep_in_mixed_precision.hpp"
+
+#include "common_test_utils/ov_test_utils.hpp"
+#include "openvino/opsets/opset1.hpp"
+#include "openvino/pass/manager.hpp"
+#include "ov_ops/rotary_positional_embeddings.hpp"
+#include "transformations/rt_info/disable_fp16_compression.hpp"
+
+TEST_F(TransformationTestsF, MarkRopeInputsToKeepInMixedPrecisionTest) {
+    /*
+    The 2nd/3rd inputs of ROPE is marked as FP32
+                Param2  Param3
+                  \       /
+                   \     /
+                  Matmul(FP32)
+                      |
+                 Transpose(FP32)
+                      |
+                 Concat(FP32)
+                     /   \
+                    /     \
+    Param1   Cos(FP32)   Sin(FP32)
+      \         |          /
+       \        |         /
+        \       |        /
+               ROPE
+    */
+    {
+        auto input = std::make_shared<ov::opset1::Parameter>(ov::element::f32, ov::Shape{1, 10, 8, 64});
+        auto input_a = std::make_shared<ov::opset1::Parameter>(ov::element::f32, ov::Shape{1, 32, 1});
+        auto input_b = std::make_shared<ov::opset1::Parameter>(ov::element::f32, ov::Shape{1, 1, 10});
+        auto matmul = std::make_shared<ov::opset1::MatMul>(input_a, input_b);
+        auto transpose_order =
+            ov::op::v0::Constant::create(ov::element::i32, ov::Shape{3}, std::vector<int32_t>{0, 2, 1});
+        auto transpose = std::make_shared<ov::opset1::Transpose>(matmul, transpose_order);
+        auto concat = std::make_shared<ov::opset1::Concat>(ov::NodeVector{transpose, transpose}, -1);
+        auto cos = std::make_shared<ov::opset1::Cos>(concat);
+        auto sin = std::make_shared<ov::opset1::Sin>(concat);
+        ov::op::internal::RoPE::Config config;
+        auto rope =
+            std::make_shared<ov::op::internal::RoPE>(ov::OutputVector{input->output(0), cos->output(0), sin->output(0)},
+                                                     config);
+        model = std::make_shared<ov::Model>(rope, ov::ParameterVector{input, input_a, input_b}, "model");
+    }
+
+    manager.register_pass<ov::pass::MarkRopeInputsToKeepInMixedPrecision>();
+
+    {
+        auto input = std::make_shared<ov::opset1::Parameter>(ov::element::f32, ov::Shape{1, 10, 8, 64});
+        auto input_a = std::make_shared<ov::opset1::Parameter>(ov::element::f32, ov::Shape{1, 32, 1});
+        auto input_b = std::make_shared<ov::opset1::Parameter>(ov::element::f32, ov::Shape{1, 1, 10});
+        auto matmul = std::make_shared<ov::opset1::MatMul>(input_a, input_b);
+        auto transpose_order =
+            ov::op::v0::Constant::create(ov::element::i32, ov::Shape{3}, std::vector<int32_t>{0, 2, 1});
+        auto transpose = std::make_shared<ov::opset1::Transpose>(matmul, transpose_order);
+        auto concat = std::make_shared<ov::opset1::Concat>(ov::NodeVector{transpose, transpose}, -1);
+        auto cos = std::make_shared<ov::opset1::Cos>(concat);
+        auto sin = std::make_shared<ov::opset1::Sin>(concat);
+        disable_fp16_compression(matmul);
+        disable_fp16_compression(transpose);
+        disable_fp16_compression(concat);
+        disable_fp16_compression(cos);
+        disable_fp16_compression(sin);
+        ov::op::internal::RoPE::Config config;
+        auto rope =
+            std::make_shared<ov::op::internal::RoPE>(ov::OutputVector{input->output(0), cos->output(0), sin->output(0)},
+                                                     config);
+        model_ref = std::make_shared<ov::Model>(rope, ov::ParameterVector{input, input_a, input_b}, "model_ref");
+    }
+    comparator.enable(FunctionsComparator::CmpValues::CONST_VALUES);
+    comparator.enable(FunctionsComparator::CmpValues::RUNTIME_KEYS);
+}

src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp

@@ -35,6 +35,7 @@
 #include "transformations/common_optimizations/matmul_const_transposes_extraction.hpp"
 #include "transformations/common_optimizations/fuse_rotary_positional_embeddings.hpp"
 #include "transformations/common_optimizations/move_eltwise_up_data_movement.hpp"
+#include "transformations/common_optimizations/mark_rope_input_to_keep_in_mixed_precision.hpp"
 #include "transformations/control_flow/unroll_tensor_iterator.hpp"
 #include "transformations/fp16_compression/mark_decompression_convert_constant_folding.hpp"
 #include "transformations/op_conversions/convert_avgpool_downgrade.hpp"
@@ -853,6 +854,9 @@ void Transformations::PostLpt() {
     }
     CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::transpose_sinking::TSShapeOfForward);
     CPU_REGISTER_PASS_COMMON(postLPTPassManager, StatefulSDPAFusion);
+    // markup Rope Input when BF16/F16 inference.
+    if (one_of(inferencePrecision, ov::element::bf16, ov::element::f16))
+        CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::MarkRopeInputsToKeepInMixedPrecision);
 
     // Should be before Snippets pipeline because Ngram pattern contains eltwise nodes that can be tokenized by Snippets.
     auto symbolic_pipeline = CPU_REGISTER_PASS_COMMON(postLPTPassManager, ov::pass::SymbolicOptimizations, false);