[Snippets] Applied Ivan comments

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

@@ -12,6 +12,9 @@ namespace ov {
 namespace snippets {
 namespace lowered {
 
+class LoopInfo;
+using LoopInfoMap = std::map<LoopInfo*, std::shared_ptr<LoopInfo>>;
+
 /**
  * @interface LoopInfo
  * @brief The base class that contains the common information about a Loop in Linear Intermediate Representation (Linear IR):
@@ -27,6 +30,14 @@ class LoopInfo {
     LoopInfo(size_t work_amount, size_t increment, const std::vector<ExpressionPort>& entries, const std::vector<ExpressionPort>& exits);
     virtual ~LoopInfo() = default;
 
+    /**
+     * @brief Clone LoopInfo with new Expressions
+     * @param expr_map map of new and old expressions
+     * @param loop_map map of new and old LoopInfo.
+     * @return the copy
+     */
+    virtual std::shared_ptr<LoopInfo> clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoMap& loop_map) const = 0;
+
     /**
      * @brief Check if some parameters of Loop are dynamic (undefined)
      * @return True if some parameters of Loop are unknown, False if all parameters are static
@@ -214,11 +225,12 @@ class UnifiedLoopInfo : public LoopInfo {
                     const SpecificIterationHandlers& handlers = SpecificIterationHandlers());
 
     /**
-     * @brief Clone LoopInfo with new expressions
+     * @brief Clone LoopInfo with new Expressions
      * @param expr_map map of new and old expressions
+     * @param loop_map map of new and old LoopInfo.
      * @return the copy
      */
-    std::shared_ptr<LoopInfo> clone_with_new_expr(const ExpressionMap& expr_map) const;
+    std::shared_ptr<LoopInfo> clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoMap& loop_map) const override;
 
     /**
      * @brief Check if some parameters of Loop are dynamic (undefined)
@@ -371,12 +383,14 @@ class InnerSplittedUnifiedLoopInfo : public UnifiedLoopInfo {
                                  const SpecificIterationHandlers& handlers, LoopInfoPtr outer_splitted_loop_info);
 
     /**
-     * @brief Clone LoopInfo with new expressions
+     * @brief Clone LoopInfo with new Expressions
      * @param expr_map map of new and old expressions
-     * @param new_outer_splitted_loop_info new outer splitted loop info
+     * @param loop_map map of new and old LoopInfo.
+     *        If `loop_map` contains cloned outer splitted loop -info, we take it from there.
+     *        Otherwise we manually clone it and add to this map.
      * @return the copy
      */
-    std::shared_ptr<LoopInfo> clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoPtr new_outer_splitted_loop_info) const;
+    std::shared_ptr<LoopInfo> clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoMap& loop_map) const override;
 
     /**
      * @brief Returns work amount of the Loop.
@@ -420,12 +434,14 @@ class ExpandedLoopInfo : public LoopInfo {
                      std::vector<int64_t> ptr_increments, std::vector<int64_t> final_offsets, std::vector<int64_t> data_sizes,
                      SpecificLoopIterType type, UnifiedLoopInfoPtr unified_loop_info, bool evaluate_once = false);
     /**
-     * @brief Clone LoopInfo with new expressions
+     * @brief Clone LoopInfo with new Expressions
      * @param expr_map map of new and old expressions
-     * @param new_unified_loop_info new unified loop info
+     * @param loop_map map of new and old LoopInfo.
+     *        If `loop_map` contains cloned unified loop -info, we take it from there.
+     *        Otherwise we manually clone it and add to this map.
      * @return the copy
      */
