From 321f479c2505f9a04c17e244fb302787e171d16d Mon Sep 17 00:00:00 2001
From: Alexandra Sidorova <alexandra.sidorova@intel.com>
Date: Thu, 23 Jan 2025 14:30:00 +0400
Subject: [PATCH] [Snippets][CPU] Optimized reg spills in BrgemmCopyBKernel

---
 .../snippets/x64/kernel_executors/brgemm_copy_b.cpp    | 10 +++++++++-
 .../snippets/x64/kernel_executors/brgemm_copy_b.hpp    |  2 ++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/plugins/intel_cpu/src/emitters/snippets/x64/kernel_executors/brgemm_copy_b.cpp b/src/plugins/intel_cpu/src/emitters/snippets/x64/kernel_executors/brgemm_copy_b.cpp
index 30d95ed6a2bf7a..433f6c2bf16b58 100644
--- a/src/plugins/intel_cpu/src/emitters/snippets/x64/kernel_executors/brgemm_copy_b.cpp
+++ b/src/plugins/intel_cpu/src/emitters/snippets/x64/kernel_executors/brgemm_copy_b.cpp
@@ -252,7 +252,7 @@ void BrgemmCopyBKernel::emit_brgemm_copy_b_kernel_call(size_t N,
                                                        size_t offset_out,
                                                        size_t offset_comp) {
     EmitABIRegSpills spill(this);
-    spill.preamble();
+    spill.preamble(get_live_regs());
 
     const auto add_offset = [&](Xbyak::Reg64 reg, size_t bytes_offset) {
         if (bytes_offset)
@@ -295,6 +295,14 @@ void BrgemmCopyBKernel::emit_brgemm_copy_b_kernel_call(size_t N,
     spill.postamble();
 }
 
+std::set<snippets::Reg> BrgemmCopyBKernel::get_live_regs() const {
+    std::set<snippets::Reg> live_regs = {{snippets::RegType::gpr, static_cast<size_t>(src_reg.getIdx())},
+                                         {snippets::RegType::gpr, static_cast<size_t>(tr_src_reg.getIdx())}};
+    if (is_with_comp)
+        live_regs.insert({snippets::RegType::gpr, static_cast<size_t>(comp_reg.getIdx())});
+    return live_regs;
+}
+
 void BrgemmCopyBKernel::execute(matmul::jit_brgemm_matmul_copy_b_t* kernel,
                                 const void* src,
                                 const void* dst,
diff --git a/src/plugins/intel_cpu/src/emitters/snippets/x64/kernel_executors/brgemm_copy_b.hpp b/src/plugins/intel_cpu/src/emitters/snippets/x64/kernel_executors/brgemm_copy_b.hpp
index 5ef740067f2035..594702fed4ae3c 100644
--- a/src/plugins/intel_cpu/src/emitters/snippets/x64/kernel_executors/brgemm_copy_b.hpp
+++ b/src/plugins/intel_cpu/src/emitters/snippets/x64/kernel_executors/brgemm_copy_b.hpp
@@ -170,6 +170,8 @@ struct BrgemmCopyBKernel : public RepackedInputKernel, public dnnl::impl::cpu::x
     void init_brgemm_copy_b_kernel(std::unique_ptr<dnnl::impl::cpu::x64::matmul::jit_brgemm_matmul_copy_b_t>& kernel,
                                    const BrgemmCopyBKernelConfig& conf) const;
 
+    std::set<snippets::Reg> get_live_regs() const;
+
     static constexpr auto abi_param_regs = dnnl::impl::cpu::x64::abi_param_regs;
     const Xbyak::Reg64 src_reg = abi_param2;
     const Xbyak::Reg64 tr_src_reg = abi_param3;