Address comments

Author: aobolensk

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.hpp

@@ -70,6 +70,8 @@ class jit_emitter : public ov::snippets::Emitter {
     static std::set<std::vector<element::Type>> get_supported_precisions(
         const std::shared_ptr<ov::Node>& node = nullptr);
 
+    static constexpr int sp_alignment = 16;
+
 protected:
     static size_t get_max_vecs_count();
     static int32_t get_vec_length();
@@ -155,6 +157,10 @@ class jit_emitter : public ov::snippets::Emitter {
         }
     }
 
+    int32_t get_gpr_length() const {
+        return h->x0.getBit() / 8;
+    }
+
 private:
     mutable std::vector<size_t> preserved_vec_idxs;
     mutable std::vector<size_t> preserved_gpr_idxs;
@@ -179,10 +185,6 @@ class jit_emitter : public ov::snippets::Emitter {
         return 32;
     }
 
-    int32_t get_gpr_length() const {
-        return h->x0.getBit() / 8;
-    }
-
     void store_context(const std::vector<size_t>& gpr_regs,
                        const std::vector<size_t>& vec_regs,
                        const std::unordered_set<size_t>& ignore_vec_regs = {}) const;

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

@@ -18,8 +18,8 @@
 
 #include "emitters/snippets/aarch64/kernel_executors/gemm_copy_b.hpp"
 #include "emitters/snippets/aarch64/utils.hpp"
-#include "emitters/snippets/utils/utils.hpp"
 #include "emitters/snippets/jit_snippets_call_args.hpp"
+#include "emitters/snippets/utils/utils.hpp"
 #include "emitters/utils.hpp"
 #include "openvino/core/node.hpp"
 #include "openvino/core/type.hpp"
