From b71be5e8de85fb84e402c812283ef0a7bce90463 Mon Sep 17 00:00:00 2001
From: xuchen-intel <chen.xu@intel.com>
Date: Fri, 7 Jun 2024 07:39:42 +0200
Subject: [PATCH] Add Mish for Snippets tokenization

---
 src/common/snippets/src/pass/collapse_subgraph.cpp         | 1 +
 .../src/emitters/snippets/aarch64/cpu_generator.cpp        | 1 +
 .../src/transformations/transformation_pipeline.cpp        | 7 ++++++-
 .../custom/single_layer_tests/classes/activation.cpp       | 3 +++
 4 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/common/snippets/src/pass/collapse_subgraph.cpp b/src/common/snippets/src/pass/collapse_subgraph.cpp
index b63b9612f81..9afdb340afb 100644
--- a/src/common/snippets/src/pass/collapse_subgraph.cpp
+++ b/src/common/snippets/src/pass/collapse_subgraph.cpp
@@ -124,6 +124,7 @@ auto is_supported_op(const std::shared_ptr<const Node> &n) -> bool {
             || ov::is_type<ov::op::v0::Erf>(n)
             || ov::is_type<ov::op::v0::Exp>(n)
             || ov::is_type<ov::op::v1::LogicalNot>(n)
+            || ov::is_type<ov::op::v4::Mish>(n)
             || ov::is_type<ov::op::v0::Negative>(n)
             || ov::is_type<ov::op::v0::Relu>(n)
             || ov::is_type<ov::op::v5::Round>(n)
diff --git a/src/plugins/intel_cpu/src/emitters/snippets/aarch64/cpu_generator.cpp b/src/plugins/intel_cpu/src/emitters/snippets/aarch64/cpu_generator.cpp
index 8c95a78ec2f..1f2029422ff 100644
--- a/src/plugins/intel_cpu/src/emitters/snippets/aarch64/cpu_generator.cpp
+++ b/src/plugins/intel_cpu/src/emitters/snippets/aarch64/cpu_generator.cpp
@@ -101,6 +101,7 @@ CPUTargetMachine::CPUTargetMachine(dnnl::impl::cpu::aarch64::cpu_isa_t host_isa)
     jitters[ov::op::v0::Exp::get_type_info_static()] = CREATE_CPU_EMITTER(jit_exp_emitter);
     jitters[ov::op::v0::Floor::get_type_info_static()] = CREATE_CPU_EMITTER(jit_floor_emitter);
     jitters[ov::op::v4::HSwish::get_type_info_static()] = CREATE_CPU_EMITTER(jit_hswish_emitter);
+    jitters[ov::op::v4::Mish::get_type_info_static()] = CREATE_CPU_EMITTER(jit_mish_emitter);
     jitters[ov::op::v0::Relu::get_type_info_static()] = CREATE_CPU_EMITTER(jit_relu_emitter);
     jitters[ov::op::v0::Sigmoid::get_type_info_static()] = CREATE_CPU_EMITTER(jit_sigmoid_emitter);
     jitters[ov::intel_cpu::SwishNode::get_type_info_static()] = CREATE_CPU_EMITTER(jit_swish_emitter);
diff --git a/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp b/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp
index b2711438dc8..42cb6cc888e 100644
--- a/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp
+++ b/src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp
@@ -955,6 +955,7 @@ void Transformations::MainSnippets(void) {
                 ov::is_type<ov::op::v4::HSwish>(n) ||
                 ov::is_type<ov::op::v1::Maximum>(n) ||
                 ov::is_type<ov::op::v1::Minimum>(n) ||
+                ov::is_type<ov::op::v4::Mish>(n) ||
                 ov::is_type<ov::op::v1::Mod>(n) ||
                 ov::is_type<ov::op::v1::Multiply>(n) ||
                 ov::is_type<ov::op::v0::Relu>(n) ||
@@ -968,6 +969,10 @@ void Transformations::MainSnippets(void) {
             return ov::is_type<const ov::op::v4::Swish>(n) && n->inputs().size() > 1 &&
                    !ov::is_type<const ov::op::v0::Constant>(n->get_input_node_shared_ptr(1));
         };
+        // CPU Plugin does not support the following ops for x64
+        auto is_unsupported_by_x64 = [](const std::shared_ptr<const ov::Node> &n) {
+            return ov::is_type<const ov::op::v4::Mish>(n);
+        };
         // todo: general tokenization flow is not currently supported for these operations.
         // they can be tokenized only as a part of complex patterns
         auto is_unsupported_by_common_tokenization = [](const std::shared_ptr<const ov::Node> &n) {
@@ -980,7 +985,7 @@ void Transformations::MainSnippets(void) {
                     ov::is_type<const ov::op::v1::ReduceMax>(n) ||
                     ov::is_type<const ov::op::v1::ReduceSum>(n));
         };
-        return !is_unsupported_swish(n) && !is_unsupported_by_common_tokenization(n);
+        return !is_unsupported_swish(n) && !is_unsupported_by_x64(n) && !is_unsupported_by_common_tokenization(n);
 #endif
     };
 
diff --git a/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp b/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp
index 67801d2396a..cce667cc2e1 100644
--- a/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp
+++ b/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp
@@ -241,6 +241,9 @@ const std::map<utils::ActivationTypes, std::vector<std::vector<float>>>& activat
         {Floor,       {{}}},
         {Relu,        {{}}},
         {HSwish,      {{}}},
+#if defined(OPENVINO_ARCH_ARM64)
+        {Mish,        {{}}},
+#endif
         {Sigmoid,     {{}}},
         {Swish,       {{0.1f}}},
         {Tanh,        {{}}},