-    std::shared_ptr<LoopInfo> clone_with_new_expr(const ExpressionMap& expr_map, UnifiedLoopInfoPtr new_unified_loop_info) const;
+    std::shared_ptr<LoopInfo> clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoMap& loop_map) const override;
 
     /**
      * @brief Check if some parameters of Loop are dynamic (undefined)

src/common/snippets/src/lowered/loop_info.cpp

@@ -200,7 +200,7 @@ UnifiedLoopInfo::UnifiedLoopInfo(size_t work_amount, size_t increment,
     sort_ports();
 }
 
-std::shared_ptr<LoopInfo> UnifiedLoopInfo::clone_with_new_expr(const ExpressionMap& expr_map) const {
+std::shared_ptr<LoopInfo> UnifiedLoopInfo::clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoMap& loop_map) const {
     const auto& new_input_ports = clone_loop_ports(expr_map, m_input_ports);
     const auto& new_output_ports = clone_loop_ports(expr_map, m_output_ports);
 
@@ -363,12 +363,17 @@ InnerSplittedUnifiedLoopInfo::InnerSplittedUnifiedLoopInfo(size_t increment, con
     OPENVINO_ASSERT(m_outer_splitted_loop_info != nullptr, "Outer Splitted Loop Info is missed!");
 }
 
-std::shared_ptr<LoopInfo> InnerSplittedUnifiedLoopInfo::clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoPtr new_outer_splitted_loop_info) const {
+std::shared_ptr<LoopInfo> InnerSplittedUnifiedLoopInfo::clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoMap& loop_map) const {
+    if (loop_map.count(m_outer_splitted_loop_info.get()) == 0)
+        loop_map[m_outer_splitted_loop_info.get()] = m_outer_splitted_loop_info->clone_with_new_expr(expr_map, loop_map);
+
+    const auto cloned_outer_splitted_loop_info = loop_map.at(m_outer_splitted_loop_info.get());
     const auto& new_input_ports = clone_loop_ports(expr_map, m_input_ports);
     const auto& new_output_ports = clone_loop_ports(expr_map, m_output_ports);
 
     return std::make_shared<InnerSplittedUnifiedLoopInfo>(m_increment, new_input_ports, new_output_ports,
-                                                          m_input_port_descs, m_output_port_descs, m_handlers, new_outer_splitted_loop_info);
+                                                          m_input_port_descs, m_output_port_descs, m_handlers,
+                                                          std::move(cloned_outer_splitted_loop_info));
 }
 
 size_t InnerSplittedUnifiedLoopInfo::get_work_amount() const {
@@ -400,14 +405,17 @@ ExpandedLoopInfo::ExpandedLoopInfo(size_t work_amount, size_t increment,
     sort_ports();
 }
 
-std::shared_ptr<LoopInfo> ExpandedLoopInfo::clone_with_new_expr(const ExpressionMap& expr_map, UnifiedLoopInfoPtr new_unified_loop_info) const {
+std::shared_ptr<LoopInfo> ExpandedLoopInfo::clone_with_new_expr(const ExpressionMap& expr_map, LoopInfoMap& loop_map) const {
+    if (loop_map.count(m_unified_loop_info.get()) == 0)
+        loop_map[m_unified_loop_info.get()] = m_unified_loop_info->clone_with_new_expr(expr_map, loop_map);
+
+    const auto cloned_unified_loop_info = ov::as_type_ptr<UnifiedLoopInfo>(loop_map.at(m_unified_loop_info.get()));
     const auto& new_input_ports = clone_loop_ports(expr_map, m_input_ports);
     const auto& new_output_ports = clone_loop_ports(expr_map, m_output_ports);
 
-    OPENVINO_ASSERT(new_unified_loop_info, "Failed to copy ExpandedLoopInfo: new unified loop info is nullptr!");
     return std::make_shared<ExpandedLoopInfo>(m_work_amount, m_increment, new_input_ports, new_output_ports,
                                               m_ptr_increments, m_finalization_offsets, m_data_sizes, m_type,
-                                              std::move(new_unified_loop_info), m_evaluate_once);
+                                              std::move(cloned_unified_loop_info), m_evaluate_once);
 }
 
 bool ExpandedLoopInfo::is_dynamic() const {

src/common/snippets/src/lowered/loop_manager.cpp

@@ -18,36 +18,14 @@ namespace ov {
 namespace snippets {
 namespace lowered {
 
-namespace {
-LoopInfoPtr clone_loop_info(const LoopInfoPtr& loop_info, const ExpressionMap& expr_map, std::map<LoopInfoPtr, LoopInfoPtr>& cloned_map) {
-    if (cloned_map.count(loop_info))
-        return cloned_map.at(loop_info);
-
-    LoopInfoPtr cloned = nullptr;
-    if (const auto inner_splitted_loop_info = ov::as_type_ptr<InnerSplittedUnifiedLoopInfo>(loop_info)) {
-        auto new_outer = clone_loop_info(inner_splitted_loop_info->get_outer_splitted_loop_info(), expr_map, cloned_map);
-        cloned = inner_splitted_loop_info->clone_with_new_expr(expr_map, new_outer);
-    } else if (const auto expanded = ov::as_type_ptr<ExpandedLoopInfo>(loop_info)) {
-        auto new_unified = ov::as_type_ptr<UnifiedLoopInfo>(clone_loop_info(expanded->get_unified_loop_info(), expr_map, cloned_map));
-        cloned = expanded->clone_with_new_expr(expr_map, new_unified);
-    } else if (const auto unified = ov::as_type_ptr<UnifiedLoopInfo>(loop_info)) {
-        cloned = unified->clone_with_new_expr(expr_map);
-    } else {
-        OPENVINO_THROW("Undefined LoopInfo type");
-    }
-    cloned_map[loop_info] = cloned;
-    return cloned;
-}
-}  // namespace
-
 std::shared_ptr<LoopManager> LoopManager::clone_with_new_expr(const ExpressionMap& expr_map) const {
     auto new_loop_manager = std::make_shared<LoopManager>();
     // To fully cloned all LoopInfo we have to create this map [old LoopInfo -> cloned LoopInfo],
     // because some LoopInfo types contains pointer to another LoopInfo
     // so we should recurrently make a cloning of LoopInfos'
-    std::map<LoopInfoPtr, LoopInfoPtr> cloned_loop_info_map; // [ old - > cloned ]
+    LoopInfoMap loop_info_map; // [ old - > cloned ]
     for (const auto& id_info : m_map)
-        new_loop_manager->m_map.insert({id_info.first, clone_loop_info(id_info.second, expr_map, cloned_loop_info_map)});
+        new_loop_manager->m_map.insert({id_info.first, id_info.second->clone_with_new_expr(expr_map, loop_info_map)});
     new_loop_manager->next_id = next_id;
     return new_loop_manager;
 }

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

@@ -155,7 +155,7 @@ bool SplitLoops::TransformInnerSplitLoop::run(LinearIR& linear_ir, LinearIR::con
         // There is already ExpandedLoopInfo
         const auto& inner_expanded_loop_info = loop_manager->get_loop_info<ExpandedLoopInfo>(inner_loop_end->get_id());
         const auto inner_unified_loop_info = ov::as_type_ptr<InnerSplittedUnifiedLoopInfo>(inner_expanded_loop_info->get_unified_loop_info());
-        if (!inner_unified_loop_info)
+        if (!inner_unified_loop_info || inner_unified_loop_info->get_outer_splitted_loop_info() != outer_loop_info->get_unified_loop_info())
             continue;
 
         OPENVINO_ASSERT(current_dim_idx == inner_unified_loop_info->get_dim_idx(), "Incorrect processing dim index of splitted loops");

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

@@ -59,23 +59,17 @@ void ValidateExpandedLoops::validate_loop_information(const LinearIR& linear_ir)
         const auto& current_unified_loop_info = expanded_loop_info->get_unified_loop_info();
         INFORMATIVE_ASSERT(current_unified_loop_info, "expects non nullptr UnifiedLoopInfo in ExpandedLoopInfo");
 
-
-        if (initializated_info_map.count(current_unified_loop_info) == 0) {
-            auto& current_info = initializated_info_map[current_unified_loop_info];
-
-            INFORMATIVE_ASSERT(current_unified_loop_info->get_input_count() == expanded_loop_info->get_input_count() &&
-                               current_unified_loop_info->get_output_count() == expanded_loop_info->get_output_count(),
-                               "incompatible loop ports with UnifiedLoopInfo");
-
-            current_info.work_amount = 0;
-            current_info.num_ports = expanded_loop_info->get_input_count() + expanded_loop_info->get_output_count();
-            current_info.finalization_offsets.clear();
+        auto& current_info = initializated_info_map[current_unified_loop_info];
+        if (current_info.num_ports == 0) { // the info was just default constructed
+            current_info.num_ports = current_unified_loop_info->get_input_count() + current_unified_loop_info->get_output_count();
             current_info.finalization_offsets.resize(current_info.num_ports, 0);
         }
 
-        auto& current_info = initializated_info_map[current_unified_loop_info];
+        INFORMATIVE_ASSERT(current_unified_loop_info->get_input_count() == expanded_loop_info->get_input_count() &&
+                           current_unified_loop_info->get_output_count() == expanded_loop_info->get_output_count(),
+                           "incompatible loop ports with UnifiedLoopInfo");
+
         current_info.work_amount = utils::dynamic_safe_add(current_info.work_amount, expanded_loop_info->get_work_amount());
-        INFORMATIVE_ASSERT(current_unified_loop_info, "expects non nullptr current UnifiedLoopInfo");
         INFORMATIVE_ASSERT(current_unified_loop_info->get_ptr_increments() == expanded_loop_info->get_ptr_increments(),
                            "incompatible pointer increments with UnifiedLoopInfo");
 

src/common/snippets/src/runtime_configurator.cpp

@@ -65,20 +65,19 @@ void RuntimeConfigurator::initialization(const lowered::LinearIRCPtr& linear_ir)
     m_optimizer.init(linear_ir, m_io_descs, m_in_num);
 
     // InnerSplittedLoops should be inited after OuterSplittedLoops
-    m_ordered_loop_ids.clear();
-    size_t offset = 0;
     const auto& loop_map = linear_ir->get_loop_manager()->get_map();
+    m_ordered_loop_ids.clear();
+    m_ordered_loop_ids.reserve(loop_map.size());
+    std::vector<size_t> loops_must_be_last;
     for (const auto& p : loop_map) {
         const auto loop_id = p.first;
         const auto& expanded_loop_info = ov::as_type_ptr<lowered::ExpandedLoopInfo>(p.second);
         OPENVINO_ASSERT(expanded_loop_info, "UpdateLoopInfo expects ExpandedLoopInfo in LoopManager");
-        if (ov::is_type<lowered::InnerSplittedUnifiedLoopInfo>(expanded_loop_info->get_unified_loop_info())) {
-            m_ordered_loop_ids.push_back(loop_id);
-        } else {
-            m_ordered_loop_ids.insert(m_ordered_loop_ids.begin() + offset, loop_id);
-            ++offset;
-        }
+        const auto& unified_loop_info = expanded_loop_info->get_unified_loop_info();
+        auto& collection = ov::is_type<lowered::InnerSplittedUnifiedLoopInfo>(unified_loop_info) ? loops_must_be_last : m_ordered_loop_ids;
+        collection.push_back(loop_id);
     }
+    m_ordered_loop_ids.insert(m_ordered_loop_ids.end(), loops_must_be_last.cbegin(), loops_must_be_last.cend());
 }
 
 void RuntimeConfigurator::update(const lowered::LinearIRCPtr& linear_ir) {
@@ -238,10 +237,9 @@ void RuntimeConfigurator::update_buffer_scratchpad_size(const lowered::LinearIRC
     m_config->buffer_scratchpad_size = linear_ir->get_static_buffer_scratchpad_size();
 
     auto is_not_executed = [&loop_manager](const lowered::ExpressionPtr& buffer_expr) {
-        for (const auto& loop_id : buffer_expr->get_loop_ids())
-            if (loop_manager->get_loop_info(loop_id)->get_work_amount() == 0)
-                return true;
-        return false;
+        const auto& loop_ids = buffer_expr->get_loop_ids();
+        return std::any_of(loop_ids.cbegin(), loop_ids.cend(),
+                          [&loop_manager](size_t loop_id) { return loop_manager->get_loop_info(loop_id)->get_work_amount() == 0; });
     };
 
     for (const auto& p : m_dynamic_buffer_clusters) {