[CPU] [ARM64] emitter pipeline fix: aux vectors store/restore + gpr r…

…egisters update
openvinotoolkit · Mar 21, 2024 · af36ad5 · af36ad5
1 parent 12404fc
commit af36ad5
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 6 deletions.
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
@@ -117,12 +117,39 @@ void jit_emitter::emitter_preamble(const std::vector<size_t>& in_idxs,
         h->str(Xbyak_aarch64::XReg(preserved_gpr_idxs[i]), pre_ptr(h->sp, -16));
     }
 
+    const size_t aux_vec_idxs_size = aux_vec_idxs.size();
+    if (aux_vec_idxs_size > 1ull) {
+        for (size_t i = 0; i < (aux_vec_idxs_size - 1); i += 2) {
+            h->stp(Xbyak_aarch64::XReg(aux_vec_idxs[i]),
+                   Xbyak_aarch64::XReg(aux_vec_idxs[i + 1]),
+                   pre_ptr(h->sp, -get_vec_length() * 2));
+        }
+    }
+    if (aux_vec_idxs_size % 2) {
+        h->str(Xbyak_aarch64::XReg(aux_vec_idxs[aux_vec_idxs_size - 1]),
+               pre_ptr(h->sp, -get_vec_length()));
+    }
+
     if (!entry_map_.empty()) {
         load_table_addr();
     }
 }
 
 void jit_emitter::emitter_postamble() const {
+    const int aux_vec_idxs_size = static_cast<int>(aux_vec_idxs.size());
+    if (aux_vec_idxs_size % 2) {
+        h->ldr(Xbyak_aarch64::XReg(aux_vec_idxs[aux_vec_idxs_size - 1]),
+               post_ptr(h->sp, get_vec_length()));
+    }
+    if (aux_vec_idxs_size > 1) {
+        const int begin = aux_vec_idxs_size - ((aux_vec_idxs_size % 2) ? 2 : 1);
+        for (int i = begin; i >= 0; i -= 2) {
+            h->ldp(Xbyak_aarch64::XReg(aux_vec_idxs[i - 1]),
+                   Xbyak_aarch64::XReg(aux_vec_idxs[i]),
+                   post_ptr(h->sp, get_vec_length() * 2));
+        }
+    }
+
     const int size = static_cast<int>(preserved_gpr_idxs.size());
     for (int i = (size - 1); i >= 0; --i) {
         h->ldr(Xbyak_aarch64::XReg(preserved_gpr_idxs[i]), post_ptr(h->sp, 16));

diff --git a/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.hpp b/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.hpp
@@ -138,12 +138,12 @@ struct jit_uni_eltwise_generic : public jit_uni_eltwise_kernel, jit_generator {
     // X20    | src ptr
     // X21    | src ptr
     // X22    | src ptr
-    // X23    | src ptr
+    // X23    | temporary & kernel used (oneDNN: X_TMP_0)
     // X24    | src ptr
     // X25    | src ptr
-    // X26    | temporary
+    // X26    | src ptr
     // X27    | temporary
-    // X28    | kernel used (X_DEFAULT_ADDR)
+    // X28    | temporary & kernel used (oneDNN: X_DEFAULT_ADDR)
 
     // X29    | [not used: The Frame Pointer (FP)]
     // X30    | [not used: The Link Register (LR)]
@@ -156,14 +156,26 @@ struct jit_uni_eltwise_generic : public jit_uni_eltwise_kernel, jit_generator {
         if (idx > MAX_ELTWISE_INPUTS) {
             OPENVINO_THROW("source vector ptr register " + std::to_string(idx) + " is not supported");
         }
-        return XReg(19 + idx);
+
+        const uint32_t base = 19;
+        if ((base + idx) == 23) {
+            idx++;
+        }
+
+        return XReg(base + idx);
     }
 
     inline XReg get_aux_gpr(const uint32_t idx) {
-        if (idx > 2) {
+        if (idx > 3) {
             OPENVINO_THROW("aux gpr register " + std::to_string(idx) + " is not supported");
         }
-        return XReg(26 + idx);
+
+        if (idx == 0) {
+            return XReg(23);
+        }
+
+        const uint32_t base = 27;
+        return XReg(base + idx - 1);
     }
 
     // Vector registers mapping