Address Alexandra's review comments 3

aobolensk · aobolensk · commit e7031da87b93 · 2025-08-07T10:34:43.000+02:00
diff --git a/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.cpp b/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_emitter.cpp
@@ -272,38 +272,40 @@ void jit_emitter::store_context(const std::vector<size_t>& gpr_regs,
         }
     }
 
-    // 2. SIMD and Floating-Point registers
-    // 2.1. store pair registers
-    int prev_reg_idx = -1;
-    size_t ignore_registers_count = 0;
-    for (const auto reg_idx : vec_regs) {
-        if (ignore_vec_regs.find(reg_idx) != ignore_vec_regs.end()) {
-            ignore_registers_count++;
-            continue;
+    // 2. SIMD and Floating-Point registers - optimized to allocate stack space once
+    const auto store_vec_regs_size = vec_regs.size() - ignore_vec_regs.size();
+    if (store_vec_regs_size > 0) {
+        // Calculate total stack space needed for all vector registers (align once)
+        const auto total_vec_shift = ov::intel_cpu::rnd_up(get_vec_length() * store_vec_regs_size, sp_alignment);
+
+        // Single stack allocation for all vector registers
+        h->sub(h->sp, h->sp, total_vec_shift);
+
+        // Store vector registers using stack offset (preserving original order)
+        const auto last = store_vec_regs_size % 2;
+        int32_t current_offset = 0;
+
+        // Collect non-ignored registers
+        std::vector<size_t> active_regs;
+        for (const auto reg_idx : vec_regs) {
+            if (ignore_vec_regs.find(reg_idx) == ignore_vec_regs.end()) {
+                active_regs.push_back(reg_idx);
+            }
         }
-        if (prev_reg_idx == -1) {
-            prev_reg_idx = static_cast<int>(reg_idx);
-            continue;
+
+        // Store pairs
+        for (size_t i = 0; i < (active_regs.size() - last); i += 2) {
+            h->stp(Xbyak_aarch64::QReg(active_regs[i]),
+                   Xbyak_aarch64::QReg(active_regs[i + 1]),
+                   Xbyak_aarch64::ptr(h->sp, current_offset));
+            current_offset += static_cast<int32_t>(get_vec_length() * 2);
         }
-        const auto shift = ov::intel_cpu::rnd_up(get_vec_length() * 2, sp_alignment);
-        h->stp(Xbyak_aarch64::QReg(prev_reg_idx),
-               Xbyak_aarch64::QReg(reg_idx),
-               pre_ptr(h->sp, -static_cast<int32_t>(shift)));
-        prev_reg_idx = -1;
-    }
 
-    // 2.1. store the remaining register
-    if (prev_reg_idx != -1) {
-        if (ignore_vec_regs.find(prev_reg_idx) == ignore_vec_regs.end()) {
-            const auto shift = ov::intel_cpu::rnd_up(get_vec_length(), sp_alignment);
-            h->str(Xbyak_aarch64::QReg(prev_reg_idx), pre_ptr(h->sp, -static_cast<int32_t>(shift)));
-        } else {
-            ignore_registers_count++;
+        // Store the remaining register
+        if (last != 0) {
+            h->str(Xbyak_aarch64::QReg(active_regs[active_regs.size() - 1]), Xbyak_aarch64::ptr(h->sp, current_offset));
         }
     }
-
-    OPENVINO_ASSERT(ignore_registers_count == ignore_vec_regs.size(),
-                    "ignored registers size is not equal actual ignored registers count");
 }
 
 void jit_emitter::restore_context(const std::unordered_set<size_t>& ignore_vec_regs) const {
@@ -313,42 +315,44 @@ void jit_emitter::restore_context(const std::unordered_set<size_t>& ignore_vec_r
 void jit_emitter::restore_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 {
-    // 1. SIMD and Floating-Point registers
-    // 1.1. restore the remaining register
-    auto v_last = (vec_regs.size() - ignore_vec_regs.size()) % 2;
-    if (v_last != 0) {
-        for (size_t i = 0; i < vec_regs.size(); i++) {
-            const auto reg_idx = vec_regs.size() - 1 - i;
-            if (ignore_vec_regs.find(reg_idx) != ignore_vec_regs.end()) {
-                v_last++;
-                continue;
+    // 1. SIMD and Floating-Point registers - optimized to deallocate stack space once
+    const auto save_vec_regs_size = vec_regs.size() - ignore_vec_regs.size();
+    if (save_vec_regs_size > 0) {
+        // Restore vector registers using stack offset (reverse order to match original behavior)
+        const auto last = save_vec_regs_size % 2;
+        if (last != 0) {
+            int32_t current_offset = get_vec_length() * save_vec_regs_size - get_vec_length();
+            // Find the last non-ignored register
+            for (size_t i = 0; i < vec_regs.size(); i++) {
+                const auto reg_idx = vec_regs.size() - 1 - i;
+                if (ignore_vec_regs.find(reg_idx) != ignore_vec_regs.end()) {
+                    continue;
+                }
+                h->ldr(Xbyak_aarch64::QReg(reg_idx), Xbyak_aarch64::ptr(h->sp, current_offset));
+                break;
             }
-
-            const auto shift = ov::intel_cpu::rnd_up(get_vec_length(), sp_alignment);
-            h->ldr(Xbyak_aarch64::QReg(reg_idx), post_ptr(h->sp, shift));
-            break;
         }
-    }
-    // 1.2. restore pair registers
-    size_t ignore_registers_count = 0;
-    int prev_reg_idx = -1;
-    for (size_t i = v_last; i < vec_regs.size(); i++) {
-        const auto reg_idx = vec_regs.size() - 1 - i;
-        if (ignore_vec_regs.find(reg_idx) != ignore_vec_regs.end()) {
-            ignore_registers_count++;
-            continue;
+
+        // Collect non-ignored registers
+        std::vector<size_t> active_regs;
+        for (const auto reg_idx : vec_regs) {
+            if (ignore_vec_regs.find(reg_idx) == ignore_vec_regs.end()) {
+                active_regs.push_back(reg_idx);
+            }
         }
-        if (prev_reg_idx == -1) {
-            prev_reg_idx = static_cast<int>(reg_idx);
-            continue;
+
+        // Restore pairs in reverse order
+        for (size_t i = last; i < active_regs.size(); i += 2) {
+            int32_t current_offset = get_vec_length() * (active_regs.size() - (i + 2));
+            h->ldp(Xbyak_aarch64::QReg(active_regs[active_regs.size() - 1 - (i + 1)]),
+                   Xbyak_aarch64::QReg(active_regs[active_regs.size() - 1 - i]),
+                   Xbyak_aarch64::ptr(h->sp, current_offset));
         }
-        const auto shift = ov::intel_cpu::rnd_up(get_vec_length() * 2, sp_alignment);
-        h->ldp(Xbyak_aarch64::QReg(reg_idx), Xbyak_aarch64::QReg(prev_reg_idx), post_ptr(h->sp, shift));
-        prev_reg_idx = -1;
-    }
 
-    OPENVINO_ASSERT(ignore_registers_count == ignore_vec_regs.size(),
-                    "ignored registers size is not equal actual ignored registers count");
+        const auto total_vec_shift = ov::intel_cpu::rnd_up(get_vec_length() * save_vec_regs_size, sp_alignment);
+        // Single stack deallocation for all vector registers
+        h->add(h->sp, h->sp, total_vec_shift);
+    }
 
     // 2. General-purpose Registers - optimized to deallocate stack space once
     const auto save_gpr_regs_size = gpr_regs.size();
diff --git a/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp b/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_load_store_emitters.cpp
@@ -29,19 +29,18 @@ 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, SReg>) {
+        max_offset = 16380;
+        alignment = 4;
     } else if constexpr (std::is_same_v<RegType, HReg>) {
         max_offset = 8190;
         alignment = 2;
diff --git a/src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp b/src/plugins/intel_cpu/src/emitters/snippets/aarch64/jit_gemm_emitter.cpp
@@ -79,8 +79,6 @@ void jit_gemm_emitter::emit_impl(const std::vector<size_t>& in, const std::vecto
 
 void jit_gemm_emitter::emit_call(const std::vector<size_t>& mem_ptrs_idxs) const {
     const auto& call_address_reg = get_call_address_reg();
-    const auto& callee_saved_reg = get_callee_saved_reg();
-
     std::unordered_set<size_t> exclude_spill = {};
     store_context(exclude_spill);
 
@@ -96,7 +94,7 @@ 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 with auxiliary registers
-    std::vector<size_t> used_gpr_idxs = {call_address_reg.getIdx(), callee_saved_reg.getIdx()};
+    std::vector<size_t> used_gpr_idxs = {call_address_reg.getIdx()};
     for (const auto& reg : mem_ptrs) {
         used_gpr_idxs.push_back(reg.getIdx());
     }
diff --git a/src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.cpp b/src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.cpp
@@ -82,7 +82,6 @@ void push_ptr_with_runtime_offset_on_stack(dnnl::impl::cpu::aarch64::jit_generat
                                            const Xbyak_aarch64::XReg& ptr_reg,
                                            const std::vector<Xbyak_aarch64::XReg>& aux_regs,
                                            size_t runtime_offset) {
-    // Safety assertions as suggested
     OV_CPU_JIT_EMITTER_ASSERT(aux_regs.size() >= 3, "aux_regs must contain at least 3 registers");
 
     // Assert that ptr_reg is not in aux_regs
@@ -154,21 +153,16 @@ void push_and_load_ptrs_with_offsets(dnnl::impl::cpu::aarch64::jit_generator* h,
     const auto sp_size = rnd_up(mem_ptrs.size() * gpr_length, sp_alignment);
     h->sub(h->sp, h->sp, sp_size);
 
-    // Push all pointers with offsets onto stack
+    // Generate stack offsets for sequential storage
+    std::vector<size_t> stack_offsets;
+    stack_offsets.reserve(mem_ptrs.size());
     for (size_t i = 0; i < mem_ptrs.size(); i++) {
-        const auto& ptr_reg = mem_ptrs[i];
-        int32_t stack_offset = i * gpr_length;
-
-        if (ov::snippets::utils::is_dynamic_value(memory_offsets[i])) {
-            // 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 {
-            // Static offset: add compile-time constant
-            push_ptr_with_static_offset_on_stack(h, stack_offset, ptr_reg, aux_regs, memory_offsets[i]);
-        }
+        stack_offsets.push_back(i * gpr_length);
     }
 
+    // Use the common function to push pointers with offsets to stack
+    push_ptrs_with_offsets_to_stack(h, mem_ptrs, memory_offsets, buffer_ids, aux_regs, stack_offsets);
+
     // Load back the adjusted pointers to specified registers
     for (size_t i = 0; i < load_regs.size() && i < mem_ptrs.size(); i++) {
         h->ldr(load_regs[i], Xbyak_aarch64::ptr(h->sp, static_cast<int32_t>(i * gpr_length)));
@@ -187,22 +181,17 @@ void push_ptrs_with_offsets_to_stack(dnnl::impl::cpu::aarch64::jit_generator* h,
     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];
         auto 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);
-            }
+        if (ov::snippets::utils::is_dynamic_value(memory_offsets[i])) {
+            OPENVINO_ASSERT(!ov::snippets::utils::is_dynamic_value(buffer_ids[i]), "In dynamic case Buffer ID must be defined");
+            // 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 {
             // Static offset: add compile-time constant
             size_t offset = (i < memory_offsets.size()) ? memory_offsets[i] : 0;
diff --git a/src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.hpp b/src/plugins/intel_cpu/src/emitters/snippets/aarch64/utils.hpp
@@ -85,6 +85,7 @@ void push_ptr_with_static_offset_on_stack(dnnl::impl::cpu::aarch64::jit_generato
 
 /**
  * @brief Push multiple data pointers on stack with offsets and load them back to specified registers
+ * Note: This helper doesn't allocate stack space - the user should guarantee allocated space on stack
  * @param h generator
  * @param mem_ptrs vector of registers containing data pointers
  * @param memory_offsets vector of memory offsets (can be dynamic or static)
@@ -103,6 +104,7 @@ void push_and_load_ptrs_with_offsets(dnnl::impl::cpu::aarch64::jit_generator* h,
  * @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.
+ * Note: This helper doesn't allocate stack space - the user should guarantee allocated space on stack
  * @param h generator
  * @param mem_ptrs vector of registers containing data pointers
  * @param memory_offsets vector of memory offsets (can be dynamic or static)
