Address Alexandra's comments

Author: aobolensk

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

@@ -252,27 +252,20 @@ void jit_emitter::store_context(const std::vector<size_t>& gpr_regs,
     // 1. General-purpose Registers - optimized to allocate stack space once
     const auto store_gpr_regs_size = gpr_regs.size();
     if (store_gpr_regs_size > 0) {
-        // Calculate total stack space needed for all GPR registers
-        const auto last = store_gpr_regs_size % 2;
-        auto total_gpr_shift = 0U;
-        for (size_t i = 0; i < (store_gpr_regs_size - last); i += 2) {
-            total_gpr_shift += ov::intel_cpu::rnd_up(get_gpr_length() * 2, sp_alignment);
-        }
-        if (last != 0) {
-            total_gpr_shift += ov::intel_cpu::rnd_up(get_gpr_length(), sp_alignment);
-        }
+        // Calculate total stack space needed for all GPR registers (align once)
+        const auto total_gpr_shift = ov::intel_cpu::rnd_up(get_gpr_length() * store_gpr_regs_size, sp_alignment);
 
         // Single stack allocation for all GPR registers
         h->sub(h->sp, h->sp, total_gpr_shift);
 
         // Store GPR registers using stack offset (preserving original order)
+        const auto last = store_gpr_regs_size % 2;
         auto current_offset = 0U;
         for (size_t i = 0; i < (store_gpr_regs_size - last); i += 2) {
-            const auto shift = ov::intel_cpu::rnd_up(get_gpr_length() * 2, sp_alignment);
             h->stp(Xbyak_aarch64::XReg(gpr_regs[i]),
                    Xbyak_aarch64::XReg(gpr_regs[i + 1]),
                    Xbyak_aarch64::ptr(h->sp, static_cast<int32_t>(current_offset)));
-            current_offset += shift;
+            current_offset += get_gpr_length() * 2;
         }
         if (last != 0) {
             h->str(Xbyak_aarch64::XReg(gpr_regs[store_gpr_regs_size - 1]),
@@ -362,27 +355,19 @@ void jit_emitter::restore_context(const std::vector<size_t>& gpr_regs,
     const auto save_gpr_regs_size = gpr_regs.size();
     if (save_gpr_regs_size > 0) {
         // Calculate total stack space (must match store_context calculation)
-        const auto last = save_gpr_regs_size % 2;
-        auto total_gpr_shift = 0U;
-        for (size_t i = 0; i < (save_gpr_regs_size - last); i += 2) {
-            total_gpr_shift += ov::intel_cpu::rnd_up(get_gpr_length() * 2, sp_alignment);
-        }
-        if (last != 0) {
-            total_gpr_shift += ov::intel_cpu::rnd_up(get_gpr_length(), sp_alignment);
-        }
+        const auto total_gpr_shift = ov::intel_cpu::rnd_up(get_gpr_length() * save_gpr_regs_size, sp_alignment);
 
         // Restore GPR registers using stack offset (reverse order to match original behavior)
-        auto current_offset = total_gpr_shift;
+        const auto last = save_gpr_regs_size % 2;
+        auto current_offset = get_gpr_length() * save_gpr_regs_size;
         if (last != 0) {
-            const auto shift = ov::intel_cpu::rnd_up(get_gpr_length(), sp_alignment);
-            current_offset -= shift;
+            current_offset -= get_gpr_length();
             h->ldr(Xbyak_aarch64::XReg(gpr_regs[save_gpr_regs_size - 1]),
                    Xbyak_aarch64::ptr(h->sp, static_cast<int32_t>(current_offset)));
         }
 
         for (size_t i = last; i < save_gpr_regs_size; i += 2) {
-            const auto shift = ov::intel_cpu::rnd_up(get_gpr_length() * 2, sp_alignment);
-            current_offset -= shift;
+            current_offset -= get_gpr_length() * 2;
             h->ldp(Xbyak_aarch64::XReg(gpr_regs[save_gpr_regs_size - 1 - (i + 1)]),
                    Xbyak_aarch64::XReg(gpr_regs[save_gpr_regs_size - 1 - i]),
                    Xbyak_aarch64::ptr(h->sp, static_cast<int32_t>(current_offset)));

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp

@@ -41,6 +41,7 @@ static void load_with_offset_check(jit_generator* h, const RegType& dst, const X
         }
     } else {
         // Manual offset handling for other register types
+        // Note: read about LDR (immediate) in the manual Arm A-profile A64 Instruction Set Architecture
         int max_offset = 4095;  // Default fallback
         int alignment = 1;      // Default fallback
         if constexpr (std::is_same_v<RegType, SReg>) {
@@ -58,7 +59,7 @@ static void load_with_offset_check(jit_generator* h, const RegType& dst, const X
         }
 
         if (offset >= 0 && offset <= max_offset && (offset % alignment) == 0) {
-            h->ldr(dst, ptr(src, offset));
+            h->ldr(dst, ptr(src, static_cast<uint32_t>(offset)));
         } else {
             h->add_imm(h->X_DEFAULT_ADDR, src, offset, h->X_TMP_0);
             h->ldr(dst, ptr(h->X_DEFAULT_ADDR));

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

@@ -70,7 +70,7 @@ void jit_binary_call_emitter::init_binary_call_regs(size_t num_binary_args,
     }
 
     // Add special registers that should not be allocated
-    std::vector<size_t> reserved_regs = {
+    static const std::vector<size_t> reserved_regs = {
         18,  // Platform register (should not be used)
         29,  // Frame pointer (FP)
         30,  // Link register (LR)
@@ -129,8 +129,7 @@ void jit_binary_call_emitter::emit_stack_restore(size_t stack_size) const {
     OV_CPU_JIT_EMITTER_ASSERT(m_stack_preserved, "emit_stack_restore called without corresponding emit_stack_preserve");
 
     // ARM64 requires 16-byte stack alignment
-    const size_t alignment = 16;
-    stack_size = ov::intel_cpu::rnd_up(stack_size, alignment);
+    stack_size = ov::intel_cpu::rnd_up(stack_size, sp_alignment);
 
     if (stack_size > 0) {
         h->add(h->sp, h->sp, stack_size);

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

@@ -23,7 +23,6 @@
 #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"
@@ -88,19 +87,12 @@ void jit_gemm_copy_b_emitter::emit_impl(const std::vector<size_t>& in, const std
 }
 
 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);
 
     const auto& mem_ptrs = utils::transform_idxs_to_regs(mem_ptrs_idxs);
 

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

@@ -89,8 +89,7 @@ void jit_gemm_emitter::emit_call(const std::vector<size_t>& mem_ptrs_idxs) const
     std::unordered_set<size_t> exclude_spill = {};
     store_context(exclude_spill);
 
-    auto reserved_stack_size = sizeof(GemmKaiKernelExecutor::call_args);
-    reserved_stack_size = ov::intel_cpu::rnd_up(reserved_stack_size, sp_alignment);
+    auto reserved_stack_size = ov::intel_cpu::rnd_up(sizeof(GemmKaiKernelExecutor::call_args), sp_alignment);
     emit_stack_preserve(reserved_stack_size);
 
     const size_t A_offset = offsetof(GemmKaiKernelExecutor::call_args, A);
@@ -103,30 +102,38 @@ void jit_gemm_emitter::emit_call(const std::vector<size_t>& mem_ptrs_idxs) const
 
     for (size_t i = 0; i < mem_ptrs.size(); i++) {
         const bool is_dynamic_offset = ov::snippets::utils::is_dynamic_value(m_memory_offsets[i]);
-        const bool is_valid_buffer_id = !ov::snippets::utils::is_dynamic_value(m_buffer_ids[i]);
+        const bool is_valid_buffer_id = m_buffer_ids[i] != SIZE_MAX;
 
-        std::vector<Xbyak_aarch64::XReg> aux_regs = {call_address_reg, callee_saved_reg, h->X_TMP_1};
+        // Collect used register indices to avoid conflicts
+        std::vector<size_t> used_gpr_idxs = {call_address_reg.getIdx(),
+                                             callee_saved_reg.getIdx(),
+                                             mem_ptrs[i].getIdx()};
 
         if (is_dynamic_offset && is_valid_buffer_id) {
-            aux_regs.emplace_back(h->X_TMP_0);
+            // Get 3 auxiliary registers for dynamic offset handling (runtime offset needs at least 3)
+            auto aux_gprs = ov::intel_cpu::aarch64::utils::get_aux_gprs(used_gpr_idxs, 3);
+            std::vector<Xbyak_aarch64::XReg> aux_regs = {aux_gprs[0], aux_gprs[1], aux_gprs[2]};
             size_t runtime_offset = GET_OFF(buffer_offsets) + m_buffer_ids[i] * sizeof(size_t);
             utils::push_ptr_with_runtime_offset_on_stack(h,
                                                          gemm_args_offsets[i],
                                                          mem_ptrs[i],
                                                          aux_regs,
                                                          runtime_offset);
         } else {
-            size_t offset = is_dynamic_offset ? 0 : m_memory_offsets[i];
-            utils::push_ptr_with_static_offset_on_stack(h, gemm_args_offsets[i], mem_ptrs[i], aux_regs, offset);
+            // Static offset case (dynamic offsets with valid buffer IDs are handled above)
+            // Get 2 auxiliary registers for static offset handling (static offset needs at least 2)
+            auto aux_gprs = ov::intel_cpu::aarch64::utils::get_aux_gprs(used_gpr_idxs, 2);
+            std::vector<Xbyak_aarch64::XReg> aux_regs = {aux_gprs[0], aux_gprs[1]};
+            utils::push_ptr_with_static_offset_on_stack(h,
+                                                        gemm_args_offsets[i],
+                                                        mem_ptrs[i],
+                                                        aux_regs,
+                                                        m_memory_offsets[i]);
         }
     }
 
-    if (mem_ptrs.size() < 4) {
-        h->mov(call_address_reg, 0);
-        h->str(
-            call_address_reg,
-            Xbyak_aarch64::ptr(h->sp, static_cast<int32_t>(ov::intel_cpu::rnd_up(reserved_stack_size, sp_alignment))));
-    }
+    h->mov(call_address_reg, 0);
+    h->str(call_address_reg, Xbyak_aarch64::ptr(h->sp, static_cast<int32_t>(reserved_stack_size)));
 
     Xbyak_aarch64::XReg x0(0);
     Xbyak_aarch64::XReg x1(1);

src/plugins/intel_cpu/src/emitters/snippets/aarch64/kernel_executors/gemm.cpp

@@ -81,8 +81,8 @@ void GemmKaiKernelExecutor::execute(const GemmKaiKernelExecutor* executor, const
         const size_t dst_offset = ukernel.get_dst_offset(0, n_start, dst_stride_row);
         // in0, in1, out is point to current block memory, based on block loop info, and shift done in loop begin and
         // end emitters(adjusted copyb loop info as repack outside block loops).
-        float* rhs_ptr = const_cast<float*>(static_cast<const float*>(args->B)) + rhs_packed_offset / sizeof(float);
-        float* dst_ptr = const_cast<float*>(static_cast<const float*>(args->C)) + dst_offset / (sizeof(float));
+        const float* rhs_ptr = static_cast<const float*>(args->B) + rhs_packed_offset / sizeof(float);
+        float* dst_ptr = static_cast<float*>(args->C) + dst_offset / (sizeof(float));
         ukernel.run_matmul(M,
                            n_block_size,
                            K,