src: cpu: gemm conv: corrected 3d problem dispatching

removed extra parallel sections for data initialization by zeros

Author: akharito

src/cpu/gemm_bf16_convolution.cpp

@@ -306,24 +306,27 @@ void gemm_bf16_convolution_fwd_t<dst_data_type>::execute_forward() const {
     const size_t src_step = (size_t)jcp.ic * jcp.ih * jcp.iw * jcp.id;
     const size_t dst_step = (size_t)jcp.oc * M;
     const size_t weights_g_size = (size_t)jcp.ic * jcp.oc * jcp.ks;
+    const bool is_problem_3d = pd()->ndims() == 5;
 
     assert(IMPLICATION(
-            jcp.id != 1, jcp.oh_block == jcp.oh && jcp.ow_block == jcp.ow));
+            is_problem_3d, jcp.oh_block == jcp.oh && jcp.ow_block == jcp.ow));
     assert(IMPLICATION(jcp.ow_block != jcp.ow, jcp.oh_block == 1));
 
     const int K = jcp.ic * jcp.ks;
     const int N = jcp.oc;
 
-    if (jcp.im2col_sz && jcp.id != 1)
-        parallel_nd(jcp.im2col_sz * jcp.nthr,
-                [&](ptrdiff_t i) { col[i] = (src_data_t)0; });
-
     const int nb_oh = div_up(jcp.oh, jcp.oh_block);
     const int nb_ow = div_up(jcp.ow, jcp.ow_block);
     const size_t work_amount = (size_t)jcp.ngroups
         * jcp.mb * jcp.od * nb_oh * nb_ow;
     parallel(jcp.nthr, [&](const int ithr, const int nthr) {
         src_data_t *_col = col + (ptrdiff_t)ithr * jcp.im2col_sz;
+        if (is_problem_3d) {
+            // jit_gemm_convolution_utils::im2col_3d() requires external
+            // data initialization by zeroes
+            for (ptrdiff_t i = 0; i < jcp.im2col_sz; i++)
+                _col[i] = (src_data_t)0;
+        }
 
         int g{ 0 }, n{ 0 }, od{ 0 }, ohb{ 0 }, owb{ 0 };
         size_t start = 0, end = 0;
@@ -340,7 +343,7 @@ void gemm_bf16_convolution_fwd_t<dst_data_type>::execute_forward() const {
             const int h_step = nstl::min(jcp.oh_block, jcp.oh - oh);
             const int w_step = nstl::min(jcp.ow_block, jcp.ow - ow);
             if (jcp.im2col_sz) {
-                if (jcp.id == 1)
+                if (!is_problem_3d)
                     jit_gemm_convolution_utils::im2col<src_data_t>(
                             jcp, _src, _col, 0, jcp.os, 0, jcp.ic);
                 else
@@ -403,6 +406,7 @@ void gemm_bf16_convolution_bwd_data_t<diff_src_data_type>::
     const int LDC = jcp.im2col_sz ? m : M;
 
     const size_t work_amount = (size_t)jcp.ngroups * jcp.mb;
+    const bool is_problem_3d = pd()->ndims() == 5;
 
     parallel(jcp.nthr, [&](const int ithr, const int nthr) {
         acc_data_t *_col = col + (ptrdiff_t)ithr * jcp.im2col_sz;
@@ -419,7 +423,7 @@ void gemm_bf16_convolution_bwd_data_t<diff_src_data_type>::
                 ? acc_base + ithr * rnd_up(src_step, 16)
                 : (acc_data_t *)diff_src_local;
 
-            if (jcp.id > 1 && jcp.im2col_sz > 0) {
+            if (is_problem_3d && jcp.im2col_sz > 0) {
                 // jit_gemm_convolution_utils::col2im_3d() assumes that the
                 // accumulator is initialized by zeroes
                 for (size_t i = 0; i < src_step; i++)
@@ -437,7 +441,7 @@ void gemm_bf16_convolution_bwd_data_t<diff_src_data_type>::
                     jcp.im2col_sz ? _col: acc + od * m, &LDC);
 
                 if (jcp.im2col_sz) {
-                    if (jcp.id == 1)
+                    if (!is_problem_3d)
                         jit_gemm_convolution_utils::col2im(jcp, _col,
                             acc);
                     else
@@ -548,9 +552,7 @@ void gemm_bf16_convolution_bwd_weights_t<diff_wei_data_type>::
     const int N = jcp.oc;
     const int M = jcp.ic * jcp.ks;
     const int LDA = jcp.im2col_sz ? k : K;
-
-    parallel_nd(jcp.im2col_sz * jcp.nthr,
-            [&](ptrdiff_t i) { col[i] = (src_data_t)0; });
+    const bool is_problem_3d = pd()->ndims() == 5;
 
     parallel(jcp.nthr, [&](const int ithr, const int nthr) {
         int ithr_g, nthr_g, ithr_mb, nthr_mb;
@@ -570,6 +572,13 @@ void gemm_bf16_convolution_bwd_weights_t<diff_wei_data_type>::
             assert(IMPLICATION((g_end - g_start) > 1, need_reduction == 0));
 
             src_data_t *_col = col + (ptrdiff_t)ithr * jcp.im2col_sz;
+            if (is_problem_3d) {
+                // jit_gemm_convolution_utils::im2col_3d() requires external
+                // data initialization by zeroes
+                for (ptrdiff_t i = 0; i < jcp.im2col_sz; i++)
+                    _col[i] = (src_data_t)0;
+            }
+
             acc_data_t *weights_reduce_base = wei_reduction
                     + ithr_g * nthr_mb * weights_g_size;
             acc_data_t *weights_reduce = weights_reduce_base
@@ -586,7 +595,7 @@ void gemm_bf16_convolution_bwd_weights_t<diff_wei_data_type>::
                             + (mb * jcp.ngroups + g) * dst_step + od * k;
 
                     if (jcp.im2col_sz) {
-                        if (jcp.id == 1)
+                        if (!is_problem_3d)
                             jit_gemm_convolution_utils::im2col<src_data_t>(
                                     jcp, _src, _col, 0, jcp.os, 0, jcp.ic);
                         else

src/cpu/gemm_convolution.cpp

@@ -57,16 +57,19 @@ void gemm_convolution_fwd_t::execute_forward() const {
     const size_t src_step = jcp.ic * jcp.ih * jcp.iw * jcp.id;
     const size_t weights_oc_size = jcp.ic * jcp.ks;
     const size_t weights_g_size = weights_oc_size * jcp.oc;
+    const bool is_problem_3d = pd()->ndims() == 5;
 
     assert(IMPLICATION(
-            jcp.id != 1, jcp.os_block == jcp.os && jcp.ic_block == jcp.ic));
-
-    if (jcp.im2col_sz && jcp.id != 1)
-        parallel_nd(jcp.im2col_sz * jcp.nthr,
-                [&](ptrdiff_t i) { col[i] = (data_t)0; });
+            is_problem_3d, jcp.os_block == jcp.os && jcp.ic_block == jcp.ic));
 
     parallel(jcp.nthr, [&](const int ithr, const int nthr) {
         data_t *_col = col + (ptrdiff_t)ithr * jcp.im2col_sz;
+        if (is_problem_3d) {
+            // jit_gemm_convolution_utils::im2col_3d() requires external
+            // data initialization by zeroes
+            for (ptrdiff_t i = 0; i < jcp.im2col_sz; i++)
+                _col[i] = (data_t)0;
+        }
 
         auto inner_ker = [&](int spatial, const im_pos_t &curr, im_pos_t &prev,
                                  im_pos_t &step, const im_pos_t &end) {
@@ -82,7 +85,7 @@ void gemm_convolution_fwd_t::execute_forward() const {
             prev = curr;
 
             if (jcp.im2col_sz && do_im2col) {
-                if (jcp.id == 1)
+                if (!is_problem_3d)
                     jit_gemm_convolution_utils::im2col<float>(
                             jcp, _src, _col, curr.sp, step.sp, curr.ic, step.ic);
                 else
@@ -155,7 +158,7 @@ void gemm_convolution_fwd_t::execute_forward() const {
         im_pos_t start, end;
         end.ic = jcp.ic;
 
-        if (jcp.id == 1) {
+        if (!is_problem_3d) {
             const int sp_work = jcp.mb * jcp.ngroups * jcp.od * jcp.os;
             balance2D(nthr, ithr, sp_work, start.sp, end.sp, jcp.oc, start.oc,
                     end.oc, jcp.nthr_oc);
@@ -218,11 +221,7 @@ void gemm_convolution_bwd_data_t::execute_backward_data() const {
     const int LDC = jcp.im2col_sz ? m : M;
 
     const size_t work_amount = (size_t)jcp.ngroups * jcp.mb;
-
-    if (jcp.id > 1) {
-        const ptrdiff_t diff_src_sz = (ptrdiff_t)(work_amount * src_step);
-        parallel_nd(diff_src_sz, [&](ptrdiff_t i) { diff_src[i] = (data_t)0; });
-    }
+    const bool is_problem_3d = pd()->ndims() == 5;
 
     parallel(jcp.nthr, [&](const int ithr, const int nthr) {
         data_t *_col = col + (ptrdiff_t)ithr * jcp.im2col_sz;
@@ -234,6 +233,13 @@ void gemm_convolution_bwd_data_t::execute_backward_data() const {
         for (size_t iwork = start; iwork < end; ++iwork) {
 
             data_t *_diff_src = diff_src + (n * jcp.ngroups + g)*src_step;
+            if (is_problem_3d && jcp.im2col_sz > 0) {
+                // jit_gemm_convolution_utils::col2im_3d() assumes that the
+                // accumulator is initialized by zeroes
+                for (size_t i = 0; i < src_step; i++)
+                    _diff_src[i] = (data_t)0;
+            }
+
             const data_t *_weights = weights + g * weights_g_size;
             for (int od = 0; od < jcp.od; ++od) {
                 const data_t *_diff_dst = diff_dst + (n * jcp.ngroups + g)
@@ -245,7 +251,7 @@ void gemm_convolution_bwd_data_t::execute_backward_data() const {
                     jcp.im2col_sz ? _col:_diff_src + od * m, &LDC);
 
                 if (jcp.im2col_sz) {
-                    if (jcp.id == 1)
+                    if (!is_problem_3d)
                         jit_gemm_convolution_utils::col2im(jcp, _col,
                             _diff_src);
                     else
@@ -278,9 +284,7 @@ void gemm_convolution_bwd_weights_t::execute_backward_weights() const {
     const int N = jcp.oc;
     const int M = jcp.ic * jcp.ks;
     const int LDA = jcp.im2col_sz ? k : K;
-
-    parallel_nd(jcp.im2col_sz * jcp.nthr,
-            [&](ptrdiff_t i) { col[i] = (data_t)0; });
+    const bool is_problem_3d = pd()->ndims() == 5;
 
     parallel(jcp.nthr, [&](const int ithr, const int nthr) {
         int ithr_g, nthr_g, ithr_mb, nthr_mb;
@@ -300,6 +304,13 @@ void gemm_convolution_bwd_weights_t::execute_backward_weights() const {
             assert(IMPLICATION((g_end - g_start) > 1, need_reduction == 0));
 
             data_t *_col = col + (ptrdiff_t)ithr * jcp.im2col_sz;
+            if (is_problem_3d) {
+                // jit_gemm_convolution_utils::im2col_3d() requires external
+                // data initialization by zeroes
+                for (ptrdiff_t i = 0; i < jcp.im2col_sz; i++)
+                    _col[i] = (data_t)0;
+            }
+
             data_t *weights_reduce_base = wei_reduction
                     + ithr_g * nthr_mb * weights_g_size;
             data_t *weights_reduce = weights_reduce_base
@@ -315,7 +326,7 @@ void gemm_convolution_bwd_weights_t::execute_backward_weights() const {
                             + (mb*jcp.ngroups+g)*dst_step + od * k;
 
                     if (jcp.im2col_sz) {
-                        if (jcp.id == 1)
+                        if (!is_problem_3d)
                             jit_gemm_convolution_utils::im2col<float>(
                                     jcp, _src, _col, 0, jcp.os, 0, jcp.ic);
                         else

src/cpu/gemm_convolution_utils.cpp

@@ -652,7 +652,7 @@ status_t init_conf(jit_gemm_conv_conf_t &jcp,
         if (is_fwd) {
             const int wei_size = jcp.oc * jcp.ic * jcp.kh * jcp.kw;
             bool is_blocking_applicable = true && is_fwd && jcp.im2col_sz
-                    && jcp.id == 1 && jcp.od == 1 && jcp.dilate_h == 0
+                    && !is_3d && jcp.dilate_h == 0
                     && jcp.dilate_w == 0 && !is_depthwise && wei_size < L2 / 2;
             if (is_blocking_applicable) {
                 // looking for oh and ow blocking
@@ -808,7 +808,7 @@ status_t init_conf(jit_gemm_conv_conf_t &jcp,
             // gemm implementation which we cannot control
             bool is_blocking_applicable = true
                 && !is_bf16_conv // TODO: apply blocking to bf16
-                && jcp.id == 1 && jcp.od == 1
+                && !is_3d
                 && (!jcp.im2col_sz
                     // spatial is small
                     || spatial >= max_threads * simd_w

tests/benchdnn/inputs/test_conv_regression_general

@@ -132,3 +132,13 @@ mb1_g16ic16oc16_ih5oh2kh2sh3ph0
 
 # MFDNN-2027 2d corner case (large kernel and top padding, negative bottom padding + stride)
 --dir=BWD_W --cfg=f32 g1ic16ih7iw1oc16oh1ow1kh11kw1sh2sw1ph5pw0n
+
+# 3d problem dispatching
+--cfg=f32
+--dir=FWD_B g1ic16id1ih7iw8oc16od1oh7ow8kd3kh3kw3sd2sh1sw1pd1ph1pw1
+--dir=BWD_D g1ic16id1ih7iw8oc16od1oh7ow8kd3kh3kw3sd2sh1sw1pd1ph1pw1
+--dir=BWD_W g1ic16id1ih7iw8oc16od1oh7ow8kd3kh3kw3sd2sh1sw1pd1ph1pw1
+--cfg=bf16bf16bf16
+--dir=FWD_B g1ic16id1ih7iw8oc16od1oh7ow8kd3kh3kw3sd2sh1sw1pd1ph1pw1
+--dir=BWD_D g1ic16id1ih7iw8oc16od1oh7ow8kd3kh3kw3sd2sh1sw1pd1ph1pw1
+--dir=BWD_W g1ic16id1ih7iw8oc16od1oh7ow8kd3kh3kw3sd2sh1sw1pd1ph1pw1