Apply Vladislav's comments

Author: aobolensk

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

@@ -260,16 +260,15 @@ void jit_emitter::store_context(const std::vector<size_t>& gpr_regs,
 
         // Store GPR registers using stack offset (preserving original order)
         const auto last = store_gpr_regs_size % 2;
-        auto current_offset = 0U;
+        int32_t current_offset = 0;
         for (size_t i = 0; i < (store_gpr_regs_size - last); i += 2) {
             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 += get_gpr_length() * 2;
+                   Xbyak_aarch64::ptr(h->sp, current_offset));
+            current_offset += static_cast<int32_t>(get_gpr_length() * 2);
         }
         if (last != 0) {
-            h->str(Xbyak_aarch64::XReg(gpr_regs[store_gpr_regs_size - 1]),
-                   Xbyak_aarch64::ptr(h->sp, static_cast<int32_t>(current_offset)));
+            h->str(Xbyak_aarch64::XReg(gpr_regs[store_gpr_regs_size - 1]), Xbyak_aarch64::ptr(h->sp, current_offset));
         }
     }
 
@@ -354,25 +353,21 @@ void jit_emitter::restore_context(const std::vector<size_t>& gpr_regs,
     // 2. General-purpose Registers - optimized to deallocate stack space once
     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 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)
         const auto last = save_gpr_regs_size % 2;
-        auto current_offset = get_gpr_length() * save_gpr_regs_size;
         if (last != 0) {
-            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)));
+            int32_t current_offset = get_gpr_length() * save_gpr_regs_size - get_gpr_length();
+            h->ldr(Xbyak_aarch64::XReg(gpr_regs[save_gpr_regs_size - 1]), Xbyak_aarch64::ptr(h->sp, current_offset));
         }
 
         for (size_t i = last; i < save_gpr_regs_size; i += 2) {
-            current_offset -= get_gpr_length() * 2;
+            int32_t current_offset = get_gpr_length() * (save_gpr_regs_size - (i + 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)));
+                   Xbyak_aarch64::ptr(h->sp, current_offset));
         }
 
+        const auto total_gpr_shift = ov::intel_cpu::rnd_up(get_gpr_length() * save_gpr_regs_size, sp_alignment);
         // Single stack deallocation for all GPR registers
         h->add(h->sp, h->sp, total_gpr_shift);
     }

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

