[Snippets] InsertSpecificIterations lowered pass fix (#32305)

### Details:
- *`InsertSpecificIterations` is refactored to improve readability (some
usings are used)*
- *`InsertSpecificIterations`: fixed the issue which led to non-unique
loop indices existence: previously, when outer loop was cloned, inner
expanded loop infos kept same loop id even for the cloned parts. This PR
addresses this problem by performing inner loops cloning on the cloned
outer loop*
 - *The changes have been covered by new lowered test*

### Tickets:
 - *part of CVS-172631*

Author: v-Golubev

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

@@ -81,32 +81,30 @@ class InsertSpecificIterations : public RangedPass {
     /**
      * @brief Make a copy of Loop with ID `loop_id` and insert to LinearIR before `insert_pos`
      * @param linear_ir target Linear IR
-     * @param loop_id the target loop ID
+     * @param bounds loop bounds of current loop
      * @param insert_pos insertion position iterator
-     * @param new_entry_ports reference of vector with Loop input ports that will be updated after insertion
-     * @param new_exit_ports reference of vector with Loop output ports that will be updated after insertion
+     * @param expression_map expression map to store pairs [original_expr, new_expr]
      * @return LoopBounds: iterators of new LoopBegin and LoopEnd
      */
     static LoopManager::LoopBounds insert_copy_loop(LinearIR& linear_ir,
-                                                    size_t loop_id,
+                                                    const LoopManager::LoopBounds& bounds,
                                                     const LinearIR::constExprIt& insert_pos,
-                                                    std::vector<LoopPort>& new_entry_ports,
-                                                    std::vector<LoopPort>& new_exit_ports);
+                                                    ExpressionMap& expression_map);
     /**
      * @brief Initializes decomposed loop: update ptr arithmetic, work_amout, increment, ID
      * @param linear_ir target Linear IR
-     * @param begin iterator of LoopBegin
-     * @param end iterator of LoopEnd
+     * @param decomposed_loop_bounds decomposed loop bounds
      * @param decomposed_loop_info loop info of the corresponding decomposed loop
-     * @param unified_loop_id ID of the unified loop
+     * @param loop_id_to_replace ID of the loop which should be replaced by the decomposed one
      * @param decomposed_loop_end LoopEnd of the decomposed loop
+     * @param run_handlers flag to run handlers for the decomposed loop
      */
     static void init_decomposed_loop(LinearIR& linear_ir,
-                                     LinearIR::constExprIt begin,
-                                     LinearIR::constExprIt end,
+                                     const LoopManager::LoopBounds& decomposed_loop_bounds,
                                      const ExpandedLoopInfoPtr& decomposed_loop_info,
-                                     size_t unified_loop_id,
-                                     const std::shared_ptr<op::LoopEnd>& decomposed_loop_end);
+                                     size_t loop_id_to_replace,
+                                     const std::shared_ptr<op::LoopEnd>& decomposed_loop_end,
+                                     bool run_handlers);
 };
 
 }  // namespace ov::snippets::lowered::pass

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

@@ -9,8 +9,8 @@
 #include <cstddef>
 #include <cstdint>
 #include <iterator>
+#include <map>
 #include <memory>
-#include <tuple>
 #include <vector>
 
 #include "openvino/core/except.hpp"
