copyb

Author: aobolensk

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_copy_b_emitter.cpp

@@ -15,13 +15,15 @@
 #include <unordered_set>
 #include <vector>
 
+#include "emitters/snippets/aarch64/jit_binary_call_emitter.hpp"
 #include "emitters/snippets/aarch64/kernel_executors/gemm_copy_b.hpp"
 #include "emitters/snippets/aarch64/utils.hpp"
 #include "emitters/snippets/utils/utils.hpp"
 #include "emitters/utils.hpp"
 #include "openvino/core/node.hpp"
 #include "openvino/core/type.hpp"
 #include "openvino/core/type/element_type.hpp"
+#include "snippets/emitter.hpp"
 #include "snippets/kernel_executor_table.hpp"
 #include "snippets/lowered/expression.hpp"
 #include "snippets/utils/utils.hpp"
@@ -38,7 +40,7 @@ jit_gemm_copy_b_emitter::jit_gemm_copy_b_emitter(jit_generator* h,
                                                  cpu_isa_t isa,
                                                  const ExpressionPtr& expr,
                                                  const snippets::KernelExecutorTablePtr& kernel_table)
-    : jit_emitter(h, isa) {
+    : jit_binary_call_emitter(h, isa, expr->get_live_regs()) {
     in_out_type_ = emitter_in_out_map::gpr_to_gpr;
     const auto gemm_repack = ov::as_type_ptr<GemmCopyB>(expr->get_node());
     OV_CPU_JIT_EMITTER_ASSERT(gemm_repack, "expects GemmCopyB node");
@@ -78,17 +80,28 @@ void jit_gemm_copy_b_emitter::validate_arguments(const std::vector<size_t>& in,
 
 void jit_gemm_copy_b_emitter::emit_impl(const std::vector<size_t>& in, const std::vector<size_t>& out) const {
     validate_arguments(in, out);
-    // todo: use optimized reg spill after CVS-162498
-    std::unordered_set<size_t> exclude = {};
-    store_context(exclude);
+
+    std::vector<size_t> mem_ptrs_idxs{in[0], out[0]};
+
+    init_binary_call_regs(2, mem_ptrs_idxs);
+    emit_call(mem_ptrs_idxs);
+}
+
+void jit_gemm_copy_b_emitter::emit_call(const std::vector<size_t>& mem_ptrs_idxs) const {
+    const auto& regs_to_spill = get_regs_to_spill();
+    std::unordered_set<size_t> exclude_spill = {};
+    for (const auto& reg : regs_to_spill) {
+        if (reg.type == snippets::RegType::gpr) {
+            exclude_spill.insert(reg.idx);
+        }
+    }
+    store_context(exclude_spill);
 
     Xbyak_aarch64::XReg x0(0);
     Xbyak_aarch64::XReg x1(1);
     Xbyak_aarch64::XReg x2(2);
     Xbyak_aarch64::XReg aux_reg(3);
 
-    // Prepare memory pointers with offsets
-    std::vector<size_t> mem_ptrs_idxs{in[0], out[0]};
     const auto& mem_ptrs = utils::transform_idxs_to_regs(mem_ptrs_idxs);
 
     // Apply memory offsets and load adjusted pointers
@@ -111,11 +124,11 @@ void jit_gemm_copy_b_emitter::emit_impl(const std::vector<size_t>& in, const std
     const auto& compiled_kernel = get_compiled_kernel_ptr();
     h->mov(x0, compiled_kernel);
 
-    Xbyak_aarch64::XReg func_reg(9);
-    h->mov(func_reg, get_execute_function_ptr());
-    h->blr(func_reg);
+    const auto& call_address_reg = get_call_address_reg();
+    h->mov(call_address_reg, get_execute_function_ptr());
+    h->blr(call_address_reg);
 
-    restore_context(exclude);
+    restore_context(exclude_spill);
 }
 
 uintptr_t jit_gemm_copy_b_emitter::get_compiled_kernel_ptr() const {
@@ -125,4 +138,4 @@ uintptr_t jit_gemm_copy_b_emitter::get_compiled_kernel_ptr() const {
 uintptr_t jit_gemm_copy_b_emitter::get_execute_function_ptr() {
     return reinterpret_cast<const uintptr_t>(GemmCopyBKaiKernelExecutor::execute);
 }
-}  // namespace ov::intel_cpu::aarch64
+}  // namespace ov::intel_cpu::aarch64

src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_copy_b_emitter.hpp

@@ -4,11 +4,12 @@
 
 #pragma once
 
-#include "emitters/plugin/aarch64/jit_emitter.hpp"
+#include "emitters/snippets/aarch64/jit_binary_call_emitter.hpp"
 #include "emitters/snippets/aarch64/kernel_executors/gemm_copy_b.hpp"
+#include "snippets/emitter.hpp"
 
 namespace ov::intel_cpu::aarch64 {
-class jit_gemm_copy_b_emitter : public jit_emitter {
+class jit_gemm_copy_b_emitter : public jit_binary_call_emitter {
 public:
     jit_gemm_copy_b_emitter(dnnl::impl::cpu::aarch64::jit_generator* h,
                             dnnl::impl::cpu::aarch64::cpu_isa_t isa,
@@ -25,6 +26,7 @@ class jit_gemm_copy_b_emitter : public jit_emitter {
 protected:
     void validate_arguments(const std::vector<size_t>& in, const std::vector<size_t>& out) const override;
     void emit_impl(const std::vector<size_t>& in, const std::vector<size_t>& out) const override;
+    void emit_call(const std::vector<size_t>& mem_ptrs_idxs) const;
 
     static uintptr_t get_execute_function_ptr();
     uintptr_t get_compiled_kernel_ptr() const;