Skip to content

Commit

Permalink
[Snippets][CPU] Optimized reg spills in BrgemmCopyBKernel (#28635)
Browse files Browse the repository at this point in the history
### Details:
- *BrgemmCopyBKernel is JIT compiled kernel which emits blocking loop by
`N` dimension with binary call of the kernel
`dnnl::x64::brgemm_copy_b_kernel` inside. Currently, we spill all
registers before the each call of this kernel. It means that if the
dimension `N` is quite big, there will be many binary calls and reg
spills. These regular reg spills may lead to perf degradation for
Subgraphs with small `M`, `K,` and big `N` dimensions. This PR uses
feature of binary call optimizations - spills only live registers*

### Tickets:
 - *N/A*
  • Loading branch information
a-sidorova authored Jan 31, 2025
1 parent 9930aea commit 7e8c945
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,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) {
Expand Down Expand Up @@ -298,6 +298,16 @@ void BrgemmCopyBKernel::emit_brgemm_copy_b_kernel_call(size_t N,
spill.postamble();
}

std::set<snippets::Reg> BrgemmCopyBKernel::get_live_regs() const {
// Only the registers `src_reg`, `tr_src_reg` and `comp_reg` should be
// saved on each `jit_brgemm_matmul_copy_b_t` binary call.
// They're ABI parameter registers (caller saved). So we have to manually
// spills only them on each `jit_brgemm_matmul_copy_b_t` binary call
return {{snippets::RegType::gpr, static_cast<size_t>(src_reg.getIdx())},
{snippets::RegType::gpr, static_cast<size_t>(tr_src_reg.getIdx())},
{snippets::RegType::gpr, static_cast<size_t>(comp_reg.getIdx())}};
}

void BrgemmCopyBKernel::execute(matmul::jit_brgemm_matmul_copy_b_t* kernel,
const void* src,
const void* dst,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down

0 comments on commit 7e8c945

Please sign in to comment.