[Snippets] LIR serialization improvements (openvinotoolkit#21291)

src/common/snippets/include/snippets/lowered/linear_ir.hpp

@@ -125,7 +125,6 @@ class LinearIR {
     iterator find_after(iterator it, const ExpressionPtr& target) const;
 
     void init_emitters(const std::shared_ptr<TargetMachine>& target);
-    void serialize(const std::string& xml, const std::string& bin) const;
 
     class LoopManager;
     using LoopManagerPtr = std::shared_ptr<LoopManager>;

src/common/snippets/include/snippets/lowered/pass/serialize_base.hpp

@@ -0,0 +1,35 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "pass.hpp"
+#include "snippets/lowered/linear_ir.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+/**
+ * @interface SerializeBase
+ * @brief Base class for LinearIR serialization passes
+ * @ingroup snippets
+ */
+class SerializeBase : public Pass {
+public:
+    OPENVINO_RTTI("SerializeBase", "Pass")
+    SerializeBase(const std::string& xml_path);
+
+protected:
+    std::string get_bin_path_from_xml(const std::string& xml_path);
+
+    const std::string m_xml_path;
+    const std::string m_bin_path;
+};
+
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

src/common/snippets/include/snippets/lowered/pass/serialize_control_flow.hpp

@@ -0,0 +1,30 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "serialize_base.hpp"
+#include "snippets/lowered/linear_ir.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+/**
+ * @interface SerializeControlFlow
+ * @brief Serializes control flow graph of LinearIR
+ * @ingroup snippets
+ */
+class SerializeControlFlow : public SerializeBase {
+public:
+    OPENVINO_RTTI("SerializeControlFlow", "Pass", SerializeBase)
+    SerializeControlFlow(const std::string& xml_path) : SerializeBase(xml_path) {}
+    bool run(LinearIR& linear_ir) override;
+};
+
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

src/common/snippets/include/snippets/lowered/pass/serialize_data_flow.hpp

@@ -0,0 +1,32 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "serialize_base.hpp"
+#include "snippets/lowered/linear_ir.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+/**
+ * @interface SerializeDataFlow
+ * @brief Serializes data flow graph of LinearIR
+ * @attention - This pass can not be run on the LinearIR after tail loop insertion.
+ * @attention - Control flow operations (e.g. LoopBegin/LoopEnd) are not serialized
+ * @ingroup snippets
+ */
+class SerializeDataFlow : public SerializeBase {
+public:
+    OPENVINO_RTTI("SerializeDataFlow", "Pass", SerializeBase)
+    SerializeDataFlow(const std::string& xml_path) : SerializeBase(xml_path) {}
+    bool run(LinearIR& linear_ir) override;
+};
+
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

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

@@ -19,17 +19,28 @@ namespace op {
  */
 class SerializationNode : public ov::op::Op {
 public:
-    OPENVINO_OP("SerializationNode", "SnippetsOpset");
-
+    enum SerializationMode { DATA_FLOW, CONTROL_FLOW };
     SerializationNode() = default;
-    SerializationNode(const ov::OutputVector& args, const std::shared_ptr<lowered::Expression>& expr);
+    SerializationNode(const ov::OutputVector& args,
+                      const std::shared_ptr<lowered::Expression>& expr,
+                      SerializationMode mode = SerializationMode::CONTROL_FLOW);
 
     void validate_and_infer_types() override;
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector &new_args) const override;
     bool visit_attributes(AttributeVisitor &visitor) override;
 
+    _OPENVINO_HIDDEN_METHOD static const DiscreteTypeInfo& get_type_info_static() {
+        static ::ov::DiscreteTypeInfo type_info_static{"SerializationNode", "SnippetsOpset"};
+        return type_info_static;
+    }
+
+    const ::ov::DiscreteTypeInfo& get_type_info() const override {
+        return m_expr->get_node()->get_type_info();
+    }
+
 private:
     std::shared_ptr<lowered::Expression> m_expr;
+    SerializationMode m_mode;
 };
 
 } // namespace op

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

@@ -124,7 +124,6 @@ class Subgraph : public ov::op::util::SubGraphOp {
 
     void print() const;
 
-    void serialize() const;
     VectorDims infer_master_shape();
 
     static auto wrap_node_as_subgraph(const std::shared_ptr<ov::Node>& node) -> std::shared_ptr<Subgraph>;

src/common/snippets/src/lowered/linear_ir.cpp

@@ -8,7 +8,6 @@
 
 #include "snippets/lowered/loop_manager.hpp"
 #include "snippets/lowered/expression_factory.hpp"
-#include "snippets/op/serialization_node.hpp"
 
 #include "openvino/core/graph_util.hpp"
 #include "openvino/core/type.hpp"
@@ -86,38 +85,6 @@ ov::NodeVector LinearIR::get_ordered_ops(const std::shared_ptr<ov::Model>& m) {
     return ov::topological_sort(nodes);
 }
 
-void LinearIR::serialize(const std::string& xml, const std::string& bin) const {
-    auto first_node = std::make_shared<ov::op::v0::Parameter>(element::f32, Shape{});
-    first_node->set_friendly_name("Start");
-    first_node->get_rt_info()["execTimeMcs"] = 0;
-    std::shared_ptr<Node> serialization_node = first_node;
-
-    // This map allows to get LoopBegin serialization node by original LoopBegin node
-    // It is used to draw an edge between LoopBegin and LoopEnd serialization nodes
-    std::map<std::shared_ptr<snippets::op::LoopBegin>, std::shared_ptr<Node>> loops_map;
-    for (const auto& expr : m_expressions) {
-        const auto node = expr->get_node();
-        if (auto loop_end = ov::as_type_ptr<snippets::op::LoopEnd>(node)) {
-            OPENVINO_ASSERT(loops_map.count(loop_end->get_loop_begin()),
-                            "Serialization can't find LoopBegin that corresponds to LoopEnd with friendly name ",
-                            loop_end->get_friendly_name());
-            auto loop_begin_serialization_node = loops_map.at(loop_end->get_loop_begin());
-            serialization_node = std::make_shared<op::SerializationNode>(ov::OutputVector{serialization_node, loop_begin_serialization_node}, expr);
-        } else {
-            serialization_node = std::make_shared<op::SerializationNode>(ov::OutputVector{serialization_node}, expr);
-            if (auto loop_begin = ov::as_type_ptr<snippets::op::LoopBegin>(node)) {
-                loops_map[loop_begin] = serialization_node;
-            }
-        }
-    }
-    auto last_node = std::make_shared<ov::op::v0::Result>(serialization_node);
-    last_node->set_friendly_name("End");
-    const auto tmp_model = std::make_shared<ov::Model>(ResultVector {last_node},
-                                                       ParameterVector {first_node},
-                                                       "Lowered_IR_Serialization");
-    ov::pass::Serialize(xml, bin).run_on_model(tmp_model);
-}
-
 LinearIR::container LinearIR::deep_copy_range(LinearIR::container::const_iterator begin,
                                               LinearIR::container::const_iterator end,
                                               ExressionMap& expression_map) {

src/common/snippets/src/lowered/pass/serialize_base.cpp

@@ -0,0 +1,29 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "snippets/lowered/pass/serialize_base.hpp"
+
+#include "snippets/itt.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+SerializeBase::SerializeBase(const std::string& xml_path)
+    : m_xml_path(xml_path),
+      m_bin_path(get_bin_path_from_xml(xml_path)) {}
+
+std::string SerializeBase::get_bin_path_from_xml(const std::string& xml_path) {
+#if defined(__linux__)
+    return "/dev/null";
+#else
+    return "";
+#endif
+}
+
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

src/common/snippets/src/lowered/pass/serialize_control_flow.cpp

@@ -0,0 +1,55 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "snippets/lowered/pass/serialize_control_flow.hpp"
+
+#include "openvino/pass/serialize.hpp"
+#include "snippets/itt.hpp"
+#include "snippets/lowered/linear_ir.hpp"
+#include "snippets/op/serialization_node.hpp"
+#include "snippets/snippets_isa.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+bool SerializeControlFlow::run(LinearIR& linear_ir) {
+    OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::SerializeControlFlow")
+    if (linear_ir.empty())
+        return false;
+
+    auto first_node = std::make_shared<ov::op::v0::Parameter>(element::f32, Shape{});
+    first_node->set_friendly_name("Start");
+    first_node->get_rt_info()["execTimeMcs"] = 0;
+    std::shared_ptr<Node> serialization_node = first_node;
+
+    // This map allows to get LoopBegin serialization node by original LoopBegin node
+    // It is used to draw an edge between LoopBegin and LoopEnd serialization nodes
+    std::map<std::shared_ptr<snippets::op::LoopBegin>, std::shared_ptr<Node>> loops_map;
+    for (const auto& expr : linear_ir) {
+        const auto node = expr->get_node();
+        if (auto loop_end = ov::as_type_ptr<snippets::op::LoopEnd>(node)) {
+            OPENVINO_ASSERT(loops_map.count(loop_end->get_loop_begin()),
+                            "Serialization can't find LoopBegin that corresponds to LoopEnd with friendly name ",
+                            loop_end->get_friendly_name());
+            auto loop_begin_serialization_node = loops_map.at(loop_end->get_loop_begin());
+            serialization_node = std::make_shared<op::SerializationNode>(ov::OutputVector{serialization_node, loop_begin_serialization_node}, expr);
+        } else {
+            serialization_node = std::make_shared<op::SerializationNode>(ov::OutputVector{serialization_node}, expr);
+            if (auto loop_begin = ov::as_type_ptr<snippets::op::LoopBegin>(node)) {
+                loops_map[loop_begin] = serialization_node;
+            }
+        }
+    }
+    auto last_node = std::make_shared<ov::op::v0::Result>(serialization_node);
+    last_node->set_friendly_name("End");
+    const auto model = std::make_shared<ov::Model>(ResultVector{last_node}, ParameterVector{first_node}, "Lowered_IR_Control_Flow");
+    return ov::pass::Serialize(m_xml_path, m_bin_path).run_on_model(model);
+}
+
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

src/common/snippets/src/lowered/pass/serialize_data_flow.cpp

@@ -0,0 +1,57 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "snippets/lowered/pass/serialize_data_flow.hpp"
+
+#include "openvino/pass/serialize.hpp"
+#include "snippets/itt.hpp"
+#include "snippets/lowered/linear_ir.hpp"
+#include "snippets/op/serialization_node.hpp"
+#include "snippets/snippets_isa.hpp"
+
+namespace ov {
+namespace snippets {
+namespace lowered {
+namespace pass {
+
+bool SerializeDataFlow::run(LinearIR& linear_ir) {
+    OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::SerializeDataFlow")
+    if (linear_ir.empty())
+        return false;
+
+    ov::ResultVector results;
+    ov::ParameterVector parameters;
+    std::map<ExpressionPtr, std::shared_ptr<Node>> ops_map;
+    const auto serialization_mode = op::SerializationNode::SerializationMode::DATA_FLOW;
+    for (const auto& expr : linear_ir) {
+        const auto node = expr->get_node();
+        ov::OutputVector inputs(expr->get_input_count());
+        for (size_t i = 0; i < expr->get_input_count(); ++i) {
+            const auto& input_expr = expr->get_input_port_connector(i)->get_source().get_expr();
+            OPENVINO_ASSERT(ops_map.count(input_expr), "input node wasn't found during serialization");
+            inputs[i] = ops_map[input_expr]->output(expr->get_input_port_connector(i)->get_source().get_index());
+        }
+        if (auto ioexpr = std::dynamic_pointer_cast<IOExpression>(expr)) {
+            if (ioexpr->get_type() == IOExpression::io_type::INPUT) {
+                const auto parameter = std::make_shared<ov::op::v0::Parameter>(element::f32, Shape{});
+                ops_map[ioexpr] = parameter;
+                parameters.push_back(parameter);
+            } else {
+                const auto result = std::make_shared<ov::op::v0::Result>(inputs[0]);
+                ops_map[ioexpr] = result;
+                results.push_back(result);
+            }
+        } else {
+            const auto serialization_node = std::make_shared<op::SerializationNode>(inputs, expr, serialization_mode);
+            ops_map[expr] = serialization_node;
+        }
+    }
+    const auto model = std::make_shared<ov::Model>(results, parameters, "Lowered_IR_Data_Flow");
+    return ov::pass::Serialize(m_xml_path, m_bin_path).run_on_model(model);
+}
+
+} // namespace pass
+} // namespace lowered
+} // namespace snippets
+} // namespace ov

src/common/snippets/src/op/serialization_node.cpp

@@ -9,26 +9,34 @@ namespace ov {
 namespace snippets {
 namespace op {
 
-SerializationNode::SerializationNode(const ov::OutputVector& args, const std::shared_ptr<lowered::Expression>& expr)
-    : Op(args), m_expr(expr) {
-    if (!m_expr || !m_expr->get_node())
-        OPENVINO_THROW("SerializationNode requires a valid expression with non-null node pointer");
-    const auto &node = expr->get_node();
+SerializationNode::SerializationNode(const ov::OutputVector& args,
+                                     const std::shared_ptr<lowered::Expression>& expr,
+                                     SerializationMode mode)
+    : Op(args),
+      m_expr(expr),
+      m_mode(mode) {
+    OPENVINO_ASSERT(m_expr && m_expr->get_node(), "SerializationNode requires a valid expression with non-null node pointer");
+    const auto& node = expr->get_node();
+    set_friendly_name(node->get_friendly_name());
     std::string type = node->get_type_name();
-    std::string name = node->get_friendly_name();
-    // If node is a parameter, show another type name, so the node will be displayed correctly
     get_rt_info()["layerType"] = type == "Parameter" ? "ParameterLowered" : type;
-    set_friendly_name(name);
     constructor_validate_and_infer_types();
 }
 
 void SerializationNode::validate_and_infer_types() {
-    set_output_type(0, element::f32, ov::PartialShape{});
+    // If SerializationNode is used for control flow serialization, it always has one output
+    // (since it represents a linear execution order)
+    if (m_mode == SerializationMode::CONTROL_FLOW) {
+        set_output_type(0, element::f32, {});
+    } else if (m_mode == SerializationMode::DATA_FLOW) {
+        for (size_t i = 0; i < m_expr->get_output_count(); ++i)
+            set_output_type(i, element::f32, {});
+    }
 }
 
 std::shared_ptr<Node> SerializationNode::clone_with_new_inputs(const OutputVector &new_args) const {
     check_new_args_count(this, new_args);
-    return std::make_shared<SerializationNode>(new_args, m_expr);
+    return std::make_shared<SerializationNode>(new_args, m_expr, m_mode);
 }
 
 bool SerializationNode::visit_attributes(AttributeVisitor &visitor) {