From 1c03bc601e9888a6aa0490e8435f4d027ed0546b Mon Sep 17 00:00:00 2001
From: Alexandra Sidorova <alexandra.sidorova@intel.com>
Date: Fri, 23 Feb 2024 14:35:56 +0400
Subject: [PATCH] [CPU] Updated aux GPR count for Load/Store Emitters

---
 .../plugin/x64/jit_load_store_emitters.cpp    |  2 +-
 .../intel_cpu/tests/unit/jit_kernel_test.cpp  | 35 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/plugins/intel_cpu/src/emitters/plugin/x64/jit_load_store_emitters.cpp b/src/plugins/intel_cpu/src/emitters/plugin/x64/jit_load_store_emitters.cpp
index 283396a0818c00..e25120515988c9 100644
--- a/src/plugins/intel_cpu/src/emitters/plugin/x64/jit_load_store_emitters.cpp
+++ b/src/plugins/intel_cpu/src/emitters/plugin/x64/jit_load_store_emitters.cpp
@@ -62,7 +62,7 @@ size_t store_emitter_params::hash() const {
 }
 
 static int get_aux_regs_as_temp(const size_t byte_size, const bool is_fill = false) {
-    if (!one_of(byte_size, 64u, 32u, 16u))
+    if (one_of(byte_size % 16, 1u, 2u, 3u))
         return 1;
     if (mayiuse(cpu::x64::avx512_core) && is_fill)
         return 1;
diff --git a/src/plugins/intel_cpu/tests/unit/jit_kernel_test.cpp b/src/plugins/intel_cpu/tests/unit/jit_kernel_test.cpp
index 608ca1b86b041f..a4de792a305227 100644
--- a/src/plugins/intel_cpu/tests/unit/jit_kernel_test.cpp
+++ b/src/plugins/intel_cpu/tests/unit/jit_kernel_test.cpp
@@ -324,13 +324,23 @@ TEST(JitKernel, variable_load_and_store) {
     {
         jit_variable_load_store_test_kernel<uint8_t, float> kernel;
         if (mayiuse(cpu_isa_t::avx512_core)) {
-            kernel.test<16, 11, false>();
+            kernel.test<16, 16, false>();
+            kernel.test<16, 15, false>();
+            kernel.test<16, 10, false>();
+            kernel.test<16, 1, false>();
         }
         if (mayiuse(cpu_isa_t::avx2)) {
+            kernel.test<8, 8, false>();
+            kernel.test<8, 7, false>();
+            kernel.test<8, 6, false>();
             kernel.test<8, 5, false>();
+            kernel.test<8, 4, false>();
         }
         if (mayiuse(cpu_isa_t::sse41)) {
+            kernel.test<4, 4, false>();
             kernel.test<4, 3, false>();
+            kernel.test<4, 2, false>();
+            kernel.test<4, 1, false>();
         }
     }
 
@@ -360,6 +370,29 @@ TEST(JitKernel, variable_load_and_store) {
         }
     }
 
+    {
+        jit_variable_load_store_test_kernel<float, uint8_t> kernel;
+        if (mayiuse(cpu_isa_t::avx512_core)) {
+            kernel.test<16, 16, true>();
+            kernel.test<16, 10, true>();
+            kernel.test<16, 2, true>();
+            kernel.test<16, 1, true>();
+        }
+        if (mayiuse(cpu_isa_t::avx2)) {
+            kernel.test<8, 8, true>();
+            kernel.test<8, 7, true>();
+            kernel.test<8, 6, true>();
+            kernel.test<8, 5, true>();
+            kernel.test<8, 4, true>();
+        }
+        if (mayiuse(cpu_isa_t::sse41)) {
+            kernel.test<4, 4, true>();
+            kernel.test<4, 3, true>();
+            kernel.test<4, 2, true>();
+            kernel.test<4, 1, true>();
+        }
+    }
+
     {
         jit_variable_load_store_test_kernel<int32_t, bfloat16_t> kernel;
         if (mayiuse(cpu_isa_t::avx512_core)) {