@@ -94,13 +94,12 @@ void jit_gemm_copy_b_emitter::emit_impl(const std::vector<size_t>& in, const std
     const auto& mem_ptrs = utils::transform_idxs_to_regs(mem_ptrs_idxs);
 
     // Apply memory offsets to input/output pointers
-    h->sub(h->sp, h->sp, 2 * get_vec_length());
-
+    const auto sp_size = rnd_up(2 * get_gpr_length(), sp_alignment);
+    h->sub(h->sp, h->sp, sp_size);
     // Apply offsets and store pointers on stack
     for (size_t i = 0; i < mem_ptrs.size(); i++) {
         const auto& ptr_reg = mem_ptrs[i];
-        int32_t stack_offset = i * get_vec_length();
-
+        int32_t stack_offset = i * get_gpr_length();
         if (ov::snippets::utils::is_dynamic_value(m_memory_offsets[i])) {
             // Dynamic offset: read from runtime parameters
             size_t runtime_offset = GET_OFF(buffer_offsets) + m_buffer_ids[i] * sizeof(size_t);
@@ -110,13 +109,12 @@ void jit_gemm_copy_b_emitter::emit_impl(const std::vector<size_t>& in, const std
             utils::push_ptr_with_static_offset_on_stack(h, stack_offset, ptr_reg, m_memory_offsets[i]);
         }
     }
-
     // Load back the adjusted pointers for function call
     h->ldr(x1, Xbyak_aarch64::ptr(h->sp));                    // input pointer
-    h->ldr(x2, Xbyak_aarch64::ptr(h->sp, get_vec_length()));  // output pointer
+    h->ldr(x2, Xbyak_aarch64::ptr(h->sp, get_gpr_length()));  // output pointer
 
     // Restore stack pointer
-    h->add(h->sp, h->sp, 2 * get_vec_length());
+    h->add(h->sp, h->sp, sp_size);
 
     // Set up executor pointer as first argument
     const auto& compiled_kernel = get_compiled_kernel_ptr();

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

@@ -18,8 +18,8 @@
 
 #include "emitters/snippets/aarch64/kernel_executors/gemm.hpp"
 #include "emitters/snippets/aarch64/utils.hpp"
-#include "emitters/snippets/utils/utils.hpp"
 #include "emitters/snippets/jit_snippets_call_args.hpp"
+#include "emitters/snippets/utils/utils.hpp"
 #include "emitters/utils.hpp"
 #include "openvino/core/node.hpp"
 #include "openvino/core/type/element_type.hpp"
@@ -87,13 +87,12 @@ void jit_gemm_emitter::emit_impl(const std::vector<size_t>& in, const std::vecto
     const auto& mem_ptrs = utils::transform_idxs_to_regs(mem_ptrs_idxs);
 
     // Apply memory offsets to input/output pointers
-    h->sub(h->sp, h->sp, 3 * get_vec_length());
-
+    const auto sp_size = rnd_up(3 * get_gpr_length(), sp_alignment);
+    h->sub(h->sp, h->sp, sp_size);
     // Apply offsets and store pointers on stack
     for (size_t i = 0; i < mem_ptrs.size(); i++) {
         const auto& ptr_reg = mem_ptrs[i];
-        int32_t stack_offset = i * get_vec_length();
-
+        int32_t stack_offset = i * get_gpr_length();
         if (ov::snippets::utils::is_dynamic_value(m_memory_offsets[i])) {
             // Dynamic offset: read from runtime parameters
             size_t runtime_offset = GET_OFF(buffer_offsets) + m_buffer_ids[i] * sizeof(size_t);
@@ -103,14 +102,13 @@ void jit_gemm_emitter::emit_impl(const std::vector<size_t>& in, const std::vecto
             utils::push_ptr_with_static_offset_on_stack(h, stack_offset, ptr_reg, m_memory_offsets[i]);
         }
     }
-
     // Load back the adjusted pointers for function call
     h->ldr(x1, Xbyak_aarch64::ptr(h->sp));                        // matrix A (in0)
-    h->ldr(x2, Xbyak_aarch64::ptr(h->sp, get_vec_length()));      // matrix B (in1)
-    h->ldr(x3, Xbyak_aarch64::ptr(h->sp, 2 * get_vec_length()));  // matrix C (out)
+    h->ldr(x2, Xbyak_aarch64::ptr(h->sp, get_gpr_length()));      // matrix B (in1)
+    h->ldr(x3, Xbyak_aarch64::ptr(h->sp, 2 * get_gpr_length()));  // matrix C (out)
 
     // Restore stack pointer
-    h->add(h->sp, h->sp, 3 * get_vec_length());
+    h->add(h->sp, h->sp, sp_size);
 
     // Set up executor pointer as first argument
     const auto& compiled_kernel = get_compiled_kernel_ptr();

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

@@ -27,7 +27,6 @@ using namespace dnnl::impl::cpu::aarch64;
 
 namespace ov::intel_cpu::aarch64::utils {
 
-
 Xbyak_aarch64::XReg get_aux_gpr(const std::vector<size_t>& used_gpr_idxs) {
     // SP - stack pointer should be preserved, X0 and X1 - runtime parameter registers in the kernel
     // X18 - platform register should not be used
@@ -74,16 +73,11 @@ void push_ptr_with_runtime_offset_on_stack(dnnl::impl::cpu::aarch64::jit_generat
     // Load the runtime offset from abi_param1 (X0) and add it to the pointer
     Xbyak_aarch64::XReg abi_param1(0);
     Xbyak_aarch64::XReg offset_reg(4);
+    Xbyak_aarch64::XReg temp_reg(h->X_TMP_0);
 
-    // Handle large runtime offsets by using a temporary register
-    if (runtime_offset > 4095) {
-        Xbyak_aarch64::XReg temp_offset_reg(6);
-        h->mov(temp_offset_reg, static_cast<uint64_t>(runtime_offset));
-        h->add(temp_offset_reg, abi_param1, temp_offset_reg);
-        h->ldr(offset_reg, Xbyak_aarch64::ptr(temp_offset_reg));
-    } else {
-        h->ldr(offset_reg, Xbyak_aarch64::ptr(abi_param1, static_cast<int32_t>(runtime_offset)));
-    }
+    // Load the offset value from the runtime parameter location
+    h->add_imm(temp_reg, abi_param1, runtime_offset, Xbyak_aarch64::XReg(h->X_TMP_1));
+    h->ldr(offset_reg, Xbyak_aarch64::ptr(temp_reg));
 
     h->add(aux_reg, aux_reg, offset_reg);
 
@@ -102,17 +96,9 @@ void push_ptr_with_static_offset_on_stack(dnnl::impl::cpu::aarch64::jit_generato
     }
 
     // For non-zero offsets, apply the offset and then store
-    Xbyak_aarch64::XReg temp_reg(4);
+    Xbyak_aarch64::XReg temp_reg(h->X_TMP_0);
     h->mov(temp_reg, ptr_reg);
-
-    // For large offsets, use a register to hold the offset value
-    if (ptr_offset > 4095) {  // 12-bit immediate limit for add instruction
-        Xbyak_aarch64::XReg offset_reg(6);
-        h->mov(offset_reg, static_cast<uint64_t>(ptr_offset));
-        h->add(temp_reg, temp_reg, offset_reg);
-    } else {
-        h->add(temp_reg, temp_reg, static_cast<int32_t>(ptr_offset));
-    }
+    h->add_imm(temp_reg, temp_reg, ptr_offset, Xbyak_aarch64::XReg(h->X_TMP_1));
 
     // Store the adjusted pointer on stack
     h->str(temp_reg, Xbyak_aarch64::ptr(h->sp, stack_offset));

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

@@ -61,6 +61,7 @@ void push_ptr_with_runtime_offset_on_stack(dnnl::impl::cpu::aarch64::jit_generat
 
 /**
  * @brief Push data pointer on stack adding static offset `ptr_offset`
+ * Note: This helper doesn't allocate stack space - the user should guarantee allocated space on stack
  * @param h generator
  * @param stack_offset stack offset
  * @param ptr_reg register contains data pointer