From cd1eac8da7b1d9676ab6ccf52ff746921ec2664c Mon Sep 17 00:00:00 2001
From: Ilya Lavrenov <ilya.lavrenov@intel.com>
Date: Wed, 14 Feb 2024 18:37:32 +0400
Subject: [PATCH] Fixed deserialization with type relaxed operations (#22824)

### Details:
- `TypeRelaxed` operations to have dedicated opset during serialize /
deserialize to avoid mixing with official OpenVINO operations.

### Issues:
- CVS-130765
---
 src/core/src/pass/serialize.cpp         |  10 ++
 src/plugins/intel_cpu/src/extension.cpp | 117 +++++++++++++++---------
 2 files changed, 83 insertions(+), 44 deletions(-)

diff --git a/src/core/src/pass/serialize.cpp b/src/core/src/pass/serialize.cpp
index 4d68132f4e6051..1b45730c2ed7db 100644
--- a/src/core/src/pass/serialize.cpp
+++ b/src/core/src/pass/serialize.cpp
@@ -709,6 +709,16 @@ const std::vector<Edge> create_edge_mapping(const std::unordered_map<ov::Node*,
 
 std::string get_opset_name(const ov::Node* n) {
     OPENVINO_ASSERT(n != nullptr);
+
+    // TODO: remove it one day: try to find opset name from RT info
+    // It's a dirty hack to TypeRelaxed and similar template internal operations
+    auto opset_it = n->get_rt_info().find("opset");
+    if (opset_it != n->get_rt_info().end()) {
+        if (opset_it->second.is<std::string>()) {
+            return opset_it->second.as<std::string>();
+        }
+    }
+
     return n->get_type_info().version_id == nullptr ? "experimental" : n->get_type_info().version_id;
 }
 
diff --git a/src/plugins/intel_cpu/src/extension.cpp b/src/plugins/intel_cpu/src/extension.cpp
index 41d91bfc382681..bf6df382eedd6d 100644
--- a/src/plugins/intel_cpu/src/extension.cpp
+++ b/src/plugins/intel_cpu/src/extension.cpp
@@ -26,8 +26,37 @@
 #include "transformations/snippets/x64/op/perf_count_rdtsc.hpp"
 #include "transformations/snippets/x64/op/store_convert.hpp"
 
+namespace {
+
+template <typename Op>
+class TypeRelaxedExtension : public ov::OpExtension<ov::op::TypeRelaxed<Op>> {
+public:
+    TypeRelaxedExtension()
+        : m_ext_type(Op::get_type_info_static().name, "type_relaxed_opset") {}
+    ~TypeRelaxedExtension() override = default;
+
+    const ov::DiscreteTypeInfo& get_type_info() const override {
+        return m_ext_type;
+    }
+
+    ov::OutputVector create(const ov::OutputVector& inputs, ov::AttributeVisitor& visitor) const override {
+        return ov::OpExtension<ov::op::TypeRelaxed<Op>>::create(inputs, visitor);
+    }
+
+    std::vector<ov::Extension::Ptr> get_attached_extensions() const override {
+        return {};
+    }
+
+private:
+    ov::DiscreteTypeInfo m_ext_type;
+};
+
+}  // namespace
+
 #define OP_EXTENSION(NAME) std::make_shared<ov::OpExtension<NAME>>(),
 
+#define TYPE_RELAXED_OP_EXTENSION(NAME) std::make_shared<TypeRelaxedExtension<NAME>>(),
+
 #if defined(OPENVINO_ARCH_X86_64)
 #    define OP_EXTENSION_X64(NAME) OP_EXTENSION(NAME)
 #else
@@ -55,50 +84,50 @@
     OP_EXTENSION_X64(ov::intel_cpu::BrgemmCPU)                              \
     OP_EXTENSION_X64(ov::intel_cpu::BrgemmCopyB)
 
-#define TYPE_RELAXED_EXTENSIONS                                                 \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Add>)                          \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::AvgPool>)                      \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::Clamp>)                        \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::Concat>)                       \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Convolution>)                  \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::ConvolutionBackpropData>)      \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::DepthToSpace>)                 \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Equal>)                        \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::FakeQuantize>)                 \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Greater>)                      \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::GreaterEqual>)                 \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::GroupConvolution>)             \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::GroupConvolutionBackpropData>) \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::Interpolate>)                  \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v4::Interpolate>)                  \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Less>)                         \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::LessEqual>)                    \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::LogicalAnd>)                   \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::LogicalNot>)                   \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::LogicalOr>)                    \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::LogicalXor>)                   \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::MatMul>)                       \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::MaxPool>)                      \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Multiply>)                     \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::NormalizeL2>)                  \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::NotEqual>)                     \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::PRelu>)                        \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::Relu>)                         \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::ReduceMax>)                    \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::ReduceLogicalAnd>)             \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::ReduceLogicalOr>)              \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::ReduceMean>)                   \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::ReduceMin>)                    \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::ReduceSum>)                    \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Reshape>)                      \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Select>)                       \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::ShapeOf>)                      \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::ShuffleChannels>)              \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::Squeeze>)                      \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v1::Subtract>)                     \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::Unsqueeze>)                    \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v0::MVN>)                          \
-    OP_EXTENSION(ov::op::TypeRelaxed<ov::op::v6::MVN>)
+#define TYPE_RELAXED_EXTENSIONS                                         \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Add)                          \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::AvgPool)                      \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::Clamp)                        \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::Concat)                       \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Convolution)                  \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::ConvolutionBackpropData)      \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::DepthToSpace)                 \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Equal)                        \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::FakeQuantize)                 \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Greater)                      \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::GreaterEqual)                 \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::GroupConvolution)             \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::GroupConvolutionBackpropData) \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::Interpolate)                  \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v4::Interpolate)                  \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Less)                         \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::LessEqual)                    \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::LogicalAnd)                   \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::LogicalNot)                   \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::LogicalOr)                    \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::LogicalXor)                   \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::MatMul)                       \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::MaxPool)                      \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Multiply)                     \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::NormalizeL2)                  \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::NotEqual)                     \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::PRelu)                        \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::Relu)                         \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::ReduceMax)                    \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::ReduceLogicalAnd)             \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::ReduceLogicalOr)              \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::ReduceMean)                   \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::ReduceMin)                    \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::ReduceSum)                    \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Reshape)                      \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Select)                       \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::ShapeOf)                      \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::ShuffleChannels)              \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::Squeeze)                      \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v1::Subtract)                     \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::Unsqueeze)                    \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v0::MVN)                          \
+    TYPE_RELAXED_OP_EXTENSION(ov::op::v6::MVN)
 
 #ifdef SNIPPETS_DEBUG_CAPS
 #    define SNIPPETS_DEBUG_CAPS_EXTENSIONS                   \