diff --git a/src/common/snippets/include/snippets/lowered/pass/insert_tail_loop.hpp b/src/common/snippets/include/snippets/lowered/pass/insert_tail_loop.hpp index faafd8186b8448..bd255339543e01 100644 --- a/src/common/snippets/include/snippets/lowered/pass/insert_tail_loop.hpp +++ b/src/common/snippets/include/snippets/lowered/pass/insert_tail_loop.hpp @@ -24,7 +24,9 @@ class InsertTailLoop : public Pass { public: OPENVINO_RTTI("InsertTailLoop", "Pass") bool run(LinearIR& linear_ir) override; - static LinearIR::container copy_loop(const LinearIR& linear_ir, const size_t loop_id); + static LinearIR::constExprIt insert_copy_loop(LinearIR& linear_ir, + const size_t loop_id, + const LinearIR::constExprIt& insert_pos); static constexpr size_t existing_subtensor_value = SIZE_MAX; static void propagate_updated_subtensor_through_loop(const LinearIR& linear_ir, diff --git a/src/common/snippets/src/lowered/loop_manager.cpp b/src/common/snippets/src/lowered/loop_manager.cpp index ee151fc2c78760..e7e83361ee0a39 100644 --- a/src/common/snippets/src/lowered/loop_manager.cpp +++ b/src/common/snippets/src/lowered/loop_manager.cpp @@ -358,6 +358,15 @@ size_t LinearIR::LoopManager::replace_with_new_loop(const LinearIR& linear_ir, for (auto expr_it = loop_begin_pos; expr_it != loop_end_pos; ++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()); + // 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) { + this->remove_loop_info(old_id); + } return loop_id; } diff --git a/src/common/snippets/src/lowered/pass/insert_tail_loop.cpp b/src/common/snippets/src/lowered/pass/insert_tail_loop.cpp index cc685c1851157a..6c23d77536a0de 100644 --- a/src/common/snippets/src/lowered/pass/insert_tail_loop.cpp +++ b/src/common/snippets/src/lowered/pass/insert_tail_loop.cpp @@ -127,12 +127,14 @@ void InsertTailLoop::propagate_updated_subtensor_through_loop(const LinearIR& li (*expr_it)->updateShapes(); } -LinearIR::container InsertTailLoop::copy_loop(const LinearIR& linear_ir, const size_t loop_id) { +LinearIR::constExprIt InsertTailLoop::insert_copy_loop(LinearIR& linear_ir, const size_t loop_id, const LinearIR::constExprIt& insert_pos) { const auto& loop_manager = linear_ir.get_loop_manager(); LinearIR::constExprIt loop_begin_pos, loop_end_pos; loop_manager->get_loop_bounds(linear_ir, loop_id, loop_begin_pos, loop_end_pos, true); ExressionMap expression_map; const auto& loop_copy_range = LinearIR::deep_copy_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 = insert_pos; const auto original_loop_info = loop_manager->get_loop_info(loop_id); std::vector new_entry_points, new_exit_points; @@ -156,11 +158,9 @@ LinearIR::container InsertTailLoop::copy_loop(const LinearIR& linear_ir, const s loop_manager->update_loops_port(outer_loop_ids, expr->get_output_port(i), {expr->get_output_port(i), new_expr->get_output_port(i)}, false); } - const auto new_loop_begin_pos = loop_copy_range.begin(); - const auto new_loop_end_pos = loop_copy_range.end(); const auto new_id = loop_manager->replace_with_new_loop(linear_ir, - std::next(new_loop_begin_pos), - std::prev(new_loop_end_pos), + new_loop_begin_pos, + new_loop_end_pos, original_loop_info->get_work_amount(), original_loop_info->get_increment(), new_entry_points, @@ -169,7 +169,7 @@ LinearIR::container InsertTailLoop::copy_loop(const LinearIR& linear_ir, const s const auto loop_end = ov::as_type_ptr(std::prev(new_loop_end_pos)->get()->get_node()); OPENVINO_ASSERT(loop_end, "Cloned Loop does not contain LoopEnd op at the expected place."); loop_end->set_id(new_id); - return loop_copy_range; + return new_loop_begin_pos; } void InsertTailLoop::create_tail_loop(LinearIR& linear_ir, @@ -186,17 +186,16 @@ void InsertTailLoop::create_tail_loop(LinearIR& linear_ir, auto original_loop_info = loop_manager->get_loop_info(original_loop_id); auto tail_loop_info = original_loop_info; if (need_vector_loop) { - const auto new_loop_range = copy_loop(linear_ir, original_loop_id); - const auto new_loop_end = ov::as_type_ptr(std::prev(new_loop_range.end())->get()->get_node()); - OPENVINO_ASSERT(new_loop_end, "Cloned Loop does not contain LoopEnd op at the expected place."); - tail_loop_info = original_loop_info; - original_loop_info = loop_manager->get_loop_info(new_loop_end->get_id()); - // Note: new loop body is inserted before the original loop // So new loop becomes a main vector loop, the original loop becomes tail loop // This is done in such way to have original ops from the main body at the end: // this allows us to conveniently interact with outer loops in further passes - linear_ir.insert(begin, new_loop_range.begin(), new_loop_range.end()); + const auto new_loop_begin_pos = insert_copy_loop(linear_ir, original_loop_id, begin); + const auto new_loop_begin = ov::as_type_ptr(new_loop_begin_pos->get()->get_node()); + OPENVINO_ASSERT(new_loop_begin, "Cloned Loop does not contain LoopBegin op at the expected place."); + const auto new_loop_end = new_loop_begin->get_loop_end(); + tail_loop_info = original_loop_info; + original_loop_info = loop_manager->get_loop_info(new_loop_end->get_id()); const auto new_vector_loop_wa = original_loop_info->get_work_amount() - tail_size; original_loop_info->set_work_amount(new_vector_loop_wa); diff --git a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp index fc9aeeac10ee92..cf6752c8fb67e3 100644 --- a/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp +++ b/src/plugins/intel_cpu/src/transformations/snippets/x64/pass/lowered/brgemm_blocking.cpp @@ -160,16 +160,18 @@ bool BrgemmBlocking::run(LinearIR& linear_ir) { if (work_amount <= increment) return false; - auto new_loop_range = snippets::lowered::pass::InsertTailLoop::copy_loop(linear_ir, loop_id); - const auto firt_iter_loop_end = ov::as_type_ptr(std::prev(new_loop_range.end())->get()->get_node()); + const auto loop_begin_it = linear_ir.find(linear_ir.get_expr_by_node(loop_end->get_loop_begin())); + const auto new_loop_begin_pos = + snippets::lowered::pass::InsertTailLoop::insert_copy_loop(linear_ir, loop_id, loop_begin_it); + const auto new_loop_begin = + ov::as_type_ptr(new_loop_begin_pos->get()->get_node()); + OPENVINO_ASSERT(new_loop_begin, "Cloned Loop does not contain LoopBegin op at the expected place."); + const auto firt_iter_loop_end = new_loop_begin->get_loop_end(); auto first_iter_loop_info = loop_manager->get_loop_info(firt_iter_loop_end->get_id()); firt_iter_loop_end->set_work_amount(increment); first_iter_loop_info->set_work_amount(increment); firt_iter_loop_end->set_finalization_offsets(std::vector(loop_end->get_finalization_offsets().size(), 0)); - const auto loop_begin_it = linear_ir.find(linear_ir.get_expr_by_node(loop_end->get_loop_begin())); - linear_ir.insert(loop_begin_it, new_loop_range.begin(), new_loop_range.end()); - const auto new_work_amount = work_amount - increment; loop_info->set_work_amount(new_work_amount); loop_end->set_work_amount(new_work_amount);