@@ -25,43 +25,50 @@ namespace ov::intel_cpu::aarch64 {
 using jit_generator = dnnl::impl::cpu::aarch64::jit_generator;
 using cpu_isa_t = dnnl::impl::cpu::aarch64::cpu_isa_t;
 
+// Helper function to get max_offset and alignment for different register types
+template <typename RegType>
+static std::pair<int, int> get_load_store_limits() {
+    // Default fallback
+    int max_offset = 4095;
+    int alignment = 1;
+
+    if constexpr (std::is_same_v<RegType, VReg> || std::is_same_v<RegType, QReg>) {
+        max_offset = 65520;  // 4095 * 16
+        alignment = 16;
+    } else if constexpr (std::is_same_v<RegType, SReg>) {
+        max_offset = 16380;
+        alignment = 4;
+    } else if constexpr (std::is_same_v<RegType, DReg>) {
+        max_offset = 32760;
+        alignment = 8;
+    } else if constexpr (std::is_same_v<RegType, HReg>) {
+        max_offset = 8190;
+        alignment = 2;
+    } else if constexpr (std::is_same_v<RegType, BReg>) {
+        max_offset = 4095;
+        alignment = 1;
+    }
+
+    return {max_offset, alignment};
+}
+
 // Helper function to load with large offset handling
 template <typename RegType>
 static void load_with_offset_check(jit_generator* h, const RegType& dst, const XReg& src, int offset) {
-    if constexpr (std::is_same_v<RegType, VReg> || std::is_same_v<RegType, QReg>) {
-        // Manual offset handling for VReg/QReg due to uni_ldr limitations
-        const int off_mod = offset % 16;
-        const int off_mul_vl = offset / 16;
+    const auto [max_offset, alignment] = get_load_store_limits<RegType>();
 
-        if (off_mod == 0 && off_mul_vl >= 0 && off_mul_vl <= 4095) {
+    if (offset >= 0 && offset <= max_offset && (offset % alignment) == 0) {
+        if constexpr (std::is_same_v<RegType, VReg> || std::is_same_v<RegType, QReg>) {
             h->ldr(QReg(dst.getIdx()), ptr(src, static_cast<uint32_t>(offset)));
         } else {
-            h->add_imm(h->X_DEFAULT_ADDR, src, offset, h->X_TMP_0);
-            h->ldr(QReg(dst.getIdx()), ptr(h->X_DEFAULT_ADDR));
+            h->ldr(dst, ptr(src, static_cast<uint32_t>(offset)));
         }
     } 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>) {
-            max_offset = 16380;
-            alignment = 4;
-        } else if constexpr (std::is_same_v<RegType, DReg>) {
-            max_offset = 32760;
-            alignment = 8;
-        } else if constexpr (std::is_same_v<RegType, HReg>) {
-            max_offset = 8190;
-            alignment = 2;
-        } else if constexpr (std::is_same_v<RegType, BReg>) {
-            max_offset = 4095;
-            alignment = 1;
-        }
-
-        if (offset >= 0 && offset <= max_offset && (offset % alignment) == 0) {
-            h->ldr(dst, ptr(src, static_cast<uint32_t>(offset)));
+        // Use add_imm which handles register allocation internally
+        h->add_imm(h->X_DEFAULT_ADDR, src, offset, h->X_TMP_0);
+        if constexpr (std::is_same_v<RegType, VReg> || std::is_same_v<RegType, QReg>) {
+            h->ldr(QReg(dst.getIdx()), ptr(h->X_DEFAULT_ADDR));
         } else {
-            h->add_imm(h->X_DEFAULT_ADDR, src, offset, h->X_TMP_0);
             h->ldr(dst, ptr(h->X_DEFAULT_ADDR));
         }
     }
@@ -70,39 +77,20 @@ static void load_with_offset_check(jit_generator* h, const RegType& dst, const X
 // Helper function to store with large offset handling
 template <typename RegType>
 static void store_with_offset_check(jit_generator* h, const RegType& src, const XReg& dst, int offset) {
-    if constexpr (std::is_same_v<RegType, VReg> || std::is_same_v<RegType, QReg>) {
-        // Manual offset handling for VReg/QReg due to uni_str limitations
-        const int off_mod = offset % 16;
-        const int off_mul_vl = offset / 16;
+    const auto [max_offset, alignment] = get_load_store_limits<RegType>();
 
-        if (off_mod == 0 && off_mul_vl >= 0 && off_mul_vl <= 4095) {
+    if (offset >= 0 && offset <= max_offset && (offset % alignment) == 0) {
+        if constexpr (std::is_same_v<RegType, VReg> || std::is_same_v<RegType, QReg>) {
             h->str(QReg(src.getIdx()), ptr(dst, static_cast<uint32_t>(offset)));
         } else {
-            h->add_imm(h->X_DEFAULT_ADDR, dst, offset, h->X_TMP_0);
-            h->str(QReg(src.getIdx()), ptr(h->X_DEFAULT_ADDR));
+            h->str(src, ptr(dst, static_cast<uint32_t>(offset)));
         }
     } else {
-        // Manual offset handling for other register types
-        int max_offset = 4095;  // Default fallback
-        int alignment = 1;      // Default fallback
-        if constexpr (std::is_same_v<RegType, SReg>) {
-            max_offset = 16380;
-            alignment = 4;
-        } else if constexpr (std::is_same_v<RegType, DReg>) {
-            max_offset = 32760;
-            alignment = 8;
-        } else if constexpr (std::is_same_v<RegType, HReg>) {
-            max_offset = 8190;
-            alignment = 2;
-        } else if constexpr (std::is_same_v<RegType, BReg>) {
-            max_offset = 4095;
-            alignment = 1;
-        }
-
-        if (offset >= 0 && offset <= max_offset && (offset % alignment) == 0) {
-            h->str(src, ptr(dst, static_cast<uint32_t>(offset)));
+        // Use add_imm which handles register allocation internally
+        h->add_imm(h->X_DEFAULT_ADDR, dst, offset, h->X_TMP_0);
+        if constexpr (std::is_same_v<RegType, VReg> || std::is_same_v<RegType, QReg>) {
+            h->str(QReg(src.getIdx()), ptr(h->X_DEFAULT_ADDR));
         } else {
-            h->add_imm(h->X_DEFAULT_ADDR, dst, offset, h->X_TMP_0);
             h->str(src, ptr(h->X_DEFAULT_ADDR));
         }
     }

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

@@ -121,10 +121,9 @@ void jit_binary_call_emitter::init_binary_call_regs(size_t num_binary_args,
 void jit_binary_call_emitter::emit_stack_preserve(size_t stack_size) const {
     OV_CPU_JIT_EMITTER_ASSERT(!m_stack_preserved, "emit_stack_preserve called twice without emit_stack_restore");
 
-    // ARM64 requires 16-byte stack alignment
-    stack_size = ov::intel_cpu::rnd_up(stack_size, sp_alignment);
-
     if (stack_size > 0) {
+        // ARM64 requires 16-byte stack alignment
+        stack_size = ov::intel_cpu::rnd_up(stack_size, sp_alignment);
         h->sub(h->sp, h->sp, stack_size);
     }
 
@@ -134,10 +133,9 @@ void jit_binary_call_emitter::emit_stack_preserve(size_t stack_size) const {
 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
-    stack_size = ov::intel_cpu::rnd_up(stack_size, sp_alignment);
-
     if (stack_size > 0) {
+        // ARM64 requires 16-byte stack alignment
+        stack_size = ov::intel_cpu::rnd_up(stack_size, sp_alignment);
         h->add(h->sp, h->sp, stack_size);
     }
 

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

@@ -101,42 +101,19 @@ void jit_gemm_emitter::emit_call(const std::vector<size_t>& mem_ptrs_idxs) const
 
     const auto& mem_ptrs = utils::transform_idxs_to_regs(mem_ptrs_idxs);
 
-    // Collect used register indices to avoid conflicts
+    // Collect used register indices to avoid conflicts with auxiliary registers
     std::vector<size_t> used_gpr_idxs = {call_address_reg.getIdx(), callee_saved_reg.getIdx()};
-
-    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 = m_buffer_ids[i] != SIZE_MAX;
-
-        // Add current register to avoid conflicts in auxiliary register allocation
-        used_gpr_idxs.push_back(mem_ptrs[i].getIdx());
-
-        if (is_dynamic_offset && is_valid_buffer_id) {
-            OPENVINO_ASSERT(is_valid_buffer_id, "In dynamic case Buffer ID must be defined");
-            // 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 {
-            // 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]);
-        }
+    for (const auto& reg : mem_ptrs) {
+        used_gpr_idxs.push_back(reg.getIdx());
     }
 
-    h->mov(call_address_reg, 0);
-    h->str(call_address_reg, Xbyak_aarch64::ptr(h->sp, static_cast<int32_t>(reserved_stack_size)));
+    // Get auxiliary registers for the helper function (needs at least 3 for dynamic offsets)
+    auto aux_gprs = ov::intel_cpu::aarch64::utils::get_aux_gprs(used_gpr_idxs, 3);
+
+    // Use the new helper function to push all pointers with offsets to their stack locations
+    utils::push_ptrs_with_offsets_to_stack(h, mem_ptrs, m_memory_offsets, m_buffer_ids, aux_gprs, gemm_args_offsets);
+
+    // Note: scratch field was removed per earlier review feedback, so we don't need to zero it
 
     Xbyak_aarch64::XReg x0(0);
     Xbyak_aarch64::XReg x1(1);

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

@@ -86,7 +86,7 @@ void GemmKaiKernelExecutor::execute(const GemmKaiKernelExecutor* executor, const
         ukernel.run_matmul(M,
                            n_block_size,
                            K,
-                           static_cast<const void*>(args->A),
+                           args->A,
                            lhs_stride,
                            rhs_ptr,
                            dst_ptr,

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

@@ -178,4 +178,37 @@ void push_and_load_ptrs_with_offsets(dnnl::impl::cpu::aarch64::jit_generator* h,
     h->add(h->sp, h->sp, sp_size);
 }
 
+void push_ptrs_with_offsets_to_stack(dnnl::impl::cpu::aarch64::jit_generator* h,
+                                     const std::vector<Xbyak_aarch64::XReg>& mem_ptrs,
+                                     const std::vector<size_t>& memory_offsets,
+                                     const std::vector<size_t>& buffer_ids,
+                                     const std::vector<Xbyak_aarch64::XReg>& aux_regs,
+                                     const std::vector<size_t>& stack_offsets) {
+    OV_CPU_JIT_EMITTER_ASSERT(mem_ptrs.size() == memory_offsets.size(), "mem_ptrs and memory_offsets size mismatch");
+    OV_CPU_JIT_EMITTER_ASSERT(mem_ptrs.size() == buffer_ids.size(), "mem_ptrs and buffer_ids size mismatch");
+    OV_CPU_JIT_EMITTER_ASSERT(mem_ptrs.size() == stack_offsets.size(), "mem_ptrs and stack_offsets size mismatch");
+    OV_CPU_JIT_EMITTER_ASSERT(aux_regs.size() >= 3, "aux_regs must contain at least 3 registers");
+
+    // Store all pointers with offsets to their specific stack locations
+    for (size_t i = 0; i < mem_ptrs.size(); i++) {
+        const auto& ptr_reg = mem_ptrs[i];
+        int32_t stack_offset = static_cast<int32_t>(stack_offsets[i]);
+
+        if (i < memory_offsets.size() && ov::snippets::utils::is_dynamic_value(memory_offsets[i])) {
+            if (i < buffer_ids.size() && !ov::snippets::utils::is_dynamic_value(buffer_ids[i]) && buffer_ids[i] < 24) {
+                // Dynamic offset: read from runtime parameters
+                size_t runtime_offset = GET_OFF(buffer_offsets) + buffer_ids[i] * sizeof(size_t);
+                push_ptr_with_runtime_offset_on_stack(h, stack_offset, ptr_reg, aux_regs, runtime_offset);
+            } else {
+                // Invalid buffer ID, store with zero offset
+                push_ptr_with_static_offset_on_stack(h, stack_offset, ptr_reg, aux_regs, 0);
+            }
+        } else {
+            // Static offset: add compile-time constant
+            size_t offset = (i < memory_offsets.size()) ? memory_offsets[i] : 0;
+            push_ptr_with_static_offset_on_stack(h, stack_offset, ptr_reg, aux_regs, offset);
+        }
+    }
+}
+
 }  // namespace ov::intel_cpu::aarch64::utils

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

@@ -99,4 +99,22 @@ void push_and_load_ptrs_with_offsets(dnnl::impl::cpu::aarch64::jit_generator* h,
                                      const std::vector<Xbyak_aarch64::XReg>& aux_regs,
                                      const std::vector<Xbyak_aarch64::XReg>& load_regs);
 
+/**
+ * @brief Push multiple data pointers to specific stack offsets with applied memory offsets
+ * This is designed for struct-based calling conventions where pointers need to be stored
+ * at specific stack locations rather than loaded back to registers.
+ * @param h generator
+ * @param mem_ptrs vector of registers containing data pointers
+ * @param memory_offsets vector of memory offsets (can be dynamic or static)
+ * @param buffer_ids vector of buffer IDs for runtime offset calculation
+ * @param aux_regs vector of available auxiliary registers (must contain >= 3 registers, no overlap with mem_ptrs)
+ * @param stack_offsets vector of stack offsets where adjusted pointers should be stored
+ */
+void push_ptrs_with_offsets_to_stack(dnnl::impl::cpu::aarch64::jit_generator* h,
+                                     const std::vector<Xbyak_aarch64::XReg>& mem_ptrs,
+                                     const std::vector<size_t>& memory_offsets,
+                                     const std::vector<size_t>& buffer_ids,
+                                     const std::vector<Xbyak_aarch64::XReg>& aux_regs,
+                                     const std::vector<size_t>& stack_offsets);
+
 }  // namespace ov::intel_cpu::aarch64::utils