[Snippets][CPU] Optimized reg spills in BrgemmCopyBKernel

openvinotoolkit · Jan 23, 2025 · 321f479 · 321f479
1 parent 8dde87c
commit 321f479
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 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
@@ -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
@@ -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;