Add "if" statement for loop unrolling in rms kernel. (#27215)

### Details: - *Add "if" statement for loop unrolling in rms kernel to fix Segmentation Fault in tiny-random-sd3 model* ### Tickets: - *CVS-152057*
openvinotoolkit · Oct 28, 2024 · 78864ca · 78864ca
1 parent 65dd174
commit 78864ca
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 15 deletions.
diff --git a/src/plugins/intel_cpu/src/nodes/kernels/x64/rms_kernel.cpp b/src/plugins/intel_cpu/src/nodes/kernels/x64/rms_kernel.cpp
@@ -123,22 +123,25 @@ void jit_rms_kernel<isa>::generate() {
     // x * 1/Sqrt(ReduceMean(x^2,axes)+eps) * gamma
     // sum(x^2)
     align(16);
-    Xbyak::Label loop_4reg;
-    L(loop_4reg);
-    {
-        load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false);
-        vfmadd231ps(vmm_sum0, vmm_src, vmm_src);
-        load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false, vec_size * m_jcp.src_prc.size() * 1);
-        vfmadd231ps(vmm_sum1, vmm_src, vmm_src);
-        load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false, vec_size * m_jcp.src_prc.size() * 2);
-        vfmadd231ps(vmm_sum2, vmm_src, vmm_src);
-        load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false, vec_size * m_jcp.src_prc.size() * 3);
-        vfmadd231ps(vmm_sum3, vmm_src, vmm_src);
-
-        add(reg_src, vec_size * m_jcp.src_prc.size() * 4);
-        dec(reg_size);
-        jnz(loop_4reg);
+    if ((m_jcp.data_size / (vec_size * 4)) != 0) {
+        Xbyak::Label loop_4reg;
+        L(loop_4reg);
+        {
+            load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false);
+            vfmadd231ps(vmm_sum0, vmm_src, vmm_src);
+            load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false, vec_size * m_jcp.src_prc.size() * 1);
+            vfmadd231ps(vmm_sum1, vmm_src, vmm_src);
+            load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false, vec_size * m_jcp.src_prc.size() * 2);
+            vfmadd231ps(vmm_sum2, vmm_src, vmm_src);
+            load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false, vec_size * m_jcp.src_prc.size() * 3);
+            vfmadd231ps(vmm_sum3, vmm_src, vmm_src);
+
+            add(reg_src, vec_size * m_jcp.src_prc.size() * 4);
+            dec(reg_size);
+            jnz(loop_4reg);
+        }
     }
+
     // 1 ~ 3 vmm
     for (size_t i = m_jcp.data_size / (vec_size * 4) * 4; i < m_jcp.data_size / vec_size; i++) {
         load(vmm_src, reg_src, m_jcp.src_prc, vec_size, false);

diff --git a/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/instances/x64/rms_norm.cpp b/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/instances/x64/rms_norm.cpp
@@ -24,6 +24,17 @@ const std::vector<std::vector<InputShape>> shapes{
             {ov::Shape{1024 + 16 + 1}, ov::Shape{1024 + 16 + 1}}}
         },
     },
+    // small data size
+    {
+        // data shape
+        {ov::test::InputShape{ov::PartialShape{-1, -1, 31},
+            {ov::Shape{1, 8, 31}, ov::Shape{2, 3, 31}}}
+        },
+        // scale shape
+        {ov::test::InputShape{ov::PartialShape{31},
+            {ov::Shape{31}, ov::Shape{31}}}
+        },
+    },
     // scale is scalar
     {
         // data shape