Lazy initialization for reg spiller. Snippets tests are green

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

@@ -28,8 +28,10 @@ bool InitLiveRanges::run(LinearIR& linear_ir) {
             || ov::is_type<op::PerfCountBeginBase>(op)
             || ov::is_type<op::PerfCountEndBase>(op)
 #endif
-            )
+            ) {
+            m_reg_manager.set_live_regs(expr, {});
             continue;
+        }
         OPENVINO_ASSERT(expr->get_output_count() == op->get_output_size(), "Incorrect count of output port descriptors!");
         const double start = expr->get_exec_num();
         // Remove all regs that expired before start

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

@@ -21,6 +21,7 @@ namespace pass {
 
 bool InsertRegSpills::run(LinearIR& linear_ir) {
     OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::InsertRegSpills")
+    return false;
 
     auto needs_reg_spill = [](const ExpressionPtr& expr) {
         return ov::is_type<snippets::op::Brgemm>(expr->get_node());

src/plugins/intel_cpu/src/emitters/plugin/x64/utils.cpp

@@ -34,55 +34,58 @@ inline snippets::Reg Xbyak2SnippetsReg(const Xbyak::Reg& xb_reg) {
 template<cpu_isa_t isa,
         typename std::enable_if<dnnl::impl::utils::one_of(isa, sse41, avx2, avx512_core), bool>::type = true>
 struct regs_to_spill {
-    static std::vector<Xbyak::Reg> get(jit_generator* h) {
-        std::vector<Xbyak::Reg> regs_to_spill = {h->r8, h->r9, h->r10, h->r11, h->r12, h->r13, h->r14, h->r15,
-                                                 h->rax, h->rbx, h->rcx, h->rdx, h->rdi, h->rsi, h->rbp};
+    static std::vector<Xbyak::Reg> get(const std::set<snippets::Reg>& live_regs) {
+        std::vector<Xbyak::Reg> regs_to_spill;
+        auto push_if_live = [&live_regs, &regs_to_spill](Xbyak::Reg&& reg) {
+            if (live_regs.empty() || live_regs.count(Xbyak2SnippetsReg(reg)))
+                regs_to_spill.emplace_back(reg);
+        };
+        // todo: according to ABI, caller needs to save only RAX, RCX, RDX, R8, R9, R10, R11
+        // 0    1    2    3    4    5    6    7
+        // rax, rcx, rdx, rbx, rsp, rbp, rsi, rdi, r8, r9, r10, r11, r12, r13, r14, r15
+        for (int i = 0; i < 16; i++) {
+            // do not spill rsp;
+            if (i != 4)
+                push_if_live(Reg64(i));
+        }
+
         for (int i = 0; i < cpu_isa_traits<isa>::n_vregs; ++i)
-            regs_to_spill.push_back(typename cpu_isa_traits<isa>::Vmm(i));
+            push_if_live(typename cpu_isa_traits<isa>::Vmm(i));
+
         const int num_k_mask = isa == cpu_isa_t::avx512_core ? 8 : 0;
-        for (size_t i = 0; i < num_k_mask; ++i)
-            regs_to_spill.push_back(Xbyak::Opmask(i));
+        for (int i = 0; i < num_k_mask; ++i)
+            push_if_live(Xbyak::Opmask(i));
         return regs_to_spill;
     }
 };
-std::vector<Xbyak::Reg> get_regs_to_spill(jit_generator* h, cpu_isa_t isa) {
+std::vector<Xbyak::Reg> get_regs_to_spill(cpu_isa_t isa, const std::set<snippets::Reg>& live_regs) {
     switch (isa) {
-        case sse41: return regs_to_spill<sse41>::get(h);
-        case avx2: return regs_to_spill<avx2>::get(h);
-        case avx512_core: return regs_to_spill<avx512_core>::get(h);
+        case sse41: return regs_to_spill<sse41>::get(live_regs);
+        case avx2: return regs_to_spill<avx2>::get(live_regs);
+        case avx512_core: return regs_to_spill<avx512_core>::get(live_regs);
         default: OPENVINO_THROW("Unhandled isa in get_regs_to_spill");
     }
 }
 
 
 
 EmitABIRegSpills::EmitABIRegSpills(jit_generator* h_arg) : h(h_arg), isa(get_isa()) {
-    // all regs to spill according to ABI
-    m_regs_to_spill = get_regs_to_spill(h, isa);
-    for (const auto& reg : m_regs_to_spill) {
-        const auto reg_bit_size = reg.getBit();
-        OPENVINO_ASSERT(reg_bit_size % 8 == 0, "Unexpected reg bit size");
-        m_bytes_to_spill += reg_bit_size / 8;
-    }
-}
-
-void EmitABIRegSpills::limit_to_live_regs(const std::set<snippets::Reg>& live_regs) {
-    auto not_live = [&live_regs](const Xbyak::Reg& reg) -> bool {
-        return live_regs.count(Xbyak2SnippetsReg(reg)) == 0;
-    };
-    auto remove_it = std::remove_if(m_regs_to_spill.begin(), m_regs_to_spill.end(), not_live);
-    for (auto it = remove_it; it != m_regs_to_spill.end(); it++)
-        m_bytes_to_spill -= it->getBit() / 8;
-    m_regs_to_spill.erase(remove_it, m_regs_to_spill.end());
 }
 
 EmitABIRegSpills::~EmitABIRegSpills() {
     OPENVINO_ASSERT(spill_status, "postamble or preamble is missed");
     OPENVINO_ASSERT(rsp_status, "rsp_align or rsp_restore is missed");
 }
 
-void EmitABIRegSpills::preamble() {
+void EmitABIRegSpills::preamble(const std::set<snippets::Reg>& live_regs) {
     OPENVINO_ASSERT(spill_status, "Attempt to spill ABI registers twice in a row");
+    // all regs to spill according to ABI
+    m_regs_to_spill = get_regs_to_spill(isa, live_regs);
+    for (const auto& reg : m_regs_to_spill) {
+        const auto reg_bit_size = reg.getBit();
+        OPENVINO_ASSERT(reg_bit_size % 8 == 0, "Unexpected reg bit size");
+        m_bytes_to_spill += reg_bit_size / 8;
+    }
     h->sub(h->rsp, m_bytes_to_spill);
     uint32_t byte_stack_offset = 0;
     for (const auto& reg : m_regs_to_spill) {

src/plugins/intel_cpu/src/emitters/plugin/x64/utils.hpp

@@ -14,12 +14,10 @@ namespace intel_cpu {
 class EmitABIRegSpills {
 public:
     EmitABIRegSpills(dnnl::impl::cpu::x64::jit_generator* h);
-
-    void limit_to_live_regs(const std::set<snippets::Reg>& live_regs );
     ~EmitABIRegSpills();
 
     // push (save) all registers on the stack
-    void preamble();
+    void preamble(const std::set<snippets::Reg>& live_regs = {});
     // pop (take) all registers from the stack
     void postamble();
 

src/plugins/intel_cpu/src/emitters/snippets/x64/jit_reg_spill_emitters.cpp

@@ -44,8 +44,7 @@ void jit_reg_spill_begin_emitter::emit_code(const std::vector<size_t> &in, const
 void jit_reg_spill_begin_emitter::emit_impl(const std::vector<size_t>& in, const std::vector<size_t>& out) const {
     const auto& reg_info = m_reg_spill_begin_expr->get_reg_info();
     OPENVINO_ASSERT(reg_info.second.empty(), "Out regs should be empty for reg_spill_begin emitter");
-    m_abi_reg_spiller->limit_to_live_regs({reg_info.first.begin(), reg_info.first.end()});
-    m_abi_reg_spiller->preamble();
+    m_abi_reg_spiller->preamble({reg_info.first.begin(), reg_info.first.end()});
 }
 
 /* ============================================================== */