[FIX] [1D] Enlarge support

Author: alexey-varyzgin

src/cpu/reorder/simple_reorder.hpp

@@ -1322,8 +1322,10 @@ struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
 
 template <SIMPLE_REORDER_TEMPL_DECL>
 struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
-        typename utils::enable_if<(tag_i == format_tag::nchw
-                                          && tag_o == format_tag::nChw16c)
+        typename utils::enable_if<((tag_i == format_tag::nchw
+                                          && tag_o == format_tag::nChw16c) ||
+                                    (tag_i == format_tag::ncw
+                                          && tag_o == format_tag::nCw16c))
                 && type_i == data_type::f32
                 && type_o == data_type::bf16>::type> {
     static bool is_applicable(const memory_desc_wrapper &input_d,
@@ -1339,23 +1341,25 @@ struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
 
     static size_t get_scratchpad_size(const memory_desc_wrapper &input_d,
             const memory_desc_wrapper &output_d) {
+        constexpr int ndims = tag_traits<tag_i>::ndims;
         const size_t blksize = 16;
-        const size_t W = input_d.dims()[3];
+        const size_t W = input_d.dims()[ndims - 1];
         return sizeof(float) * blksize * W * dnnl_get_max_threads();
     }
 
     static status_t execute(const cpu_reorder_pd_t *pd, const exec_ctx_t &ctx) {
         DECLARE_COMMON_PARAMS();
 
         const dim_t blksize = 16;
+        const dim_t ndims = tag_traits<tag_i>::ndims;
 
         const auto &flat_d = input_d;
         const auto &dims = input_d.dims();
         const auto &pdims = output_d.padded_dims();
 
         const dim_t C = dims[1];
-        const dim_t H = dims[2];
-        const dim_t W = dims[3];
+        const dim_t H = ndims == 3 ? 1 : dims[ndims - 2];
+        const dim_t W = dims[ndims - 1];
 
         const dim_t wsp_size = W * blksize;
         float *wspace = scratchpad.template get<float>(
@@ -1368,7 +1372,7 @@ struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
                 for (c = 0; c < curr_c_block; ++c) {
                     const ptrdiff_t flat_off = 0
                             + c * flat_d.blocking_desc().strides[1]
-                            + w * flat_d.blocking_desc().strides[3];
+                            + w * flat_d.blocking_desc().strides[ndims - 1];
                     o[w * blksize + c] = i[flat_off];
                 }
                 for (/* continue */; c < c_block; ++c) {

src/cpu/x64/jit_avx512_core_x8s8s32x_conv_kernel.cpp

@@ -1491,7 +1491,7 @@ status_t jit_avx512_core_x8s8s32x_fwd_kernel::init_conf(jit_conv_conf_t &jcp,
             return status::unimplemented;
     }
 
-    if (jcp.with_input_zp && jcp.is_depthwise && ndims != 4)
+    if (jcp.with_input_zp && jcp.is_depthwise && ndims == 5)
         return status::unimplemented;
 
     if (jcp.with_weights_zp)

src/cpu/x64/jit_avx512_core_x8s8s32x_convolution.cpp

@@ -97,9 +97,9 @@ status_t jit_avx512_core_x8s8s32x_convolution_fwd_t::execute_forward_1d(
     size_t ch_offset = jcp.is_depthwise ? jcp.nb_ch * jcp.ch_block
                                         : jcp.ngroups * jcp.oc;
     auto w = const_cast<char *>(weights);
-    int32_t *compensation = (jcp.signed_input)
-            ? reinterpret_cast<int32_t *>(&w[extra_data_offset])
-            : nullptr;
+    int32_t *compensation = (jcp.signed_input) ? reinterpret_cast<int32_t *>(&w[extra_data_offset]) :
+                            (jcp.with_input_zp) ? pd()->attr()->output_compensations_.shifts_ : nullptr;
+    const uint8_t* input_zp = pd()->attr()->input_zero_points_.shifts_;
     int32_t *zp_compensation = jcp.src_zero_point
             ? reinterpret_cast<int32_t *>(&w[extra_data_offset])
                     + (jcp.signed_input ? ch_offset : 0)
@@ -147,7 +147,7 @@ status_t jit_avx512_core_x8s8s32x_convolution_fwd_t::execute_forward_1d(
 
             p.bias = bias ? bias + (bias_d.blk_off(g_oc) * bia_dt_size)
                           : nullptr;
-            p.compensation = (jcp.signed_input) ? compensation + g_oc : nullptr;
+            p.compensation = (jcp.signed_input || jcp.with_input_zp) ? compensation + g_oc : nullptr;
             p.zp_compensation
                     = jcp.src_zero_point ? zp_compensation + g_oc : nullptr;
             p.src_zero_point = jcp.src_zero_point ? src_zero_point : nullptr;
@@ -166,6 +166,8 @@ status_t jit_avx512_core_x8s8s32x_convolution_fwd_t::execute_forward_1d(
             p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
             p.dst_orig = dst;
             p.oc_off = g_oc * sizeof(float);
+            if (jcp.with_input_zp)
+                p.input_zp = input_zp + g_ic;
 
             (*kernel_)(&p);
 

src/cpu/x64/jit_uni_fork_dw_conv_kernel_utils.hpp

@@ -104,14 +104,14 @@ status_t jit_uni_fork_dw_conv_fwd_kernel<isa, kernel_dt>::init_conf(
     const memory_desc_wrapper bias_d(&bias_md);
 
     const int ndims = src_d.ndims();
-    // Currently this kernel only supports 2D and 3D convolutions.
-    if (ndims != 4 && ndims != 5) return status::unimplemented;
 
-    const auto blocked_tag = (ndims == 5) ? one_of(isa, avx512_core, avx512_core) ? nCdhw16c : nCdhw8c
-                                          : one_of(isa, avx512_core, avx512_core) ? nChw16c : nChw8c;
-    const auto wei_tag = (ndims == 5) ? one_of(isa, avx512_core, avx512_core) ? Goidhw16g : Goidhw8g
-                                      : one_of(isa, avx512_core, avx512_core) ? Goihw16g : Goihw8g;
-    const auto nxc_tag = (ndims == 5) ? ndhwc : nhwc;
+    const auto blocked_tag = one_of(isa, avx512_core) ?
+                             pick(ndims - 3, nCw16c, nChw16c, nCdhw16c) :
+                             pick(ndims - 3, nCw8c, nChw8c, nCdhw8c);
+    const auto wei_tag = one_of(isa, avx512_core) ?
+                         pick(ndims - 3, Goiw16g, Goihw16g, Goidhw16g) :
+                         pick(ndims - 3, Goiw8g, Goihw8g, Goidhw8g);
+    const auto nxc_tag = pick(ndims - 3, nwc, nhwc, ndhwc);
 
     jcp.with_bias = cd.bias_desc.format_kind != format_kind::undef;
 
@@ -172,29 +172,29 @@ status_t jit_uni_fork_dw_conv_fwd_kernel<isa, kernel_dt>::init_conf(
     jcp.ic = src_d.dims()[1];
 
     jcp.id = (ndims == 5) ? src_d.dims()[2] : 1;
-    jcp.ih = src_d.dims()[ndims - 2];
+    jcp.ih = (ndims == 3) ? 1 : src_d.dims()[ndims - 2];
     jcp.iw = src_d.dims()[ndims - 1];
     jcp.od = (ndims == 5) ? dst_d.dims()[2] : 1;
-    jcp.oh = dst_d.dims()[ndims - 2];
+    jcp.oh = (ndims == 3) ? 1 : dst_d.dims()[ndims - 2];
     jcp.ow = dst_d.dims()[ndims - 1];
 
     jcp.kd = (ndims == 5) ? weights_d.dims()[3] : 1;
-    jcp.kh = weights_d.dims()[ndims - 1];
+    jcp.kh = (ndims == 3) ? 1 : weights_d.dims()[ndims - 1];
     jcp.kw = weights_d.dims()[ndims];
 
     jcp.f_pad = (ndims == 5) ? cd.padding[0][0] : 0;
-    jcp.t_pad = cd.padding[0][ndims - 4];
+    jcp.t_pad = (ndims == 3) ? 0 : cd.padding[0][ndims - 4];
     jcp.l_pad = cd.padding[0][ndims - 3];
     jcp.back_pad = (ndims == 5) ? cd.padding[1][0] : 0;
-    jcp.b_pad = cd.padding[1][ndims - 4];
+    jcp.b_pad = (ndims == 3) ? 0 : cd.padding[1][ndims - 4];
     jcp.r_pad = cd.padding[1][ndims - 3];
 
     jcp.stride_d = (ndims == 5) ? cd.strides[0] : 1;
-    jcp.stride_h = cd.strides[ndims - 4];
+    jcp.stride_h = (ndims == 3) ? 1 : cd.strides[ndims - 4];
     jcp.stride_w = cd.strides[ndims - 3];
 
     jcp.dilate_d = (ndims == 5) ? cd.dilates[0] : 0;
-    jcp.dilate_h = cd.dilates[ndims - 4];
+    jcp.dilate_h = (ndims == 3) ? 0 : cd.dilates[ndims - 4];
     jcp.dilate_w = cd.dilates[ndims - 3];
 
     jcp.loop_order = loop_ngcw;

src/cpu/x64/jit_uni_fork_dw_convolution.cpp

@@ -95,9 +95,12 @@ void jit_uni_fork_dw_convolution_fwd_t<isa, src_type, dst_type>::execute_forward
         const auto ic_off_idx = is_src_layout_nxc ? ch * jcp.ch_block : ch;
         const auto oc_off_idx = is_dst_layout_nxc ? ch * jcp.ch_block : ch;
 
-        size_t src_off = (jcp.ndims == 5) ? src_d.blk_off(n, ic_off_idx, id, ih, iw) : src_d.blk_off(n, ic_off_idx, ih, iw);
-        size_t dst_off = (jcp.ndims == 5) ? dst_d.blk_off(n, oc_off_idx, od, oh, ow) : dst_d.blk_off(n, oc_off_idx, oh, ow);
-        size_t wei_off = (jcp.ndims == 5) ? weights_d.blk_off(ch, 0, 0, kd, kh, kw) : weights_d.blk_off(ch, 0, 0, kh, kw);
+        size_t src_off = (jcp.ndims == 3) ? src_d.blk_off(n, ic_off_idx, iw) :
+                         (jcp.ndims == 4) ? src_d.blk_off(n, ic_off_idx, ih, iw) : src_d.blk_off(n, ic_off_idx, id, ih, iw);
+        size_t dst_off = (jcp.ndims == 3) ? dst_d.blk_off(n, oc_off_idx, ow) :
+                         (jcp.ndims == 4) ? dst_d.blk_off(n, oc_off_idx, oh, ow) : dst_d.blk_off(n, oc_off_idx, od, oh, ow);
+        size_t wei_off = (jcp.ndims == 3) ? weights_d.blk_off(ch, 0, 0, kw) :
+                         (jcp.ndims == 4) ? weights_d.blk_off(ch, 0, 0, kh, kw) : weights_d.blk_off(ch, 0, 0, kd, kh, kw);
 
         par_conv.src = &src[src_off];
         par_conv.dst = &dst[dst_off];

src/cpu/x64/jit_uni_x8s8s32x_conv_kernel.cpp

@@ -1367,7 +1367,7 @@ status_t jit_uni_x8s8s32x_fwd_kernel<isa>::init_conf(jit_conv_conf_t &jcp,
             return status::unimplemented;
     }
 
-    if (jcp.with_input_zp && jcp.is_depthwise && ndims != 4)
+    if (jcp.with_input_zp && jcp.is_depthwise && !utils::one_of(ndims, 3, 4))
         return status::unimplemented;
 
     if (jcp.with_weights_zp)

src/cpu/x64/jit_uni_x8s8s32x_convolution.cpp

@@ -267,9 +267,9 @@ status_t jit_uni_x8s8s32x_convolution_fwd_t<isa>::execute_forward_1d(
     size_t ch_offset = jcp.is_depthwise ? jcp.nb_ch * jcp.ch_block
                                         : jcp.ngroups * jcp.oc;
     auto w = const_cast<char *>(weights);
-    const int32_t *compensation = (jcp.signed_input)
-            ? reinterpret_cast<int32_t *>(&w[extra_data_offset])
-            : nullptr;
+    const int32_t *compensation = (jcp.signed_input) ? reinterpret_cast<int32_t *>(&w[extra_data_offset]) :
+                                  (jcp.with_input_zp) ? pd()->attr()->output_compensations_.shifts_ : nullptr;
+    const uint8_t* input_zp = pd()->attr()->input_zero_points_.shifts_;
     const int32_t *zp_compensation = jcp.src_zero_point
             ? reinterpret_cast<int32_t *>(&w[extra_data_offset])
                     + (jcp.signed_input ? ch_offset : 0)
@@ -316,7 +316,7 @@ status_t jit_uni_x8s8s32x_convolution_fwd_t<isa>::execute_forward_1d(
 
             p.bias = bias ? bias + (bias_d.blk_off(g_oc) * bia_dt_size)
                           : nullptr;
-            p.compensation = (jcp.signed_input) ? compensation + g_oc : nullptr;
+            p.compensation = (jcp.signed_input || jcp.with_input_zp) ? compensation + g_oc : nullptr;
             p.zp_compensation
                     = jcp.src_zero_point ? zp_compensation + g_oc : nullptr;
             p.src_zero_point = jcp.src_zero_point ? src_zero_point : nullptr;
@@ -335,6 +335,8 @@ status_t jit_uni_x8s8s32x_convolution_fwd_t<isa>::execute_forward_1d(
             p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
             p.dst_orig = dst;
             p.oc_off = g_oc * sizeof(float);
+            if (jcp.with_input_zp)
+                p.input_zp = input_zp + g_ic;
 
             (*kernel_)(&p);