[Snippets] Implement runtime loop info serialization in netron graphs

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

@@ -32,11 +32,21 @@ class LinearIRBuilder {
      * @param expression_map expression map
      * @return clone of `linear_ir`
      */
-    std::shared_ptr<LinearIR> clone(const std::shared_ptr<LinearIR>& linear_ir,  ExpressionMap& expression_map) const;
+    inline std::shared_ptr<LinearIR> clone(const std::shared_ptr<LinearIR>& linear_ir,  ExpressionMap& expression_map) const {
+        auto result = std::make_shared<LinearIR>();
+        clone(linear_ir.get(), result.get(), expression_map);
+        return result;
+    }
     inline std::shared_ptr<LinearIR> clone(const std::shared_ptr<LinearIR>& linear_ir) const {
         ExpressionMap expression_map;
         return clone(linear_ir, expression_map);
     }
+    inline LinearIR clone(const LinearIR& linear_ir) const {
+        LinearIR result;
+        ExpressionMap expression_map;
+        clone(&linear_ir, &result, expression_map);
+        return result;
+    }
     /**
      * @brief Make a copy of LinearIR range by rules described in `m_config`
      * @param begin begin iterator of the target range of LinearIR
@@ -48,6 +58,7 @@ class LinearIRBuilder {
                                     ExpressionMap& expression_map) const;
 
 private:
+    void clone(const LinearIR* src, LinearIR* dst,  ExpressionMap& expression_map) const;
     Config m_config = {};
 };
 

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

@@ -15,13 +15,24 @@ namespace pass {
 /**
  * @interface SerializeControlFlow
  * @brief Serializes control flow graph of LinearIR
+ * @param update_dynamic_ops - update state of dynamic operations (e.g. LoopEnd) based on runtime info contained in LIR.
+ * Note that this flag triggers copying of the passed LIR internally to preserve the original state of IR.
  * @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;
+    SerializeControlFlow(const std::string& xml_path, bool update_dynamic_ops = false) :
+        SerializeBase(xml_path), m_update_dynamic_ops{update_dynamic_ops} {}
+
+    bool run(LinearIR& linear_ir) override {
+        return run(const_cast<const LinearIR&>(linear_ir));
+    }
+    // We need a const method to run from functions that can't change LIR
+    bool run(const LinearIR& linear_ir);
+
+private:
+    const bool m_update_dynamic_ops = false;
 };
 
 } // namespace pass

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

@@ -23,8 +23,15 @@ class SerializationNode : public ov::op::Op {
     SerializationNode() = default;
     SerializationNode(const ov::OutputVector& args,
                       const std::shared_ptr<lowered::Expression>& expr,
+                      const std::shared_ptr<ov::Node>& node,
                       SerializationMode mode = SerializationMode::CONTROL_FLOW);
 
+    SerializationNode(const ov::OutputVector& args,
+                      const std::shared_ptr<lowered::Expression>& expr,
+                      SerializationMode mode = SerializationMode::CONTROL_FLOW) :
+                        SerializationNode(args, expr, expr->get_node(), mode) {
+    }
+
     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;
@@ -40,6 +47,7 @@ class SerializationNode : public ov::op::Op {
 
 private:
     std::shared_ptr<lowered::Expression> m_expr;
+    std::shared_ptr<ov::Node> m_node;
     SerializationMode m_mode;
 };
 

src/common/snippets/src/lowered/linear_ir_builder.cpp

@@ -65,21 +65,21 @@ std::vector<std::shared_ptr<ov::Node>> clone_nodes(const std::vector<std::shared
 }
 }  // namespace
 
-std::shared_ptr<LinearIR> LinearIRBuilder::clone(const std::shared_ptr<LinearIR>& linear_ir, ExpressionMap& expression_map) const {
-    auto cloned = std::make_shared<LinearIR>();
-    cloned->m_config = linear_ir->m_config;
+void LinearIRBuilder::clone(const LinearIR* src, LinearIR* dst,  ExpressionMap& expression_map) const {
+    OPENVINO_ASSERT(src && dst, "Invalid pointers were provided for LinearIRBuilder::clone");
+    dst->m_config = src->m_config;
 
-    cloned->m_expressions = clone_range(linear_ir->m_expressions.cbegin(), linear_ir->m_expressions.cend(), expression_map);
-    for (const auto& expr : cloned->m_expressions) {
-        cloned->register_expression(expr, true);
+    dst->m_expressions = clone_range(src->m_expressions.cbegin(), src->m_expressions.cend(), expression_map);
+    for (const auto& expr : dst->m_expressions) {
+        dst->register_expression(expr, true);
     }
 
-    cloned->m_loop_manager = linear_ir->m_loop_manager->clone_with_new_expr(expression_map);
+    dst->m_loop_manager = src->m_loop_manager->clone_with_new_expr(expression_map);
     // It's Ok to share shapeInfer factory ptr, since the factory doesn't depend on LIR in any way
-    cloned->m_shape_infer_factory = linear_ir->m_shape_infer_factory;
-    cloned->m_shape_infer = std::make_shared<LinearIR::LIRShapeInfer>(cloned->m_expressions, cloned->m_parameter_expressions, cloned->m_result_expressions);
-    cloned->m_is_dynamic = linear_ir->m_is_dynamic;
-    return cloned;
+    dst->m_shape_infer_factory = src->m_shape_infer_factory;
+    dst->m_shape_infer = std::make_shared<LinearIR::LIRShapeInfer>(dst->m_expressions, dst->m_parameter_expressions,
+                                                                   dst->m_result_expressions);
+    dst->m_is_dynamic = src->m_is_dynamic;
 }
 
 LinearIR::container LinearIRBuilder::clone_range(LinearIR::container::const_iterator begin, LinearIR::container::const_iterator end,

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

@@ -9,16 +9,22 @@
 #include "snippets/lowered/linear_ir.hpp"
 #include "snippets/op/serialization_node.hpp"
 #include "snippets/snippets_isa.hpp"
+#include "snippets/lowered/loop_manager.hpp"
+#include "snippets/lowered/linear_ir_builder.hpp"
 
 namespace ov {
 namespace snippets {
 namespace lowered {
 namespace pass {
 
-bool SerializeControlFlow::run(LinearIR& linear_ir) {
+bool SerializeControlFlow::run(const LinearIR& original_linear_ir) {
     OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::SerializeControlFlow")
-    if (linear_ir.empty())
+    if (original_linear_ir.empty())
         return false;
+    const auto& linear_ir = m_update_dynamic_ops ? LinearIRBuilder().clone(original_linear_ir) : original_linear_ir;
+
+    const auto& loop_manager = linear_ir.get_loop_manager();
+    const auto& loop_info_map = loop_manager ? loop_manager->get_map() : std::map<size_t, LoopInfoPtr>{};
 
     auto first_node = std::make_shared<ov::op::v0::Parameter>(element::f32, Shape{});
     first_node->set_friendly_name("Start");
@@ -35,7 +41,24 @@ bool SerializeControlFlow::run(LinearIR& linear_ir) {
                             "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);
+            if (m_update_dynamic_ops) {
+                OPENVINO_ASSERT(loop_info_map.count(loop_end->get_id()), "Failed to find loop id in loop info map");
+                const auto& loop_info = loop_info_map.at(loop_end->get_id());
+                loop_end->set_work_amount(loop_info->get_work_amount());
+                loop_end->set_increment(loop_info->get_increment());
+                loop_end->set_is_incremented(loop_info->get_is_incremented());
+                if (auto unified = ov::as_type_ptr<UnifiedLoopInfo>(loop_info)) {
+                    loop_end->set_ptr_increments(unified->get_ptr_increments());
+                    loop_end->set_finalization_offsets(unified->get_finalization_offsets());
+                } else if (auto expanded = ov::as_type_ptr<ExpandedLoopInfo>(loop_info)) {
+                    loop_end->set_ptr_increments(expanded->get_ptr_increments());
+                    loop_end->set_finalization_offsets(expanded->get_finalization_offsets());
+                }
+            }
+            serialization_node = std::make_shared<op::SerializationNode>(ov::OutputVector{serialization_node, loop_begin_serialization_node},
+                                                                         expr,
+                                                                         loop_end,
+                                                                         op::SerializationNode::SerializationMode::CONTROL_FLOW);
         } 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)) {

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

@@ -11,14 +11,15 @@ namespace op {
 
 SerializationNode::SerializationNode(const ov::OutputVector& args,
                                      const std::shared_ptr<lowered::Expression>& expr,
+                                     const std::shared_ptr<ov::Node>& node,
                                      SerializationMode mode)
     : Op(args),
       m_expr(expr),
+      m_node(node),
       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();
+    OPENVINO_ASSERT(m_expr && m_node, "SerializationNode requires a valid expression with non-null node pointer");
+    set_friendly_name(m_node->get_friendly_name());
+    std::string type = m_node->get_type_name();
     get_rt_info()["layerType"] = type == "Parameter" ? "ParameterLowered" : type;
     constructor_validate_and_infer_types();
 }
@@ -114,7 +115,7 @@ bool SerializationNode::visit_attributes(AttributeVisitor &visitor) {
 
     auto loop_ids = m_expr->get_loop_ids();
     visitor.on_attribute("loop_ids", loop_ids);
-    m_expr->get_node()->visit_attributes(visitor);
+    m_node->visit_attributes(visitor);
     return true;
 }