back port full fix

Author: chenhu-wang

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

@@ -85,6 +85,10 @@ class LinearIR::LoopManager {
         void set_outer_splited_loop(bool outer_splited_loop);
         void set_first_iter_handler(FirstIterHandler handler);
 
+        // Update the parameters of existing LoopPorts
+        void update_entry_points(const std::function<void(LoopPort&)>& updater);
+        void update_exit_points(const std::function<void(LoopPort&)>& updater);
+
     private:
         size_t m_work_amount = 0;
         size_t m_increment = 0;
@@ -101,8 +105,6 @@ class LinearIR::LoopManager {
     using LoopInfoPtr = std::shared_ptr<LoopInfo>;
 
     std::shared_ptr<LoopManager> clone_with_new_expr(const ExressionMap& expr_map) const;
-    size_t add_loop_info(const LoopInfoPtr& loop);
-    void remove_loop_info(size_t index);
     LoopInfoPtr get_loop_info(size_t index) const;
     size_t get_loop_count() const { return m_map.size(); }
     const std::map<size_t, LoopInfoPtr>& get_map() const;
@@ -182,23 +184,55 @@ class LinearIR::LoopManager {
     void expression_replacement(constExprIt new_expr_begin, constExprIt new_expr_end, const ExpressionPtr& decomposed_expr,
                                 size_t loop_id, const std::vector<ExpressionPort>& new_entries, const std::vector<ExpressionPort>& exits);
 
-    // Note: these methods find iterators of first entry loop point and last exit point (bounds of Loop)
-    //       If there are already inserted LoopBegin and LoopEnd in Linear IR, the methods can find them as well if `loop_ops_inserted` = true
-    void get_loop_bounds(const LinearIR& linear_ir,
-                         size_t loop_id,
-                         LinearIR::constExprIt& loop_begin_pos,
-                         LinearIR::constExprIt& loop_end_pos,
-                         bool loop_ops_inserted = false) const;
-    static void get_loop_bounds(const LinearIR& linear_ir,
-                                const std::vector<LoopPort>& entries,
-                                const std::vector<LoopPort>& exits,
-                                LinearIR::constExprIt& loop_begin_pos,
-                                LinearIR::constExprIt& loop_end_pos,
-                                size_t loop_id, bool loop_ops_inserted = false);
+    /**
+     * @brief Find bounds of Loop:
+     *        - If the explicit Loop exprs with the target `loop_id` have been inserted,
+     *          Loop bounds are these iterators of the corresponding LoopBegin and LoopEnd.
+     *        - Otherwise Loop bounds are iterators of the first entry loop port (or Scalar, VectorBuffer and another LoopBegin that
+     *          are in this Loop but have another `loop_id`) and the next iterator of the last exit loop port (or another LoopEnd that
+     *          are in this Loop but have another `loop_id`).
+     * @param linear_ir linear IR
+     * @param loop_id target Loop ID
+     * @return the pair of loop_begin_pos and loop_end_pos iterators
+     */
+    std::pair<LinearIR::constExprIt, LinearIR::constExprIt> get_loop_bounds(const LinearIR& linear_ir, size_t loop_id) const;
+    /**
+     * @brief Find bounds of Loop:
+     *        - If the explicit Loop exprs with the target `loop_id` have been inserted,
+     *          Loop bounds are these iterators of the corresponding LoopBegin and LoopEnd.
+     *        - Otherwise Loop bounds are iterators of the first entry loop port (or Scalar, VectorBuffer and another LoopBegin that
+     *          are in this Loop but have another `loop_id`) and the next iterator of the last exit loop port (or another LoopEnd that
+     *          are in this Loop but have another `loop_id`).
+     * @param linear_ir linear IR
+     * @param loop_id target Loop ID
+     * @param entries input loop ports
+     * @param exits output loop ports
+     * @return the pair of loop_begin_pos and loop_end_pos iterators
+     */
+    static std::pair<LinearIR::constExprIt, LinearIR::constExprIt> get_loop_bounds(const LinearIR& linear_ir, size_t loop_id,
+                                                                                   const std::vector<LoopPort>& entries, const std::vector<LoopPort>& exits);
 
     LoopPort get_loop_port_by_expr_port(const ExpressionPort& expr_port, const size_t loop_id);
 
 private:
+    /**
+     * @brief Add new Loop Info to the map
+     * @param loop target loop info
+     * @return the loop ID
+     */
+    size_t add_loop_info(const LoopInfoPtr& loop);
+    /**
+     * @brief Remove LoopInfo from the map
+     * @param index the target index of Loop
+     */
+    void remove_loop_info(size_t index);
+    /**
+     * @brief Find expression ports in bounds that are connected to consumers or parent that aren't in these bounds
+     * @param loop_begin_pos the first expression iterator of the Loop
+     * @param loop_end_pos the next iterator after the last expression
+     * @param entries found input expression ports
+     * @param exits found output expression ports
+     */
     static void get_io_loop_ports(LinearIR::constExprIt loop_begin_pos,
                                   LinearIR::constExprIt loop_end_pos,
                                   std::vector<ExpressionPort>& entries,
@@ -220,6 +254,7 @@ class LinearIR::LoopManager {
     //                                         for `before` the new Loop is the most outer Loop
     void insert_loop_id(const ExpressionPtr& expr, size_t new_id, bool before = true, size_t target_id = SIZE_MAX);
     void insert_loop_ids(const ExpressionPtr& expr, const std::vector<size_t>& new_ids, bool before = true, size_t target_id = SIZE_MAX);
+    static bool is_loop_id_found(const ExpressionPtr& expr, size_t id);
 
     std::map<size_t, LoopInfoPtr> m_map = {};
     size_t next_id = 0;

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

@@ -141,6 +141,14 @@ void LoopInfo::set_first_iter_handler(LoopInfo::FirstIterHandler first_iter_hand
     m_first_iter_handler = std::move(first_iter_handler);
 }
 
+void LoopInfo::update_entry_points(const std::function<void(LoopPort&)>& updater) {
+    std::for_each(m_entry_points.begin(), m_entry_points.end(), updater);
+}
+
+void LoopInfo::update_exit_points(const std::function<void(LoopPort&)>& updater) {
+    std::for_each(m_exit_points.begin(), m_exit_points.end(), updater);
+}
+
 bool operator==(const LinearIR::LoopManager::LoopPort& lhs, const LinearIR::LoopManager::LoopPort& rhs) {
     if (&lhs == &rhs)
         return true;
@@ -194,26 +202,19 @@ std::vector<size_t> LinearIR::LoopManager::get_outer_expr_loops(const Expression
     return std::vector<size_t>(loop_ids.cbegin(), it);
 }
 
-void LinearIR::LoopManager::get_loop_bounds(const LinearIR &linear_ir,
-                                            size_t loop_id,
-                                            LinearIR::constExprIt &loop_begin_pos,
-                                            LinearIR::constExprIt &loop_end_pos,
-                                            bool loop_ops_inserted) const {
+std::pair<LinearIR::constExprIt, LinearIR::constExprIt> LinearIR::LoopManager::get_loop_bounds(const LinearIR &linear_ir, size_t loop_id) const {
     const auto loop_info = get_loop_info(loop_id);
-    get_loop_bounds(linear_ir, loop_info->get_entry_points(), loop_info->get_exit_points(), loop_begin_pos, loop_end_pos, loop_id, loop_ops_inserted);
+    return get_loop_bounds(linear_ir, loop_id, loop_info->get_entry_points(), loop_info->get_exit_points());
 }
 
-void LinearIR::LoopManager::get_loop_bounds(const LinearIR &linear_ir,
-                                            const std::vector<LoopPort>& entries,
-                                            const std::vector<LoopPort>& exits,
-                                            LinearIR::constExprIt &loop_begin_pos,
-                                            LinearIR::constExprIt &loop_end_pos,
-                                            size_t loop_id, bool loop_ops_inserted) {
+std::pair<LinearIR::constExprIt, LinearIR::constExprIt> LinearIR::LoopManager::get_loop_bounds(const LinearIR &linear_ir, size_t loop_id,
+                                                                                               const std::vector<LoopPort>& entries,
+                                                                                               const std::vector<LoopPort>& exits) {
     OPENVINO_ASSERT(!entries.empty(), "Loop must have entry points");
     OPENVINO_ASSERT(!exits.empty(), "Loop must have entry points");
-    const auto& entry_expr = entries.front().expr_port->get_expr();
-    loop_begin_pos = linear_ir.find(entry_expr);
 
+    const auto& entry_expr = entries.front().expr_port->get_expr();
+    auto loop_begin_pos = linear_ir.find(entry_expr);
     // Some operations in Loop can be before first entry points: Scalars, VectorBuffer.
     // We should iterate by them till the expr is in the corresponding Loop
     auto prev_loop_ids = (*std::prev(loop_begin_pos))->get_loop_ids();
@@ -222,18 +223,28 @@ void LinearIR::LoopManager::get_loop_bounds(const LinearIR &linear_ir,
         prev_loop_ids = (*std::prev(loop_begin_pos))->get_loop_ids();
     }
 
-    if (loop_ops_inserted) {
-        const auto loop_begin = ov::as_type_ptr<op::LoopBegin>((*std::prev(loop_begin_pos))->get_node());
-        OPENVINO_ASSERT(loop_begin, "Failed explicit loop bounds getting: LoopBegin has not been found");
-        const auto loop_end = loop_begin->get_loop_end();
-        OPENVINO_ASSERT(loop_end->get_id() == loop_id, "Failed explicit loop bounds getting: Loop bounds with correct ID have not been found");
-        loop_begin_pos = std::prev(loop_begin_pos);
-        loop_end_pos = linear_ir.find_after(loop_begin_pos, linear_ir.get_expr_by_node(loop_end));
-    } else {
-        // At the moment all Loops must have exit points
-        const auto& exit_expr = exits.back().expr_port->get_expr();
-        loop_end_pos = std::next(linear_ir.find_after(loop_begin_pos, exit_expr));
+    const auto& exit_expr = exits.back().expr_port->get_expr();
+    auto loop_end_pos = std::next(linear_ir.find_after(loop_begin_pos, exit_expr));
+    // There might be LoopEnd with another `loop_id` but in the target Loop as well.
+    auto current_loop_ids = (*loop_end_pos)->get_loop_ids();
+    while (std::find(current_loop_ids.begin(), current_loop_ids.end(), loop_id) != current_loop_ids.end()) {
+        loop_end_pos = std::next(loop_end_pos);
+        current_loop_ids = (*loop_end_pos)->get_loop_ids();
     }
+
+    // Check for the existing LoopBegin/LoopEnd
+    if (const auto loop_end = ov::as_type_ptr<op::LoopEnd>((*loop_end_pos)->get_node())) {
+        if (loop_end->get_id() == loop_id) {
+            // loop_begin_pos is iterator of LoopBegin now
+            // loop_end_pos is iterator of LoopEnd now
+            loop_begin_pos = std::prev(loop_begin_pos);
+            const auto loop_begin = loop_end->get_loop_begin();
+            OPENVINO_ASSERT((*loop_begin_pos)->get_node() == loop_begin, "LoopBegin has not been found!");
+        }
+    }
+
+    OPENVINO_ASSERT(loop_begin_pos != linear_ir.cend() && loop_end_pos != linear_ir.cend(), "Loop bounds haven't been found!");
+    return std::make_pair(loop_begin_pos, loop_end_pos);
 }
 
 LinearIR::LoopManager::LoopPort LinearIR::LoopManager::get_loop_port_by_expr_port(const ExpressionPort& expr_port, const size_t loop_id) {
@@ -353,27 +364,35 @@ size_t LinearIR::LoopManager::replace_with_new_loop(const LinearIR& linear_ir,
                                                     const std::vector<LoopPort>& entries,
                                                     const std::vector<LoopPort>& exits,
                                                     const size_t old_id) {
-    const auto loop_info = std::make_shared<LoopManager::LoopInfo>(work_amount, increment, entries, exits);
-    const auto loop_id = this->add_loop_info(loop_info);
-    for (auto expr_it = loop_begin_pos; expr_it != loop_end_pos; ++expr_it) {
+    const auto is_bound_explicit_loop_begin = ov::is_type<op::LoopBegin>(loop_begin_pos->get()->get_node());
+    const auto is_bound_explicit_loop_end = ov::is_type<op::LoopEnd>(std::prev(loop_end_pos)->get()->get_node());
+    OPENVINO_ASSERT((is_bound_explicit_loop_begin && is_bound_explicit_loop_end) ||
+                        (!is_bound_explicit_loop_begin && !is_bound_explicit_loop_end),
+                    "Incorrect LoopBounds!");
+    const auto explicit_loop_bounds = is_bound_explicit_loop_begin && is_bound_explicit_loop_end;
+
+    const auto loop_id =
+        this->add_loop_info(std::make_shared<LoopManager::LoopInfo>(work_amount, increment, entries, exits));
+    const auto begin = explicit_loop_bounds ? std::next(loop_begin_pos) : loop_begin_pos;
+    const auto end = explicit_loop_bounds ? std::prev(loop_end_pos) : loop_end_pos;
+    for (auto expr_it = begin; expr_it != end; ++expr_it) {
         replace_loop_id(*expr_it, old_id, loop_id);
     }
 
-    const auto old_loop_info = this->get_loop_info(old_id);
-    const auto old_loop_begin_pos = linear_ir.find(old_loop_info->get_entry_points().front().expr_port->get_expr());
-    const auto old_loop_end_pos = linear_ir.find(old_loop_info->get_exit_points().back().expr_port->get_expr());
+    // Check that other expression in LinearIR doesn't have the old loop ID - otherwise completely removed from loop
+    // manager
+    const auto old_loop_bounds = get_loop_bounds(linear_ir, old_id);
     // If new bounds are equal to old loop bounds, this means that old Loop is removed totally from LIR
     // In this case old loop info must be completely removed from loop manager
-    if (loop_begin_pos == old_loop_begin_pos && loop_end_pos == old_loop_end_pos) {
+    if (loop_begin_pos == old_loop_bounds.first && loop_end_pos == old_loop_bounds.second) {
         this->remove_loop_info(old_id);
     }
     return loop_id;
 }
 
 void LinearIR::LoopManager::fuse_loops(const LinearIR& linear_ir, size_t loop_id_upper, size_t loop_id_lower, bool fuse_into_upper) {
-    LinearIR::constExprIt loop_begin_target, loop_end_target;
-    get_loop_bounds(linear_ir, fuse_into_upper ? loop_id_lower : loop_id_upper, loop_begin_target, loop_end_target);
-    fuse_loops(loop_begin_target, loop_end_target, loop_id_upper, loop_id_lower, fuse_into_upper);
+    const auto loop_bounds = get_loop_bounds(linear_ir, fuse_into_upper ? loop_id_lower : loop_id_upper);
+    fuse_loops(loop_bounds.first, loop_bounds.second, loop_id_upper, loop_id_lower, fuse_into_upper);
 }
 
 void LinearIR::LoopManager::fuse_loops(LinearIR::constExprIt loop_begin_target, LinearIR::constExprIt loop_end_target,
@@ -541,6 +560,7 @@ void LinearIR::LoopManager::sort_loop_ports(LinearIR::constExprIt& loop_begin_po
 
 void LinearIR::LoopManager::insert_loop_id(const ExpressionPtr& expr, size_t new_id, bool before, size_t target_id) {
     OPENVINO_ASSERT(m_map.count(new_id) == 1, "Failed marking expression by Loop ID: the Loop with this ID hasn't registered");
+    OPENVINO_ASSERT(!is_loop_id_found(expr, new_id), "Expression cannot have several the same Loop IDs");
     auto& loop_ids = expr->m_loop_ids;
     OPENVINO_ASSERT(std::find(loop_ids.cbegin(), loop_ids.cend(), new_id) == loop_ids.cend(),
                     "Expression cannot have several the same Loop IDs");
@@ -568,6 +588,7 @@ void LinearIR::LoopManager::insert_loop_ids(const ExpressionPtr& expr, const std
 
 void LinearIR::LoopManager::replace_loop_id(const ExpressionPtr& expr, size_t prev_id, size_t new_id) {
     OPENVINO_ASSERT(m_map.count(new_id), "Failed marking expression by Loop ID: the Loop with this ID hasn't registered");
+    OPENVINO_ASSERT(!is_loop_id_found(expr, new_id), "Expression cannot have several the same Loop IDs");
     auto& loop_ids = expr->m_loop_ids;
     OPENVINO_ASSERT(std::find(loop_ids.cbegin(), loop_ids.cend(), new_id) == loop_ids.cend(),
                     "Expression already has the Loop with ID " + std::to_string(new_id));
@@ -584,6 +605,11 @@ void LinearIR::LoopManager::remove_loop_id(const ExpressionPtr& expr, size_t id)
     loop_ids.erase(it);
 }
 
+bool LinearIR::LoopManager::is_loop_id_found(const ExpressionPtr& expr, size_t id) {
+    const auto loop_ids = expr->get_loop_ids();
+    return std::find(loop_ids.cbegin(), loop_ids.cend(), id) != loop_ids.cend();
+}
+
 }// namespace lowered
 }// namespace snippets
 }// namespace ov

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

@@ -68,10 +68,9 @@ void FuseLoops::move(LinearIR& linear_ir, const LinearIR::LoopManagerPtr& loop_m
     std::map<size_t, std::pair<LinearIR::constExprIt, LinearIR::constExprIt>> outer_loops;  // The map: LoopID -> [ LoopBegin, LoopEnd ]
     const auto outer_loop_ids = LinearIR::LoopManager::get_outer_expr_loops(*loop_begin_pos, loop_id);
     for (const auto& loop_id : outer_loop_ids) {
-        LinearIR::constExprIt begin, end;
-        loop_manager->get_loop_bounds(linear_ir, loop_id, begin, end);
+        const auto loop_bounds = loop_manager->get_loop_bounds(linear_ir, loop_id);
         // save previos iterator since the current iterator can be moved
-        outer_loops[loop_id] = {std::prev(begin), end};
+        outer_loops[loop_id] = {std::prev(loop_bounds.first), loop_bounds.second};
     }
      // Secondly, move expressions
     for (auto it = loop_begin_pos; it != loop_end_pos;) {
@@ -122,7 +121,7 @@ bool FuseLoops::fuse_upper_into_current(LinearIR& linear_ir, const LinearIR::Loo
         return false;
 
     LinearIR::constExprIt target_loop_begin_pos, target_loop_end_pos;
-    loop_manager->get_loop_bounds(linear_ir, target_loop_id, target_loop_begin_pos, target_loop_end_pos);
+    std::tie(target_loop_begin_pos, target_loop_end_pos) = loop_manager->get_loop_bounds(linear_ir, target_loop_id);
     loop_manager->fuse_loops(target_loop_begin_pos, target_loop_end_pos, target_loop_id, current_loop_id, false);
     // Update work_amount for Loop (increment is constant because increments must be the identical for fusion):
     loop_current->set_work_amount(std::max(loop_current->get_work_amount(), loop_target->get_work_amount()));
@@ -167,7 +166,7 @@ bool FuseLoops::fuse_lower_into_current(LinearIR& linear_ir, const LinearIR::Loo
         return false;
 
     LinearIR::constExprIt target_loop_begin_pos, target_loop_end_pos;
-    loop_manager->get_loop_bounds(linear_ir, target_loop_id, target_loop_begin_pos, target_loop_end_pos);
+    std::tie(target_loop_begin_pos, target_loop_end_pos) = loop_manager->get_loop_bounds(linear_ir, target_loop_id);
     loop_manager->fuse_loops(target_loop_begin_pos, target_loop_end_pos, current_loop_id, target_loop_id);
     // Update work_amount for Loop (increment is constant because increments must be the identical for fusion):
     loop_current->set_work_amount(std::max(loop_current->get_work_amount(), loop_target->get_work_amount()));
@@ -212,7 +211,7 @@ bool FuseLoops::run(LinearIR& linear_ir) {
 
             const auto current_loop_info = loop_manager->get_loop_info(current_loop_id);
             LinearIR::constExprIt current_loop_begin_pos, current_loop_end_pos;
-            loop_manager->get_loop_bounds(linear_ir, current_loop_id, current_loop_begin_pos, current_loop_end_pos);
+            std::tie(current_loop_begin_pos, current_loop_end_pos) = loop_manager->get_loop_bounds(linear_ir, current_loop_id);
 
             // We fuse upper Loops into the current till we can do it.
             // After that we fuse lower Loops into the current till we can do it.