revert to origin mul_cos code

Signed-off-by: HU Yuan2 <[email protected]>

Author: tiger100256-hu

src/common/transformations/src/transformations/common_optimizations/fuse_rotary_positional_embeddings.cpp

@@ -179,7 +179,8 @@ ov::pass::RoPEFusionGPTNEOX::RoPEFusionGPTNEOX(int rank) {
     // so here we use a WA, only match the path of rotate_hal(x)*sin and check the x*cos path
     // in the callback
     auto x = pattern::any_input(pattern::rank_equals(rank));
-    auto cos = pattern::any_input(pattern::rank_equals(rank));
+    auto x_or_cos1 = pattern::any_input(pattern::rank_equals(rank));
+    auto x_or_cos2 = pattern::any_input(pattern::rank_equals(rank));
     auto t_sin = pattern::any_input(pattern::rank_equals(rank));
 
     auto varsplit = pattern::wrap_type<v1::VariadicSplit>({x, rank - 1, {"half_ndims", "?"}});
@@ -192,10 +193,7 @@ ov::pass::RoPEFusionGPTNEOX::RoPEFusionGPTNEOX(int rank) {
     auto x1 = NewGenSlice(x, 0, "half_ndims", 1, rank - 1);
     auto x_rotate_half = pattern::wrap_type<v0::Concat>({x2neg, x1 | varsplit->output(0)}, {{"axis", -1}});
 
-    auto mul_cos1 = pattern::wrap_type<v1::Multiply>({x, cos}, {{"auto_broadcast", "numpy"}});
-    auto mul_cos2 = pattern::wrap_type<v1::Multiply>({cos, x}, {{"auto_broadcast", "numpy"}});
-    auto mul_cos = mul_cos1 | mul_cos2;
-
+    auto mul_cos = pattern::wrap_type<v1::Multiply>({x_or_cos1, x_or_cos2}, {{"auto_broadcast", "numpy"}});
     auto mul_sin = pattern::wrap_type<v1::Multiply>({x_rotate_half, t_sin}, {{"auto_broadcast", "numpy"}});
 
     auto result = pattern::wrap_type<v1::Add>({mul_cos, mul_sin}, {{"auto_broadcast", "numpy"}});
@@ -204,6 +202,17 @@ ov::pass::RoPEFusionGPTNEOX::RoPEFusionGPTNEOX(int rank) {
         const auto& pattern_map = m.get_pattern_value_map();
         auto root = m.get_match_root();
 
+        // check mul(x, cos) exists
+        Output<Node> v_cos;
+        if (pattern_map.at(x_or_cos1) == pattern_map.at(x)) {
+            v_cos = pattern_map.at(x_or_cos2);
+        } else if (pattern_map.at(x_or_cos2) == pattern_map.at(x)) {
+            v_cos = pattern_map.at(x_or_cos1);
+        } else {
+            // not a RoPE
+            return false;
+        }
+
         auto symbols = m.get_symbols();
         auto half_ndims = symbols["half_ndims"];
         if (!half_ndims.is_integer()) {
@@ -218,7 +227,7 @@ ov::pass::RoPEFusionGPTNEOX::RoPEFusionGPTNEOX(int rank) {
         config.rotary_ndims = 2ul * static_cast<size_t>(half_ndims.i());
 
         new_args.push_back(pattern_map.at(x));
-        new_args.push_back(pattern_map.at(cos));
+        new_args.push_back(v_cos);
         new_args.push_back(pattern_map.at(t_sin));
         auto old_node = root;
         auto new_node = std::make_shared<internal::RoPE>(new_args, config);