cpu: conv: Fix segfault in bf16 fwd_d due to push/pop imbalance

oneapi-src · Nov 21, 2019 · 4f6e3d5 · 4f6e3d5
1 parent 52d476c
commit 4f6e3d5
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/src/cpu/jit_avx512_core_bf16_conv_kernel.cpp b/src/cpu/jit_avx512_core_bf16_conv_kernel.cpp
@@ -213,7 +213,6 @@ void jit_avx512_core_bf16_fwd_kernel::compute_loop(
     Label skip_kh_loop;
 
     if (jcp.ndims == 5) {
-        push(reg_out);
         mov(reg_ki, ptr[param1 + GET_OFF(kd_padding)]);
         mov(aux_reg_ker_d, reg_ker);
         mov(aux_reg_inp_d, reg_inp);
@@ -269,6 +268,8 @@ void jit_avx512_core_bf16_fwd_kernel::compute_loop(
         jg(kh_label, T_NEAR);
     }
 
+    L(skip_kh_loop);
+
     if (jcp.ndims == 5) {
         add(aux_reg_inp_d,
                 jcp.typesize_in * (jcp.dilate_d + 1) * jcp.ih * jcp.iw
@@ -280,12 +281,8 @@ void jit_avx512_core_bf16_fwd_kernel::compute_loop(
         dec(reg_ki);
         cmp(reg_ki, 0);
         jg(kd_label, T_NEAR);
-
-        pop(reg_out);
     }
 
-    L(skip_kh_loop);
-
     // End of IC Loop
     size_t inp_step = (size_t)jcp.id * jcp.ih * jcp.iw * jcp.ic_block;
     size_t ker_step

diff --git a/src/cpu/jit_avx512_core_bf16_conv_kernel.hpp b/src/cpu/jit_avx512_core_bf16_conv_kernel.hpp
@@ -100,7 +100,7 @@ struct jit_avx512_core_bf16_fwd_kernel : public jit_generator {
     reg64_t reg_bias = rbx;
 
     reg64_t reg_kj = abi_not_param1;
-    reg64_t reg_ki = reg_out;
+    reg64_t reg_ki = reg_bias;
     reg64_t reg_oi = rdx;
     reg64_t reg_kh = rsi;
 

diff --git a/tests/benchdnn/inputs/conv/conv_regression_padding b/tests/benchdnn/inputs/conv/conv_regression_padding
@@ -21,3 +21,6 @@ mb1ic16oc16_ih10kh1_kh1ph1sh1
 
 # 3d with depth padding
 mb1ic1oc16_kd2kh1kw1_id2ih1iw1_od1oh1ow1_sd2sh1sw1_pd1ph0pw0
+
+# Skip compute on kh
+mb1ic16id1ih1iw1oc16od2oh2ow1kd3kh3kw1dd1dh1ph1pd1