Ivan's comments applied

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

@@ -21,18 +21,7 @@ class LinearIR::LoopManager {
 
     struct LoopPort {
         LoopPort() = default;
-        LoopPort(const ExpressionPort& port, bool is_incremented = true, size_t dim_idx = 0)
-            : expr_port(std::make_shared<ExpressionPort>(port)),
-              is_incremented(is_incremented),
-              dim_idx(dim_idx) {
-            OPENVINO_ASSERT(dim_idx < port.get_descriptor_ptr()->get_shape().size(),
-                            "LoopPort dim_idx (",
-                            dim_idx,
-                            ") must be less than the corresponding expression port shape rank (",
-                            port.get_descriptor_ptr()->get_shape().size(),
-                            ")");
-        }
-
+        LoopPort(const ExpressionPort& port, bool is_incremented = true, size_t dim_idx = 0);
         std::shared_ptr<LoopPort> clone_with_new_expr(const ExpressionPtr& new_expr) const;
 
         friend bool operator==(const LoopPort& lhs, const LoopPort& rhs);
@@ -51,35 +40,63 @@ class LinearIR::LoopManager {
 
     class LoopInfo {
     public:
+        enum {UNDEFINED_DIM_IDX = std::numeric_limits<size_t>::max()};
         LoopInfo() = default;
         LoopInfo(size_t work_amount, size_t increment,
                  const std::vector<LoopPort>& entries,
-                 const std::vector<LoopPort>& exits)
-            : work_amount(work_amount), increment(increment),
-              entry_points(entries), exit_points(exits), outer_splited_loop(false) {}
+                 const std::vector<LoopPort>& exits,
+                 bool outer_splited_loop = false)
+            : m_work_amount(work_amount), m_increment(increment),
+              m_entry_points(entries), m_exit_points(exits), m_outer_splited_loop(outer_splited_loop) {}
         LoopInfo(size_t work_amount, size_t increment,
                  const std::vector<ExpressionPort>& entries,
-                 const std::vector<ExpressionPort>& exits);
+                 const std::vector<ExpressionPort>& exits,
+                 bool outer_splited_loop = false);
 
         std::shared_ptr<LoopInfo> clone_with_new_expr(const ExressionMap& expr_map) const;
-        // Returns dimension index if dimension indices for all entry and exit points are equal, and SIZE_MAX otherwise
-        size_t get_dim_idx() const;
 
-        // TODO: replace this temporary solution when ticket 119851 is implemented
+        // Returns dimension index if dimension indices for all entry and exit points are equal, and UNDEFINED_DIM_IDX otherwise
+        size_t get_dim_idx() const;
+        size_t get_work_amount() const;
+        size_t get_increment() const;
+        const std::vector<LoopPort>& get_entry_points() const;
+        const std::vector<LoopPort>& get_exit_points() const;
+        bool get_outer_splited_loop() const;
+
+        /**
+         * \brief Inserts a separate body for first loop iteration processing if needed.
+         * Can also modify both main and first iter loop bodies.
+         * TODO: replace this temporary solution when ticket 119851 is implemented
+         *
+         * \param linear_ir LIR which should be modified
+         * \param loop_end_it iterator on LoopEnd expression for which the handler is called
+         *
+         * \return bool value which indicates whether the linear_ir was changed or not.
+         */
         using FirstIterHandler = std::function<bool(LinearIR&, LinearIR::constExprIt)>;
+        const FirstIterHandler& get_first_iter_handler() const;
+
+        // Sets dim_idx to all entry and exit points
+        void set_dim_idx(size_t dim_idx);
+        void set_work_amount(size_t work_amount);
+        void set_increment(size_t increment);
+        void set_entry_points(std::vector<LoopPort> entry_points);
+        void set_exit_points(std::vector<LoopPort> exit_points);
+        void set_outer_splited_loop(bool outer_splited_loop);
         void set_first_iter_handler(FirstIterHandler handler);
-        FirstIterHandler fst_iter_handler = nullptr;
 
-        size_t work_amount = 0;
-        size_t increment = 0;
+    private:
+        size_t m_work_amount = 0;
+        size_t m_increment = 0;
         // The order of entry and exit expressions is important:
         //     - The position before first entry expr is Loop Begin position
         //     - The position after last exit expr is Loop End position
         // Note: Scalars aren't entry expressions but can be before first entry expr in Linear IR
-        std::vector<LoopPort> entry_points = {};
-        std::vector<LoopPort> exit_points = {};
+        std::vector<LoopPort> m_entry_points = {};
+        std::vector<LoopPort> m_exit_points = {};
         // True if this Loop is outer Loop for nested Loops that splits the same dimension
-        bool outer_splited_loop = false;
+        bool m_outer_splited_loop = false;
+        FirstIterHandler m_first_iter_handler = nullptr;
     };
     using LoopInfoPtr = std::shared_ptr<LoopInfo>;
 
@@ -106,12 +123,7 @@ class LinearIR::LoopManager {
                      const std::vector<T>& entries,
                      const std::vector<T>& exits) {
         const auto loop_info = std::make_shared<LoopManager::LoopInfo>(work_amount, work_amount_increment, entries, exits);
-        auto set_common_dim_idx = [dim_idx](std::vector<LoopPort>& ports) {
-            for (auto& port : ports)
-                port.dim_idx = dim_idx;
-        };
-        set_common_dim_idx(loop_info->entry_points);
-        set_common_dim_idx(loop_info->exit_points);
+        loop_info->set_dim_idx(dim_idx);
         const auto loop_id = this->add_loop_info(loop_info);
         for (auto expr_it = loop_begin_pos; expr_it != loop_end_pos; ++expr_it) {
             insert_loop_id(*expr_it, loop_id);
@@ -134,13 +146,14 @@ class LinearIR::LoopManager {
         return loop_id;
     }
 
-    size_t mark_loop_with_old_loop_replacement(LinearIR::constExprIt loop_begin_pos,
-                                               LinearIR::constExprIt loop_end_pos,
-                                               size_t work_amount,
-                                               size_t increment,
-                                               const std::vector<LoopPort>& entries,
-                                               const std::vector<LoopPort>& exits,
-                                               const size_t old_id);
+    size_t replace_with_new_loop(const LinearIR& linear_ir,
+                                 LinearIR::constExprIt loop_begin_pos,
+                                 LinearIR::constExprIt loop_end_pos,
+                                 size_t work_amount,
+                                 size_t increment,
+                                 const std::vector<LoopPort>& entries,
+                                 const std::vector<LoopPort>& exits,
+                                 const size_t old_id);
 
     void fuse_loops(const LinearIR& linear_ir, size_t loop_id_upper, size_t loop_id_lower, bool fuse_into_upper = true);
     void fuse_loops(LinearIR::constExprIt loop_begin_target, LinearIR::constExprIt loop_end_target,

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

@@ -25,14 +25,9 @@ class InitLoops : public Pass {
     bool run(LinearIR& linear_ir) override;
 
 private:
-    static void init_ptr_increments(std::vector<LinearIR::LoopManager::LoopPort>& loop_inputs,
-                                    std::vector<LinearIR::LoopManager::LoopPort>& loop_outputs,
-                                    size_t work_amount);
-    static void init_finalization_offsets(std::vector<LinearIR::LoopManager::LoopPort>& loop_inputs,
-                                          std::vector<LinearIR::LoopManager::LoopPort>& loop_outputs,
-                                          size_t work_amount);
-    static void init_element_type_sizes(std::vector<LinearIR::LoopManager::LoopPort>& loop_inputs,
-                                        std::vector<LinearIR::LoopManager::LoopPort>& loop_outputs);
+    static void init_ptr_increments(const LinearIR::LoopManager::LoopInfoPtr& loop_info);
+    static void init_finalization_offsets(const LinearIR::LoopManager::LoopInfoPtr& loop_info);
+    static void init_element_type_sizes(const LinearIR::LoopManager::LoopInfoPtr& loop_info);
 };
 
 } // namespace pass

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

@@ -30,7 +30,7 @@ class BroadcastLoad : public MemoryAccess {
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
     void validate_and_infer_types() override;
     const ov::Dimension& get_bcast_dimension() {return bcast_dimension;}
-    void set_bcast_dimension(const ov::Dimension new_dim) {bcast_dimension = new_dim;}
+    void set_bcast_dimension(ov::Dimension new_dim) {bcast_dimension = std::move(new_dim);}
 
     // Note:BroadcastMove and BroadcastLoad are implemented as separate classes,
     // but have identical shapeInfer semantics. In order to avoid code duplication,

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

@@ -29,7 +29,7 @@ class BroadcastMove : public ov::op::Op {
 
     void validate_and_infer_types() override;
     const ov::Dimension& get_bcast_dimension() {return bcast_dimension;}
-    void set_bcast_dimension(const ov::Dimension new_dim) {bcast_dimension = new_dim;}
+    void set_bcast_dimension(ov::Dimension new_dim) {bcast_dimension = std::move(new_dim);}
     // Note:BroadcastMove and BroadcastLoad are implemented as separate classes,
     // but have identical shapeInfer semantics. In order to avoid code duplication,
     // we created dummy ShapeInfer classes that are essentially instantiations

src/common/snippets/src/generator.cpp

@@ -27,6 +27,13 @@ void Generator::generate(lowered::LinearIR& linear_ir, LoweringResult& result, c
         return get_op_reg_type(op);
     };
     lowered::pass::PassPipeline lowered_pipeline;
+    // Note: the order of all passes in this pipeline must not be changed since they have hard dependencies
+    //    1. InsertTailLoop must be called after AssignRegisters since tail loop expressions must have the same
+    //       assigned registers as the corresponding ops in the main body.
+    //    2. CleanupLoopOffsets must be called after InsertTailLoop to avoid violating the proportionality of the pointer increments
+    //       (this might happen if tail loop and main loop have different increments)
+    //    3. OptimizeLoopSingleEvaluation must be called after CleanupLoopOffsets
+    //       since CleanupLoopOffsets can't handle loops with evaluate_once = true
     lowered_pipeline.register_pass<lowered::pass::AssignRegisters>(reg_type_mapper);
     lowered_pipeline.register_pass<lowered::pass::InsertTailLoop>();
     lowered_pipeline.register_pass<lowered::pass::CleanupLoopOffsets>();