@@ -33,11 +33,22 @@
 namespace ov::snippets::lowered::pass {
 
 namespace {
-void connect_cloned_body_with_buffers_outside(LinearIR::constExprIt cur_begin,
-                                              LinearIR::constExprIt cur_end,
-                                              LinearIR::constExprIt res_begin,
-                                              LinearIR::constExprIt res_end,
+std::vector<LoopPort> clone_ports(const ExpressionMap& expression_map, const std::vector<LoopPort>& cur_ports) {
+    std::vector<LoopPort> new_ports(cur_ports.size());
+    for (size_t i = 0; i < cur_ports.size(); ++i) {
+        const auto& port = cur_ports[i];
+        const auto& original_expr = port.get_expr_port()->get_expr().get();
+        OPENVINO_ASSERT(expression_map.count(original_expr), "Cannot find cloned expression for: ", original_expr);
+        new_ports[i] = *port.clone_with_new_expr(expression_map.at(original_expr));
+    }
+    return new_ports;
+}
+
+void connect_cloned_body_with_buffers_outside(const LoopManager::LoopBounds& cur_bounds,
+                                              const LoopManager::LoopBounds& res_bounds,
                                               LinearIR& linear_ir) {
+    const auto& [cur_begin, cur_end] = cur_bounds;
+    const auto& [res_begin, res_end] = res_bounds;
     for (auto result_it = res_begin, original_it = cur_begin; result_it != res_end; ++result_it, ++original_it) {
         const auto& result_expr = *result_it;
         const auto& original_expr = *original_it;
@@ -136,68 +147,49 @@ size_t InsertSpecificIterations::get_decomposed_loop_increment(const UnifiedLoop
 }
 
 LoopManager::LoopBounds InsertSpecificIterations::insert_copy_loop(LinearIR& linear_ir,
-                                                                   const size_t loop_id,
+                                                                   const LoopManager::LoopBounds& bounds,
                                                                    const LinearIR::constExprIt& insert_pos,
-                                                                   std::vector<LoopPort>& new_entry_ports,
-                                                                   std::vector<LoopPort>& new_exit_ports) {
-    const auto& loop_manager = linear_ir.get_loop_manager();
-    const auto [loop_begin_pos, loop_end_pos] = loop_manager->get_loop_bounds(linear_ir, loop_id);
-
-    ExpressionMap expression_map;
+                                                                   ExpressionMap& expression_map) {
+    const auto& [loop_begin_pos, loop_end_pos] = bounds;
     const auto& cloning_config = LinearIRBuilder::Config(false);
     const auto& loop_copy_range =
         LinearIRBuilder(cloning_config).clone_range(loop_begin_pos, std::next(loop_end_pos), expression_map);
     const auto new_loop_begin_pos = linear_ir.insert(insert_pos, loop_copy_range.begin(), loop_copy_range.end());
     const auto new_loop_end_pos = std::prev(insert_pos);
-
-    // Add connections between output of cloned bodies and Buffers from the current LinearIR (Buffers are connections
-    // between Loops)
-    connect_cloned_body_with_buffers_outside(loop_begin_pos,
-                                             loop_end_pos,
-                                             new_loop_begin_pos,
-                                             new_loop_end_pos,
-                                             linear_ir);
-
-    auto clone_ports = [&expression_map](const std::vector<LoopPort>& ports, std::vector<LoopPort>& new_ports) {
-        new_ports.resize(ports.size());
-        for (size_t i = 0; i < ports.size(); ++i) {
-            const auto& port = ports[i];
-            new_ports[i] = *port.clone_with_new_expr(expression_map[port.get_expr_port()->get_expr().get()]);
-        }
-    };
-    const auto original_loop_info = loop_manager->get_loop_info(loop_id);
-    clone_ports(original_loop_info->get_input_ports(), new_entry_ports);
-    clone_ports(original_loop_info->get_output_ports(), new_exit_ports);
-
     return {new_loop_begin_pos, new_loop_end_pos};
 }
 
 void InsertSpecificIterations::init_decomposed_loop(LinearIR& linear_ir,
-                                                    LinearIR::constExprIt begin,
-                                                    LinearIR::constExprIt end,
+                                                    const LoopManager::LoopBounds& decomposed_loop_bounds,
                                                     const ExpandedLoopInfoPtr& decomposed_loop_info,
-                                                    size_t unified_loop_id,
-                                                    const std::shared_ptr<op::LoopEnd>& decomposed_loop_end) {
+                                                    size_t loop_id_to_replace,
+                                                    const std::shared_ptr<op::LoopEnd>& decomposed_loop_end,
+                                                    bool run_handlers) {
     const auto& loop_manager = linear_ir.get_loop_manager();
-    const auto new_id =
-        loop_manager->replace_with_new_loop(linear_ir, begin, std::next(end), decomposed_loop_info, unified_loop_id);
+    const auto new_id = loop_manager->replace_with_new_loop(linear_ir,
+                                                            decomposed_loop_bounds.first,
+                                                            std::next(decomposed_loop_bounds.second),
+                                                            decomposed_loop_info,
+                                                            loop_id_to_replace);
     decomposed_loop_end->set_id(new_id);
     decomposed_loop_end->set_work_amount(decomposed_loop_info->get_work_amount());
     decomposed_loop_end->set_increment(decomposed_loop_info->get_increment());
     decomposed_loop_end->set_ptr_increments(decomposed_loop_info->get_ptr_increments());
     decomposed_loop_end->set_finalization_offsets(decomposed_loop_info->get_finalization_offsets());
-    // Note: handlers must be run on the range started with the first operation in the loop body.
-    const auto handlers = decomposed_loop_info->get_handler_passes();
-    handlers.run(linear_ir, std::next(begin), end);
+    if (run_handlers) {
+        const auto handlers = decomposed_loop_info->get_handler_passes();
+        // Note: handlers must be run on the range started with the first operation in the loop body.
+        handlers.run(linear_ir, std::next(decomposed_loop_bounds.first), decomposed_loop_bounds.second);
+    }
 }
 
 bool InsertSpecificIterations::decompose(LinearIR& linear_ir,
                                          LinearIR::constExprIt begin,
                                          LinearIR::constExprIt end,
                                          const std::shared_ptr<op::LoopEnd>& loop_end) {
-    const auto loop_id = loop_end->get_id();
+    const auto unified_loop_id = loop_end->get_id();
     const auto& loop_manager = linear_ir.get_loop_manager();
-    const auto& unified_loop_info = loop_manager->get_loop_info<UnifiedLoopInfo>(loop_id);
+    const auto& unified_loop_info = loop_manager->get_loop_info<UnifiedLoopInfo>(unified_loop_id);
 
     auto remaining_work_amount = unified_loop_info->get_work_amount();
     const auto is_wa_dynamic = utils::is_dynamic_value(remaining_work_amount);
@@ -219,7 +211,7 @@ bool InsertSpecificIterations::decompose(LinearIR& linear_ir,
             }
 
             auto decomposed_loop_end = loop_end;
-            auto decomposed_loop_begin_it = begin, decomposed_loop_end_it = end;
+            LoopManager::LoopBounds decomposed_loop_bounds{begin, end};
             auto decomposed_loop_entry_ports = unified_loop_info->get_input_ports();
             auto decomposed_loop_exit_ports = unified_loop_info->get_output_ports();
             auto decomposed_ptr_increments = unified_loop_info->get_ptr_increments();
@@ -228,13 +220,19 @@ bool InsertSpecificIterations::decompose(LinearIR& linear_ir,
             // Need to copy body if there are other specific sup-loops
             // Otherwise we should update the current body
             if (remaining_work_amount > 0) {
-                std::tie(decomposed_loop_begin_it, decomposed_loop_end_it) =
-                    insert_copy_loop(linear_ir,
-                                     loop_id,
-                                     begin,
-                                     decomposed_loop_entry_ports,
-                                     decomposed_loop_exit_ports);
-                decomposed_loop_end = ov::as_type_ptr<op::LoopEnd>(decomposed_loop_end_it->get()->get_node());
+                const auto cur_bounds = loop_manager->get_loop_bounds(linear_ir, unified_loop_id);
+                ExpressionMap expression_map;
+                decomposed_loop_bounds = insert_copy_loop(linear_ir, cur_bounds, begin, expression_map);
+
+                // Add connections between output of cloned bodies and Buffers from the current LinearIR
+                // (Buffers are connections between Loops)
+                connect_cloned_body_with_buffers_outside(cur_bounds, decomposed_loop_bounds, linear_ir);
+
+                const auto original_loop_info = loop_manager->get_loop_info(unified_loop_id);
+                decomposed_loop_entry_ports = clone_ports(expression_map, original_loop_info->get_input_ports());
+                decomposed_loop_exit_ports = clone_ports(expression_map, original_loop_info->get_output_ports());
+
+                decomposed_loop_end = ov::as_type_ptr<op::LoopEnd>(decomposed_loop_bounds.second->get()->get_node());
                 OPENVINO_ASSERT(decomposed_loop_end, "Cloned Loop does not contain LoopEnd op at the expected place.");
 
                 // Only latest loop iterations must have summarized finalization offsets!
@@ -246,6 +244,59 @@ bool InsertSpecificIterations::decompose(LinearIR& linear_ir,
                                       offset = 0;
                                   }
                               });
+
+                std::map<UnifiedLoopInfoPtr, UnifiedLoopInfoPtr> unified_loop_map;
+                auto get_unified_cloned_info = [&unified_loop_map,
+                                                &expression_map](const ExpandedLoopInfoPtr& expanded_loop_info) {
+                    const auto& unified_loop_info = expanded_loop_info->get_unified_loop_info();
+                    if (unified_loop_map.count(unified_loop_info) == 0) {
+                        LoopInfoMap loop_info_map;
+                        // Note: we must clone UnifiedLoopInfo for the cloned ExpandedLoopInfos
+                        auto cloned_info = ov::as_type_ptr<UnifiedLoopInfo>(
+                            unified_loop_info->clone_with_new_expr(expression_map, loop_info_map));
+                        OPENVINO_ASSERT(cloned_info, "cloned info must be UnifiedLoopInfo");
+                        unified_loop_map[unified_loop_info] = cloned_info;
+                    }
+                    OPENVINO_ASSERT(unified_loop_map.count(unified_loop_info),
+                                    "Cloned UnifiedLoopInfo must be cloned at this stage.");
+                    return unified_loop_map[unified_loop_info];
+                };
+
+                // Note: all internal decomposed loops must be also cloned to avoid a situation
+                // when 2 loops with the same ID exist in both specific iterations of the outer loop
+                for (auto it = std::next(decomposed_loop_bounds.first); it != decomposed_loop_bounds.second; ++it) {
+                    auto internal_loop_end = ov::as_type_ptr<op::LoopEnd>(it->get()->get_node());
+                    if (!internal_loop_end) {
+                        continue;
+                    }
+                    const auto loop_begin = internal_loop_end->get_loop_begin();
+                    auto begin_it = linear_ir.find_after(std::next(decomposed_loop_bounds.first),
+                                                         linear_ir.get_expr_by_node(loop_begin));
+                    OPENVINO_ASSERT(begin_it != linear_ir.cend(),
+                                    "Cannot find LoopBegin for LoopEnd with id ",
+                                    internal_loop_end->get_id());
+                    LoopManager::LoopBounds internal_loop_bounds{begin_it, it};
+                    const auto internal_loop_id = internal_loop_end->get_id();
+                    // Note: internal loops must be already decomposed to ExpandedLoops
+                    const auto internal_loop_info = loop_manager->get_loop_info<ExpandedLoopInfo>(internal_loop_id);
+                    const auto cloned_loop_info = std::make_shared<ExpandedLoopInfo>(
+                        internal_loop_info->get_work_amount(),
+                        internal_loop_info->get_increment(),
+                        clone_ports(expression_map, internal_loop_info->get_input_ports()),
+                        clone_ports(expression_map, internal_loop_info->get_output_ports()),
+                        internal_loop_info->get_ptr_increments(),
+                        internal_loop_info->get_finalization_offsets(),
+                        internal_loop_info->get_data_sizes(),
+                        internal_loop_info->get_type(),
+                        get_unified_cloned_info(internal_loop_info),
+                        internal_loop_info->is_evaluate_once());
+                    init_decomposed_loop(linear_ir,
+                                         internal_loop_bounds,
+                                         cloned_loop_info,
+                                         internal_loop_id,
+                                         internal_loop_end,
+                                         false);
+                }
             }
 
             const auto decomposed_loop_info = std::make_shared<ExpandedLoopInfo>(work_amount,
@@ -258,11 +309,11 @@ bool InsertSpecificIterations::decompose(LinearIR& linear_ir,
                                                                                  iter_type,
                                                                                  unified_loop_info);
             init_decomposed_loop(linear_ir,
-                                 decomposed_loop_begin_it,
-                                 decomposed_loop_end_it,
+                                 decomposed_loop_bounds,
                                  decomposed_loop_info,
-                                 loop_id,
-                                 decomposed_loop_end);
+                                 unified_loop_id,
+                                 decomposed_loop_end,
+                                 true);
 
             decomposed = true;
         }