onednn-src 0.1.13

/*******************************************************************************
* Copyright 2020 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/

#include "common/dnnl_thread.hpp"
#include "common/nstl.hpp"
#include "common/utils.hpp"

#include "cpu/cpu_primitive.hpp"
#include "cpu/zero_point_utils.hpp"

#include "cpu/x64/injectors/jit_uni_postops_injector.hpp"
#include "cpu/x64/jit_uni_deconv_zp_pad_str_kernel.hpp"
#include "cpu/x64/jit_uni_x8s8s32x_deconvolution.hpp"

#define GET_OFF(field) offsetof(jit_deconv_args_t, field)

namespace dnnl {
namespace impl {
namespace cpu {
namespace x64 {

using namespace dnnl::impl::status;
using namespace dnnl::impl::memory_tracking::names;
using namespace dnnl::impl::utils;
using namespace Xbyak;

using namespace nstl;

#define wht_blk_off(d, g, ...) \
    (pd()->with_groups() ? (d).blk_off((g), __VA_ARGS__) \
                         : (d).blk_off(__VA_ARGS__))

template <cpu_isa_t isa>
status_t jit_uni_x8s8s32x_deconv_fwd_kernel_t<isa>::init_conf(
        jit_conv_conf_t &jcp, const deconvolution_desc_t &cd,
        memory_desc_t &src_md, memory_desc_t &weights_md, memory_desc_t &dst_md,
        const bool with_bias, memory_desc_t &bias_md, primitive_attr_t &attr,
        int nthreads) {
    const memory_desc_wrapper src_d(&src_md);
    const memory_desc_wrapper dst_d(&dst_md);
    const memory_desc_wrapper weights_d(&weights_md);
    const memory_desc_wrapper bias_d(&bias_md);

    // disabling verbose dispatch checks for unsupported isa for better eadability
    if (!mayiuse(isa)) return status::unimplemented;

    VDISPATCH_DECONVOLUTION_IC(
            one_of(src_d.data_type(), data_type::u8, data_type::s8),
            VERBOSE_UNSUPPORTED_DT);
    VDISPATCH_DECONVOLUTION_IC(
            weights_d.data_type() == data_type::s8, VERBOSE_UNSUPPORTED_DT);
    VDISPATCH_DECONVOLUTION_IC(
            one_of(dst_d.data_type(), data_type::f32, data_type::s32,
                    data_type::s8, data_type::u8),
            VERBOSE_UNSUPPORTED_DT);

    jcp = zero<decltype(jcp)>();
    jcp.nthr = nthreads;

    const bool with_groups = weights_d.ndims() == src_d.ndims() + 1;
    jcp.signed_input = src_d.data_type() == data_type::s8;
    const int ndims = jcp.ndims = dst_d.ndims();
    const bool is_1d = ndims == 3;
    const bool is_2d = ndims == 4;
    const bool is_3d = ndims == 5;
    const bool is_avx2 = isa == avx2;

    jcp.ngroups = with_groups ? weights_d.dims()[0] : 1;
    jcp.oc = dst_d.dims()[1] / jcp.ngroups;
    jcp.ic = src_d.dims()[1] / jcp.ngroups;
    jcp.id = is_3d ? src_d.dims()[2] : 1;
    jcp.oc_without_padding = dst_d.dims()[1] / jcp.ngroups;
    jcp.ic_without_padding = src_d.dims()[1] / jcp.ngroups;
    jcp.is_depthwise = true && with_groups
            && utils::everyone_is(
                    1, jcp.ic_without_padding, jcp.oc_without_padding);
    jcp.has_vnni = mayiuse(avx2_vnni);

    /* TODO: future work, on hold until depthwise specialized kernel is
     * implemented. */
    VDISPATCH_DECONVOLUTION_IC(
            !(jcp.is_depthwise && (jcp.signed_input || is_3d)),
            VERBOSE_UNSUPPORTED_FEATURE,
            "unsupported depthwise implementation for 3D convolution/signed "
            "input");

    VDISPATCH_DECONVOLUTION_IC(
            zero_points_valid(&attr), VERBOSE_UNSUPPORTED_ZP_CFG);
    jcp.src_zero_point = !attr.zero_points_.has_default_values(DNNL_ARG_SRC);
    jcp.dst_zero_point = !attr.zero_points_.has_default_values(DNNL_ARG_DST);
    jcp.zp_src_is_common = attr.zero_points_.get_mask(DNNL_ARG_SRC) == 0;

    format_tag_t dat_tag = utils::pick(
            ndims - 3, format_tag::nwc, format_tag::nhwc, format_tag::ndhwc);

    if (src_d.format_kind() == format_kind::any) {
        CHECK(memory_desc_init_by_tag(src_md, dat_tag));
        jcp.src_tag = dat_tag;
    } else {
        jcp.src_tag = src_d.matches_one_of_tag(dat_tag);
    }
    VDISPATCH_DECONVOLUTION_IC(
            jcp.src_tag == dat_tag, VERBOSE_UNSUPPORTED_TAG_S, "src");

    if (dst_d.format_kind() == format_kind::any) {
        CHECK(memory_desc_init_by_tag(dst_md, dat_tag));
        jcp.dst_tag = dat_tag;
    } else {
        jcp.dst_tag = dst_d.matches_one_of_tag(dat_tag);
    }
    VDISPATCH_DECONVOLUTION_IC(
            jcp.dst_tag == dat_tag, VERBOSE_UNSUPPORTED_TAG_S, "dst");

    auto set_or_check_wei_format = [&]() {
        using namespace format_tag;
        format_tag_t wei_tag;
        if (jcp.ic_block == 8 || jcp.ch_block == 8) {
            if (is_1d) {
                wei_tag = with_groups ? jcp.is_depthwise ? Goiw8g : gOIw2i8o4i
                                      : OIw2i8o4i;
            } else if (is_2d) {
                wei_tag = with_groups ? jcp.is_depthwise ? Goihw8g : gOIhw2i8o4i
                                      : OIhw2i8o4i;
            } else {
                wei_tag = with_groups ? gOIdhw2i8o4i : OIdhw2i8o4i;
            }
        } else {
            if (is_avx2) {
                assert(with_groups && jcp.ic_block == 4);
                wei_tag = is_3d ? gOIdhw4o4i : is_2d ? gOIhw4o4i : gOIw4o4i;
            } else {
                if (is_1d) {
                    wei_tag = with_groups ? jcp.is_depthwise ? Goiw4g : gOIw4o4i
                                          : OIw4o4i;
                } else if (is_2d) {
                    wei_tag = with_groups
                            ? jcp.is_depthwise ? Goihw4g : gOIhw4o4i
                            : OIhw4o4i;
                } else {
                    wei_tag = with_groups ? gOIdhw4o4i : OIdhw4o4i;
                }
            }
        }

        memory_desc_t want_wei_md = weights_md;
        CHECK_BOOL(memory_desc_init_by_tag(want_wei_md, wei_tag));

        if (jcp.signed_input && !jcp.is_depthwise) {
            want_wei_md.extra.flags = 0
                    | memory_extra_flags::compensation_conv_s8s8
                    | memory_extra_flags::scale_adjust;
            want_wei_md.extra.compensation_mask = (1 << 0)
                    + (with_groups && !jcp.is_depthwise ? (1 << 1) : 0);
            want_wei_md.extra.scale_adjust = jcp.has_vnni ? 1.f : 0.5f;
        }
        if (jcp.src_zero_point) set_zp_src_comp_flags(want_wei_md, with_groups);

        if (weights_md.format_kind == format_kind::any) {
            weights_md = want_wei_md;
            return true;
        }

        return weights_md == want_wei_md;
    };

    jcp.with_bias = with_bias;
    if (jcp.with_bias) {
        if (bias_d.format_kind() == format_kind::any)
            CHECK(memory_desc_init_by_tag(bias_md, format_tag::x));
    }

    jcp.prop_kind = cd.prop_kind;
    jcp.mb = src_d.dims()[0];
    jcp.ih = is_1d ? 1 : src_d.dims()[ndims - 2];
    jcp.iw = src_d.dims()[ndims - 1];
    jcp.od = is_3d ? dst_d.dims()[2] : 1;
    jcp.oh = is_1d ? 1 : dst_d.dims()[ndims - 2];
    jcp.ow = dst_d.dims()[ndims - 1];
    jcp.kd = is_3d ? weights_d.dims()[with_groups + 2] : 1;
    jcp.kh = is_1d ? 1 : weights_d.dims()[with_groups + ndims - 2];
    jcp.kw = weights_d.dims()[with_groups + ndims - 1];
    jcp.f_pad = is_3d ? cd.padding[0][0] : 0;
    jcp.t_pad = is_1d ? 0 : cd.padding[0][ndims - 4];
    jcp.l_pad = cd.padding[0][ndims - 3];
    jcp.stride_d = is_3d ? cd.strides[0] : 1;
    jcp.stride_h = is_1d ? 1 : cd.strides[ndims - 4];
    jcp.stride_w = cd.strides[ndims - 3];

    if (jcp.is_depthwise) {
        jcp.ch_block = is_avx2 ? 8 : 4;
        jcp.oc_block = 1;
        jcp.ic_block = 1;
    } else {
        jcp.ch_block = 1;
        jcp.oc_block = is_avx2 ? 8 : 4;
        jcp.ic_block = is_avx2 ? 8 : 4;

        if (jcp.ngroups == 1) {
            jcp.oc = utils::rnd_up(jcp.oc_without_padding, jcp.oc_block);
            jcp.ic = utils::rnd_up(jcp.ic_without_padding, jcp.ic_block);
        } else if (jcp.ngroups != 1
                && ((jcp.ic % jcp.ic_block != 0)
                        || (jcp.oc % jcp.oc_block != 0))) {
            /* For grouped convolution, oneDNN doesn't support padding.
             * When channel per group is not multiple of 8 in avx2:
             * - Use Xmm when channels per groups is multiple of 4.
             * - Otherwise return unimplemented */
            jcp.oc_block = jcp.ic_block = 4;
        }
        VDISPATCH_DECONVOLUTION_IC(
                !(jcp.ic % jcp.ic_block != 0 || jcp.oc % jcp.oc_block != 0),
                VERBOSE_BLOCKING_FAIL, "bad blocking dimensions");
    }

    VDISPATCH_DECONVOLUTION_IC(
            set_or_check_wei_format(), VERBOSE_UNSUPPORTED_TAG_S, "weights");

    jcp.dilate_d = is_3d ? cd.dilates[0] : 0;
    jcp.dilate_h = is_1d ? 0 : cd.dilates[ndims - 4];
    jcp.dilate_w = cd.dilates[ndims - 3];

    VDISPATCH_DECONVOLUTION_IC(
            !(!IMPLICATION(jcp.dilate_d, jcp.stride_d == 1)
                    || !IMPLICATION(jcp.dilate_h, jcp.stride_h == 1)
                    || !IMPLICATION(jcp.dilate_w, jcp.stride_w == 1)),
            VERBOSE_UNSUPPORTED_FEATURE,
            "unsupported shape with 'stride = 1' when 'dilate > 0'");

    const int ext_kw = calculate_extended_filter_size(jcp.kw, jcp.dilate_w);
    const int ext_kh = calculate_extended_filter_size(jcp.kh, jcp.dilate_h);
    const int ext_kd = calculate_extended_filter_size(jcp.kd, jcp.dilate_d);
    jcp.r_pad = calculate_end_padding(
            jcp.l_pad, jcp.iw, jcp.ow, jcp.stride_w, ext_kw);
    jcp.b_pad = calculate_end_padding(
            jcp.t_pad, jcp.ih, jcp.oh, jcp.stride_h, ext_kh);
    jcp.back_pad = calculate_end_padding(
            jcp.f_pad, jcp.id, jcp.od, jcp.stride_d, ext_kd);
    const bool kernel_outside_src = false || ext_kw <= jcp.l_pad
            || ext_kw <= jcp.r_pad || ext_kh <= jcp.t_pad || ext_kh <= jcp.b_pad
            || ext_kd <= jcp.f_pad || ext_kd <= jcp.back_pad;
    VDISPATCH_DECONVOLUTION_IC(!kernel_outside_src,
            VERBOSE_UNSUPPORTED_PAD_FEATURE, "weight and src size mismatch");

    CHECK(attr.set_default_formats(&dst_md));
    VDISPATCH_DECONVOLUTION_IC(
            post_ops_ok(jcp, dst_d, attr), VERBOSE_UNSUPPORTED_POSTOP);

    const auto &p = attr.post_ops_;
    const int eltwise_ind = p.find(primitive_kind::eltwise);
    jcp.with_eltwise = eltwise_ind != -1;
    if (jcp.with_eltwise) jcp.eltwise = p.entry_[eltwise_ind].eltwise;

    const int binary_ind = p.find(primitive_kind::binary);
    const int prelu_ind = p.find(primitive_kind::prelu);
    jcp.with_binary = !everyone_is(-1, binary_ind, prelu_ind);

    const int sum_ind = p.find(primitive_kind::sum);
    jcp.with_sum = sum_ind != -1;

    jcp.is_oc_scale = attr.scales_.get_mask(DNNL_ARG_WEIGHTS) > 0;
    jcp.with_src_scales = !attr.scales_.get(DNNL_ARG_SRC).has_default_values();
    jcp.with_wei_scales
            = !attr.scales_.get(DNNL_ARG_WEIGHTS).has_default_values();
    jcp.with_dst_scales = !attr.scales_.get(DNNL_ARG_DST).has_default_values();

    jcp.post_ops = p;

    jcp.dst_dt = dst_d.data_type();
    jcp.bia_dt = jcp.with_bias ? bias_d.data_type() : data_type::undef;
    jcp.typesize_bia
            = jcp.with_bias ? types::data_type_size(bias_d.data_type()) : 0;
    jcp.typesize_in = types::data_type_size(src_d.data_type());
    jcp.typesize_out = types::data_type_size(dst_d.data_type());

    jcp.nb_ch = div_up(jcp.ngroups, jcp.ch_block);
    jcp.nb_oc = jcp.oc / jcp.oc_block;
    jcp.nb_ic = jcp.ic / jcp.ic_block;

    /* kernel blocking params */
    const int regs = jcp.has_vnni ? 14 : 12;

    jcp.nb_ch_blocking = 1;
    jcp.nb_oc_blocking = nstl::min(4, jcp.nb_oc);
    for (; jcp.nb_oc_blocking > 1; jcp.nb_oc_blocking--)
        if (jcp.nb_oc % jcp.nb_oc_blocking == 0
                && jcp.l_pad <= regs / (jcp.nb_oc_blocking + 1))
            break;

    jcp.ur_w = regs / (jcp.nb_oc_blocking + 1);
    const int l_overflow = max(
            0, ((jcp.kw - 1) * (jcp.dilate_w + 1) - jcp.l_pad) / jcp.stride_w);

    if (jcp.ow < jcp.ur_w) {
        jcp.ur_w = jcp.ow;
        jcp.ur_w_tail = 0;
    } else {
        for (; jcp.ur_w >= 1; jcp.ur_w--) {
            /* ur_w should be multiple of stride_w in order
               to simplify logic for get_ow_start and get_ow_end */
            const bool is_multiple_of_stride = jcp.ur_w % jcp.stride_w == 0;

            /* boundary conditions:
               These conditions ensure all elements close to boundary
               are computed in a single call of compute loop */
            const bool left_boundary_covered
                    = jcp.ur_w >= l_overflow * jcp.stride_w;
            jcp.ur_w_tail = jcp.ow % jcp.ur_w;
            const int r_overflow_no_tail = max(0,
                    ((jcp.kw - 1) * (jcp.dilate_w + 1) - max(0, jcp.r_pad)
                            - jcp.ur_w_tail)
                            / jcp.stride_w);
            const bool right_boundary_covered
                    = jcp.ur_w >= r_overflow_no_tail * jcp.stride_w;

            if (is_multiple_of_stride && left_boundary_covered
                    && right_boundary_covered)
                break;
            else if (jcp.ur_w == 1)
                /* The boundary conditions above are also important
                   to maintain simplicity of calls to icb_loop,
                   if those conditions are not satisfied,
                   then special cases will need to be added
                   to use correct l_overflow/r_overflow values
                   when different iterations of compute loop
                   work on the locations close to boundary.
                   So to keep code simple, return unimplemented
                   for extreme case when a good ur_w cannot be found.
                 */
                return status::unimplemented;
        }
    }

    jcp.wei_adj_scale
            = (weights_d.extra().flags & memory_extra_flags::scale_adjust)
            ? weights_d.extra().scale_adjust
            : 1.f;

    jcp.loop_order = jcp.ngroups > 1 ? loop_ngc : loop_cgn;
    return status::success;
}

template <cpu_isa_t isa>
jit_uni_x8s8s32x_deconv_fwd_kernel_t<isa>::jit_uni_x8s8s32x_deconv_fwd_kernel_t(
        const jit_conv_conf_t &ajcp, const primitive_attr_t &attr,
        const memory_desc_wrapper &dst_d)
    : kernel_(nullptr) {

    const int ch_block = ajcp.is_depthwise ? ajcp.ch_block : ajcp.ic_block;
    switch (ch_block) {
        case 8:
            if (isa == avx2) {
                kernel_ = utils::make_unique<
                        _jit_avx2_x8s8s32x_deconv_fwd_kernel>(
                        ajcp, attr, dst_d);
                return;
            } else {
                assert(!"invalid channel blocking for current ISA");
            }
            break;
        case 4:
            kernel_ = utils::make_unique<
                    jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Xbyak::Xmm>>(
                    ajcp, attr, dst_d);
            return;
        default: assert(!"invalid channel blocking");
    }
}

template <cpu_isa_t isa>
jit_uni_x8s8s32x_deconv_fwd_kernel_t<
        isa>::~jit_uni_x8s8s32x_deconv_fwd_kernel_t()
        = default;

template <cpu_isa_t isa>
void jit_uni_x8s8s32x_deconv_fwd_kernel_t<isa>::init_scratchpad(
        memory_tracking::registrar_t &scratchpad, const jit_conv_conf_t &jcp,
        const primitive_attr_t &attr) {
    if (jcp.with_dst_scales) {
        const size_t n_dst_scales = static_cast<size_t>(
                jcp.is_depthwise ? jcp.ngroups : jcp.nthr);
        // See brgemm_types.hpp comment for `with_dst_scales`.
        scratchpad.book(
                key_conv_dst_scales, n_dst_scales * sizeof(float), 4096);
    }

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(jcp)) {
        const auto zp_pad_comp_size
                = static_cast<size_t>(jcp.oc_without_padding) * jcp.ngroups
                * jcp.kd * jcp.kh * jcp.kw;
        scratchpad.book<int32_t>(key_deconv_zp, zp_pad_comp_size);
    }
}

template <cpu_isa_t isa>
bool jit_uni_x8s8s32x_deconv_fwd_kernel_t<isa>::post_ops_ok(
        jit_conv_conf_t &jcp, const memory_desc_wrapper &dst_d,
        const primitive_attr_t &attr) {
    using namespace injector;

    return injector::post_ops_ok(post_ops_ok_args_t(isa, {sum, eltwise, binary},
            attr.post_ops_, &dst_d, false /*sum_at_pos_0_only*/,
            false /*sum_requires_scale_one*/, false /*sum_requires_zp_zero*/,
            true /*sum_requires_same_params*/,
            {broadcasting_strategy_t::per_oc,
                    broadcasting_strategy_t::scalar}));
}

template <cpu_isa_t isa, typename Vmm>
jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa,
        Vmm>::jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t(const jit_conv_conf_t
                                                               &ajcp,
        const primitive_attr_t &attr, const memory_desc_wrapper &dst_d)
    : jit_generator_t(jit_name(), isa)
    , jcp_(ajcp)
    , postops_injector_(nullptr)
    , ker_max_regs_(jcp_.has_vnni ? 14 : 12) {

    if (jcp_.with_eltwise || jcp_.with_binary || jcp_.with_sum) {
        const std::size_t tail_size = get_tail_size();

        static constexpr bool preserve_gpr = true;
        static constexpr bool preserve_vmm = true;
        static constexpr bool use_exact_tail_scalar_bcast = false;
        static constexpr size_t vmm_helper_idx = 15;

        const binary_injector::rhs_arg_static_params_t rhs_sp {vmm_helper_idx,
                this->r14, this->r15, this->r13, preserve_gpr, preserve_vmm,
                GET_OFF(post_ops_binary_rhs_arg_vec), GET_OFF(dst_orig), dst_d,
                tail_size, Xbyak::Opmask(2), use_exact_tail_scalar_bcast};
        const binary_injector::static_params_t bsp {this->param1_, rhs_sp};

        postops_injector_ = utils::make_unique<
                injector::jit_uni_postops_injector_t<isa, Vmm>>(
                this, jcp_.post_ops, bsp);
    }
}

template <cpu_isa_t isa, typename Vmm>
jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa,
        Vmm>::~jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t()
        = default;

template <cpu_isa_t isa, typename Vmm>
Vmm jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::vmm_out(
        int i_ur, int i_oc) const {
    const int idx = i_ur * jcp_.nb_oc_blocking + i_oc;
    assert(idx < ker_max_regs_);
    /* remap the reg indices to avoid using xmm0 in eltwise injector */
    return Vmm(15 - idx);
}

template <cpu_isa_t isa, typename Vmm>
Vmm jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::vmm_inp(
        int i_ic, int nb_x_blocking) const {
    const int idx = i_ic + nb_x_blocking * jcp_.ur_w;
    assert(idx < ker_max_regs_);
    return Vmm(15 - idx);
}

template <cpu_isa_t isa, typename Vmm>
int jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::get_ow_start(
        int ki, int l_overflow) const noexcept {
    int res = (jcp_.ow - 1 + jcp_.r_pad) % jcp_.stride_w
            + l_overflow * jcp_.stride_w
            - (jcp_.kw - 1 - ki) * (jcp_.dilate_w + 1);
    while (res < 0)
        res += jcp_.stride_w;
    return res;
}

template <cpu_isa_t isa, typename Vmm>
int jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::get_ow_end(
        int ur_w, int ki, int r_overflow) const noexcept {
    if (utils::one_of(ur_w, jcp_.ow, jcp_.ur_w_tail))
        ur_w += nstl::min(0, jcp_.r_pad); // remove negative padding
    int res = (ur_w - 1 + jcp_.l_pad) % jcp_.stride_w
            + r_overflow * jcp_.stride_w - ki * (jcp_.dilate_w + 1);
    while (res < 0)
        res += jcp_.stride_w;
    return ur_w - res;
}

template <cpu_isa_t isa, typename Vmm>
int jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::get_blocking_size()
        const noexcept {
    return jcp_.is_depthwise ? jcp_.ch_block : jcp_.oc_block;
}

template <cpu_isa_t isa, typename Vmm>
int jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::get_tail_size()
        const noexcept {
    return jcp_.is_depthwise ? jcp_.ngroups % jcp_.ch_block
                             : jcp_.oc_without_padding % jcp_.oc_block;
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::compute(
        const Vmm vreg_acc, const Vmm vreg_wei, const Vmm vreg_src) {

    if (jcp_.has_vnni) {
        vpdpbusd(vreg_acc, vreg_src, vreg_wei, Xbyak::VexEncoding);
    } else if (jcp_.is_depthwise) {
        uni_vmovups(vmm_tmp_, vreg_src);
        uni_vpmulld(vmm_tmp_, vmm_tmp_, vreg_wei);
        uni_vpaddd(vreg_acc, vreg_acc, vmm_tmp_);
    } else {
        uni_vpmaddubsw(vmm_tmp_, vreg_src, vreg_wei);
        uni_vpmaddwd(vmm_tmp_, vmm_tmp_, vmm_one_);
        uni_vpaddd(vreg_acc, vreg_acc, vmm_tmp_);
    }
}

template <cpu_isa_t isa, typename Vmm>
std::function<Vmm()> jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa,
        Vmm>::prepare_round_robin_vmm_inp_generator(int ur_w) const noexcept {

    const int start_vmm_idx = vmm_inp(ur_w - 1, jcp_.nb_oc_blocking).getIdx();
    const int end_vmm_idx = vmm_inp(0, jcp_.nb_oc_blocking).getIdx() + 1;
    int current_vmm_idx = start_vmm_idx;

    return [current_vmm_idx, start_vmm_idx, end_vmm_idx]() mutable {
        const Vmm vmm {static_cast<int>(current_vmm_idx++)};

        if (current_vmm_idx == end_vmm_idx) current_vmm_idx = start_vmm_idx;

        return vmm;
    };
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa,
        Vmm>::apply_zp_src_pad_str_comp(int ur_w, int l_overflow,
        int r_overflow, bool h_padded) {
    Xbyak::Label end_zp_pad, no_tail;

    // apply once per icb loop, zp src stride paddding compensation calculate as
    // zp_pad_str_compensation = conv(1, weights_s8) * zero_point_source
    cmp(reg_icb_, jcp_.nb_ic);
    jne(end_zp_pad, T_NEAR);

    if (jcp_.ngroups % jcp_.ch_block
            || jcp_.oc_without_padding % jcp_.oc_block) {
        if (jcp_.is_depthwise)
            cmp(reg_oc_blocks_, jcp_.nb_ch - 1);
        else
            cmp(reg_oc_blocks_, jcp_.nb_oc - jcp_.nb_oc_blocking);
        jne(no_tail, T_NEAR);

        append_zp_src_pad_str_comp(
                ur_w, l_overflow, r_overflow, h_padded, true /*last_oc_block*/);
        jmp(end_zp_pad, T_NEAR);
    }

    L(no_tail);
    append_zp_src_pad_str_comp(
            ur_w, l_overflow, r_overflow, h_padded, false /*last_oc_block*/);

    L(end_zp_pad);
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa,
        Vmm>::append_zp_src_pad_str_comp(int ur_w, int l_overflow,
        int r_overflow, bool h_padded, bool last_oc_block) {

    const auto &reg_zp_src_pad_comp = reg_scratch_;
    const auto get_next_comp_vmm = prepare_round_robin_vmm_inp_generator(ur_w);
    bool base_comp_addr_loaded = false;

    const auto load_base_zp_src_pad_comp_addr = [&]() {
        if (!base_comp_addr_loaded) {
            if (jcp_.ndims == 5) mov(reg_scratch_preserved_, reg_scratch_);

            if (jcp_.ndims > 3)
                mov(reg_zp_src_pad_comp, zp_src_pad_comp_addr_);
            else
                mov(reg_zp_src_pad_comp,
                        qword[param1_ + GET_OFF(zp_src_pad_str_compensation)]);

            base_comp_addr_loaded = true;
        }
    };

    const auto load_zp_src_pad_comp
            = [&](const Vmm zp_pad_comp_vmm, const Xbyak::Address &comp_addr,
                      const int ocb) {
        const bool is_tail = last_oc_block && ocb == jcp_.nb_oc_blocking - 1;

        if (is_tail)
            load_data(data_type::s32, zp_pad_comp_vmm, comp_addr,
                    get_tail_size());
        else
            uni_vmovups(zp_pad_comp_vmm, comp_addr);
    };

    const auto get_zp_src_comp_pad_off = [&](int it_kw, int ocb) {
        const auto kw_offset = it_kw * jcp_.oc_without_padding * jcp_.ngroups;
        const auto oc_offset = ocb * jcp_.oc_block;

        return (kw_offset + oc_offset) * sizeof(int32_t);
    };

    for (int it_kw = 0; it_kw < jcp_.kw; ++it_kw) {
        const int ow_start = get_ow_start(it_kw, l_overflow);
        const int ow_end = get_ow_end(ur_w, it_kw, r_overflow);

        for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
            Vmm zp_src_comp_pad_vmm; // will be assigned later
            bool ocb_zp_loaded = false;

            const auto zp_src_comp_pad_off
                    = get_zp_src_comp_pad_off(it_kw, ocb);

            for (int it_ow = 0; it_ow < ur_w; ++it_ow) {

                const bool inside_padded_area = h_padded
                        || !(it_ow >= ow_start && it_ow < ow_end
                                && ((it_ow + jcp_.l_pad - it_kw) % jcp_.stride_w
                                        == 0));

                if (inside_padded_area) {
                    load_base_zp_src_pad_comp_addr();

                    if (!ocb_zp_loaded) {
                        zp_src_comp_pad_vmm = get_next_comp_vmm();
                        const auto comp_addr = ptr[reg_zp_src_pad_comp
                                + zp_src_comp_pad_off];
                        load_zp_src_pad_comp(
                                zp_src_comp_pad_vmm, comp_addr, ocb);
                        ocb_zp_loaded = true;
                    }

                    const auto vmm_dst = vmm_out(it_ow, ocb);
                    uni_vpaddd(vmm_dst, vmm_dst, zp_src_comp_pad_vmm);
                }
            }
        }
    }

    if (jcp_.ndims > 3) {
        if (!base_comp_addr_loaded) load_base_zp_src_pad_comp_addr();

        const auto kh_offset = jcp_.kw * jcp_.oc_without_padding * jcp_.ngroups
                * sizeof(int32_t);

        add(reg_zp_src_pad_comp, kh_offset);
        mov(zp_src_pad_comp_addr_, reg_zp_src_pad_comp);
    }

    if (jcp_.ndims == 5 && base_comp_addr_loaded)
        mov(reg_scratch_, reg_scratch_preserved_);
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::compute_ker(int ur_w,
        int l_overflow, int r_overflow, ker_block_t last_ic_block_flag,
        bool h_padded) {

    const bool signed_input_or_src_zp
            = (jcp_.signed_input || jcp_.src_zero_point);

    const int ch_block_all = jcp_.ch_block * jcp_.ic_block * jcp_.oc_block;
    const int ur_w_stride = signed_input_or_src_zp ? 1 : jcp_.stride_w;

    const auto src_offset = [this](int oj, int icb, int ki) {
        return jcp_.typesize_in
                * (((oj + jcp_.l_pad - ki * (jcp_.dilate_w + 1))
                           / jcp_.stride_w)
                                * jcp_.ngroups * jcp_.ic_without_padding
                        + icb * 4);
    };

    const auto kernel_offset = [this, ch_block_all](int ocb, int icb, int ki) {
        return jcp_.typesize_in
                * ((ocb * jcp_.nb_ic * jcp_.kd * jcp_.kh * jcp_.kw + ki)
                                * ch_block_all
                        + icb * jcp_.oc_block * IC_SUB_STEP);
    };

    for (int ki = 0; ki < jcp_.kw; ki++) {

        const int jj_start = get_ow_start(ki, l_overflow);
        const int jj_end = get_ow_end(ur_w, ki, r_overflow);

        const int _start = (signed_input_or_src_zp) ? 0 : jj_start;
        const int _end = (signed_input_or_src_zp) ? ur_w : jj_end;

        const int tail_size = jcp_.is_depthwise ? jcp_.ngroups % jcp_.ch_block
                                                : jcp_.ic_without_padding % 4;
        const int n_ic_blocks = jcp_.is_depthwise
                ? 1
                : (last_ic_block_flag != no_last_block
                                  ? div_up(jcp_.ic_without_padding
                                                    % jcp_.ic_block,
                                            4)
                                  : jcp_.ic_block / 4);

        for (int icb1 = 0; icb1 < n_ic_blocks; icb1++) {
            if (h_padded) {
                /* fill padded area with shifted values */
                if (jcp_.signed_input) {
                    const Vmm inp = vmm_inp(0, jcp_.nb_oc_blocking);
                    uni_vpxor(inp, inp, inp);
                    uni_vpsubb(inp, inp, vmm_shift_);
                }
            } else {

                for (int jj = _start; jj < _end; jj += ur_w_stride) {

                    const int aux_src_off = src_offset(jj, icb1, ki);
                    const auto vmm_src = vmm_inp(jj, jcp_.nb_oc_blocking);

                    if (jj >= jj_start && jj < jj_end
                            && ((jj + jcp_.l_pad - ki) % jcp_.stride_w == 0)) {
                        if (jcp_.is_depthwise) {
                            if (tail_size != 0)
                                assert(jcp_.nb_oc_blocking == 1);
                            const bool mask_flag
                                    = last_ic_block_flag != no_last_block
                                    && tail_size;
                            load_data(data_type::u8, vmm_src, aux_reg_src_,
                                    aux_src_off,
                                    mask_flag ? tail_size : jcp_.ch_block);
                        } else if ((last_ic_block_flag == last_sp_block)
                                && tail_size != 0 && icb1 == n_ic_blocks - 1) {
                            const auto vmm_inp_tmp = Xmm(vmm_src.getIdx());
                            load_bytes(vmm_inp_tmp, aux_reg_src_, aux_src_off,
                                    tail_size);
                            uni_vpbroadcastd(vmm_src, vmm_inp_tmp);
                        } else {
                            uni_vpbroadcastd(
                                    vmm_src, ptr[aux_reg_src_ + aux_src_off]);
                        }
                        if (jcp_.signed_input)
                            uni_vpsubb(vmm_src, vmm_src, vmm_shift_);
                    } else {
                        /* fill padded area with shifted values */
                        if (jcp_.signed_input) {
                            uni_vpxor(vmm_src, vmm_src, vmm_src);
                            uni_vpsubb(vmm_src, vmm_src, vmm_shift_);
                        }
                    }
                }
            }
            for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
                const int aux_filt_off = kernel_offset(ocb, icb1, ki);

                if (_end - _start > 0) {
                    if (jcp_.is_depthwise) {
                        uni_vpmovsxbd(
                                vmm_wei_, ptr[aux_reg_filt_ + aux_filt_off]);
                    } else
                        uni_vmovups(
                                vmm_wei_, ptr[aux_reg_filt_ + aux_filt_off]);
                }

                for (int jj = _start; jj < _end; jj += ur_w_stride) {

                    const bool inside_padded_area = h_padded
                            || !(jj >= jj_start && jj < jj_end
                                    && ((jj + jcp_.l_pad - ki) % jcp_.stride_w
                                            == 0));
                    const auto vmm_dst = vmm_out(jj, ocb);
                    if (jcp_.signed_input || !inside_padded_area) {
                        const Vmm inp = vmm_inp(
                                h_padded ? 0 : jj, jcp_.nb_oc_blocking);
                        compute(vmm_dst, vmm_wei_, inp);
                    }
                }
            }
        }
    }

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(jcp_))
        apply_zp_src_pad_str_comp(ur_w, l_overflow, r_overflow, h_padded);
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::kh_loop(int ur_w,
        int l_overflow, int r_overflow, ker_block_t last_ic_block_flag) {

    const bool signed_input_or_src_zp
            = (jcp_.signed_input || jcp_.src_zero_point);

    const int ch_block_all = jcp_.ch_block * jcp_.ic_block * jcp_.oc_block;
    const int shift_src_ih = jcp_.typesize_in * (jcp_.dilate_h + 1) * jcp_.iw
            * jcp_.ngroups * jcp_.ic_without_padding;
    const int shift_src_id = jcp_.typesize_in * (jcp_.dilate_d + 1) * jcp_.ih
            * jcp_.iw * jcp_.ngroups * jcp_.ic_without_padding;
    const int stride_h = signed_input_or_src_zp ? 1 : jcp_.stride_h;
    const int shift_filt_kh
            = jcp_.typesize_in * jcp_.kw * ch_block_all * stride_h;
    const int stride_d = signed_input_or_src_zp ? 1 : jcp_.stride_d;
    const int shift_filt_kd
            = jcp_.typesize_in * jcp_.kw * ch_block_all * jcp_.kh * stride_d;

    Label kd_loop_label, kh_loop_label, skip_kh_loop, skip_kd_loop;
    Label t_overflow_label, no_t_overflow_label, b_overflow_label,
            no_b_overflow_label;
    Label back_overflow_label, no_back_overflow_label, d_h_overflow_label,
            front_overflow_label, no_front_overflow_label, d_h_overflow_label2;
    if (jcp_.ndims == 5) {
        mov(aux_reg_filt_d_, reg_filt_);
        mov(aux_reg_src_d_, reg_src_);

        if (signed_input_or_src_zp) {
            mov(reg_ki_, ptr[param1_ + GET_OFF(back_overflow)]);
            cmp(reg_ki_, 0);
            je(no_back_overflow_label, T_NEAR);

            L(back_overflow_label);
            {
                mov(aux_reg_filt_, aux_reg_filt_d_);
                mov(reg_kh_, jcp_.kh);
                L(d_h_overflow_label);
                {
                    compute_ker(ur_w, 0, 0, last_ic_block_flag, true);
                    add(aux_reg_filt_, shift_filt_kh);
                    dec(reg_kh_);
                    jnz(d_h_overflow_label);
                }

                add(aux_reg_filt_d_, shift_filt_kd);
                dec(reg_ki_);
                jnz(back_overflow_label);
            }
            L(no_back_overflow_label);
        }

        mov(reg_ki_, ptr[param1_ + GET_OFF(kd_padding)]);

        if ((signed_input_or_src_zp) || (jcp_.dilate_d >= jcp_.id)
                || (jcp_.kd < jcp_.stride_d)
                || ((!signed_input_or_src_zp)
                        && ((min(jcp_.f_pad, jcp_.back_pad) < 0)
                                || ((jcp_.kd - 1) * (jcp_.dilate_d + 1)
                                        < nstl::max(
                                                jcp_.f_pad, jcp_.back_pad))))) {
            cmp(reg_ki_, 0);
            je(skip_kd_loop, T_NEAR);
        }

        L(kd_loop_label);
        mov(aux_reg_src_, aux_reg_src_d_);
        mov(aux_reg_filt_, aux_reg_filt_d_);
    } else {
        mov(aux_reg_src_, reg_src_);
        mov(aux_reg_filt_, reg_filt_);
    }

    if (signed_input_or_src_zp && jcp_.ndims > 3) {
        /* Weights are transposed, so first compute 'bottom' padding. */
        mov(reg_overflow_, ptr[param1_ + GET_OFF(b_overflow)]);
        cmp(reg_overflow_, 0);
        je(no_b_overflow_label, T_NEAR);
        L(b_overflow_label);
        {
            compute_ker(ur_w, 0, 0, last_ic_block_flag, true);

            add(aux_reg_filt_, shift_filt_kh);
            dec(reg_overflow_);
            cmp(reg_overflow_, 0);
            jg(b_overflow_label, T_NEAR);
        }
        L(no_b_overflow_label);
    }

    mov(reg_kh_, ptr[param1_ + GET_OFF(kh_padding)]);

    if ((signed_input_or_src_zp) || (jcp_.dilate_h >= jcp_.ih)
            || (jcp_.kh < jcp_.stride_h)
            || ((!signed_input_or_src_zp)
                    && ((min(jcp_.t_pad, jcp_.b_pad) < 0)
                            || ((jcp_.kh - 1) * (jcp_.dilate_h + 1)
                                    < nstl::max(jcp_.t_pad, jcp_.b_pad))))) {
        cmp(reg_kh_, 0);
        je(skip_kh_loop, T_NEAR);
    }

    L(kh_loop_label);
    {
        compute_ker(ur_w, l_overflow, r_overflow, last_ic_block_flag, false);
        sub(aux_reg_src_, shift_src_ih);
        add(aux_reg_filt_, shift_filt_kh);
        dec(reg_kh_);

        /* Insert weight compensation in stride 'holes' */
        if (signed_input_or_src_zp && jcp_.stride_h > 1) {
            Label kh_comp_loop;

            cmp(reg_kh_, 0);
            je(skip_kh_loop, T_NEAR);
            mov(reg_comp_strides_, jcp_.stride_h - 1);
            L(kh_comp_loop);
            {
                compute_ker(ur_w, 0, 0, last_ic_block_flag, true);
                add(aux_reg_filt_, shift_filt_kh);
                dec(reg_comp_strides_);
                cmp(reg_comp_strides_, 0);
                jg(kh_comp_loop, T_NEAR);
            }
        }
        cmp(reg_kh_, 0);
        jg(kh_loop_label, T_NEAR);
    }
    L(skip_kh_loop);
    if (signed_input_or_src_zp && jcp_.ndims > 3) {
        mov(reg_overflow_, ptr[param1_ + GET_OFF(t_overflow)]);
        cmp(reg_overflow_, 0);
        je(no_t_overflow_label, T_NEAR);
        L(t_overflow_label);
        {
            compute_ker(ur_w, 0, 0, last_ic_block_flag, true);

            add(aux_reg_filt_, shift_filt_kh);
            dec(reg_overflow_);
            cmp(reg_overflow_, 0);
            jg(t_overflow_label, T_NEAR);
        }
        L(no_t_overflow_label);
    }

    if (jcp_.ndims == 5) {
        sub(aux_reg_src_d_, shift_src_id);
        add(aux_reg_filt_d_, shift_filt_kd);
        dec(reg_ki_);

        /* Insert weight compensation in stride 'holes' */
        if (signed_input_or_src_zp && jcp_.stride_d > 1) {
            Label kd_comp_loop, kd_kh_comp_loop;
            cmp(reg_ki_, 0);
            jz(skip_kd_loop, T_NEAR);
            mov(reg_comp_strides_, jcp_.stride_d - 1);
            L(kd_comp_loop);
            mov(aux_reg_filt_, aux_reg_filt_d_);
            mov(reg_kh_, jcp_.kh);
            L(kd_kh_comp_loop);
            {
                compute_ker(ur_w, 0, 0, last_ic_block_flag, true);
                add(aux_reg_filt_, shift_filt_kh);
                dec(reg_kh_);
                jnz(kd_kh_comp_loop, T_NEAR);
            }
            add(aux_reg_filt_d_, shift_filt_kd);
            dec(reg_comp_strides_);
            jnz(kd_comp_loop);
        }

        cmp(reg_ki_, 0);
        jg(kd_loop_label, T_NEAR);
        L(skip_kd_loop);
        if (signed_input_or_src_zp) {
            mov(reg_ki_, ptr[param1_ + GET_OFF(f_overflow)]);
            cmp(reg_ki_, 0);
            jz(no_front_overflow_label, T_NEAR);
            L(front_overflow_label);
            {
                mov(aux_reg_filt_, aux_reg_filt_d_);
                mov(reg_kh_, jcp_.kh);
                L(d_h_overflow_label2);
                {
                    compute_ker(ur_w, 0, 0, last_ic_block_flag, true);
                    add(aux_reg_filt_, shift_filt_kh);
                    dec(reg_kh_);
                    jnz(d_h_overflow_label2);
                }
                add(aux_reg_filt_d_, shift_filt_kd);
                dec(reg_ki_);
                jnz(front_overflow_label);
            }
            L(no_front_overflow_label);
        }
    }
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::prepare_output(
        int ur_w) {
    for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
        for (int ur = 0; ur < ur_w; ur++) {
            const Vmm vmm = vmm_out(ur, ocb);
            uni_vpxor(vmm, vmm, vmm);
        }
    }
    if (jcp_.signed_input) {
        const auto xmm_shift = Xbyak::Xmm(vmm_shift_.getIdx());
        mov(reg_scratch_, 0x80808080);
        uni_vmovq(xmm_shift, reg_scratch_);
        uni_vpbroadcastd(vmm_shift_, xmm_shift);
    }
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::cvt2ps(
        data_type_t type_in, const Vmm vmm_in, const Reg64 reg, int offset,
        int load_size) {

    load_data(type_in, vmm_in, reg, offset, load_size);
    if (type_in != data_type::f32) uni_vcvtdq2ps(vmm_in, vmm_in);
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::apply_postops(int ur_w,
        bool last_oc_block, const float *p_sum_scale, const int32_t *p_sum_zp) {
    const auto sum_injector = [&]() {
        if (p_sum_scale) { // post_op: sum
            for (int k = 0; k < jcp_.nb_oc_blocking; k++) {
                const bool mask_flag
                        = last_oc_block == 1 && k == jcp_.nb_oc_blocking - 1;
                for (int j = 0; j < ur_w; j++) {
                    const int aux_output_offset = jcp_.typesize_out
                            * (k * jcp_.oc_block
                                    + j * jcp_.oc_without_padding
                                            * jcp_.ngroups);
                    cvt2ps(jcp_.dst_dt, vmm_prev_dst_, reg_dst_,
                            aux_output_offset,
                            mask_flag ? get_tail_size() : get_blocking_size());
                    if (*p_sum_zp != 0) {
                        uni_vbroadcastss(vmm_sum_zp_, ptr[reg_ptr_sum_zp_]);
                        uni_vcvtdq2ps(vmm_sum_zp_, vmm_sum_zp_);
                        uni_vsubps(vmm_prev_dst_, vmm_prev_dst_, vmm_sum_zp_);
                    }
                    const Vmm vmm = vmm_out(j, k);
                    if (*p_sum_scale == 1.f)
                        uni_vaddps(vmm, vmm, vmm_prev_dst_);
                    else {
                        uni_vbroadcastss(vmm_tmp_, ptr[reg_ptr_sum_scale_]);
                        uni_vfmadd231ps(vmm, vmm_prev_dst_, vmm_tmp_);
                    }
                }
            }
        }
    };

    if (p_sum_scale)
        postops_injector_->set_lambda_injector(
                primitive_kind::sum, sum_injector);

    binary_injector::rhs_arg_dynamic_params_t rhs_arg_params;
    if (jcp_.with_binary) {
        for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
            const bool mask_flag
                    = last_oc_block && ocb == jcp_.nb_oc_blocking - 1;
            for (int ur = 0; ur < ur_w; ur++) {
                const int vmm_idx = vmm_out(ur, ocb).getIdx();
                const size_t aux_output_offset = jcp_.typesize_out
                        * (ocb * jcp_.oc_block
                                + ur * jcp_.oc_without_padding * jcp_.ngroups);

                rhs_arg_params.vmm_idx_to_out_reg.emplace(vmm_idx, reg_dst_);
                rhs_arg_params.vmm_idx_to_out_elem_off_val.emplace(
                        vmm_idx, aux_output_offset);
                if (mask_flag) rhs_arg_params.vmm_tail_idx_.emplace(vmm_idx);
            }
        }
    }
    const int nb_oc_block
            = jcp_.is_depthwise ? jcp_.nb_ch_blocking : jcp_.nb_oc_blocking;
    postops_injector_->compute_vector_range(
            16 - nb_oc_block * ur_w, 16, rhs_arg_params);
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::store_output(
        int ur_w, bool last_oc_block) {
    mov(reg_bias_, ptr[param1_ + GET_OFF(bias)]);

    if (jcp_.signed_input)
        mov(reg_compensation_, ptr[param1_ + GET_OFF(compensation)]);

    if (jcp_.src_zero_point) {
        mov(reg_zp_src_, ptr[param1_ + GET_OFF(src_zero_point)]);
        mov(reg_zp_compensation_, ptr[param1_ + GET_OFF(zp_compensation)]);
    }

    const auto &p = jcp_.post_ops;
    const int sum_idx = p.find(primitive_kind::sum);
    const float *p_sum_scale
            = (sum_idx != -1) ? &p.entry_[sum_idx].sum.scale : nullptr;
    const int32_t *p_sum_zp
            = (sum_idx != -1) ? &p.entry_[sum_idx].sum.zero_point : nullptr;

    if (jcp_.src_zero_point) {
        const auto &vmm_src_zp = vmm_tmp_;
        const auto &vmm_zp_comp = vmm_scales_;
        uni_vbroadcastss(vmm_src_zp, ptr[reg_zp_src_]);

        for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
            const int zp_offset = sizeof(int32_t) * ocb * jcp_.oc_block;
            const bool mask_flag
                    = last_oc_block && ocb == jcp_.nb_oc_blocking - 1;
            const int load_size
                    = mask_flag ? get_tail_size() : get_blocking_size();
            load_data(data_type::s32, vmm_zp_comp, reg_zp_compensation_,
                    zp_offset, load_size);
            uni_vpmulld(vmm_zp_comp, vmm_zp_comp, vmm_src_zp);

            for (int ur = 0; ur < ur_w; ur++) {
                const auto vmm_dst = vmm_out(ur, ocb);
                uni_vpaddd(vmm_dst, vmm_dst, vmm_zp_comp);
            }
        }
    }

    for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
        const bool mask_flag = last_oc_block && ocb == jcp_.nb_oc_blocking - 1;

        // TODO: scales support is done not in the most optimal way.
        // If there're two free Vmm registers, one can be used to store
        // scale_adjust value permanently, the second one can re-use data
        // from it and multiply by src_scale that can be obtained at the
        // point of scales loading. Then it can be used when multiplying
        // by wei_scales. And further re-used for dst scales to avoid
        // reading from the same address, but reading from the Vmm instead.
        // This would save 1st and 3rd sections for every output Vmm.
        //
        // If only one Vmm is found, it will add scale_adjust overhead per
        // src_scale loading, but the second part of the idea holds.
        //
        // Note: attempts to identify these Vmms were not taken.

        // `avx2` is less flexible ISA in terms of tail and broadcast handling.
        // Thus, need to save scales values in Vmm registers.
        bool is_vmm_scales_set = false;
        if (jcp_.with_src_scales) {
            mov(reg_src_scales_, ptr[param1_ + GET_OFF(src_scales)]);
            uni_vbroadcastss(vmm_scales_, ptr[reg_src_scales_]);
            is_vmm_scales_set = true;
        }
        if (jcp_.with_wei_scales) {
            mov(reg_wei_scales_, ptr[param1_ + GET_OFF(wei_scales)]);

            if (!jcp_.is_oc_scale) {
                uni_vbroadcastss(vmm_scales_tmp_, ptr[reg_wei_scales_]);
            } else {
                const int scale_offset = jcp_.is_oc_scale
                        * (sizeof(float) * ocb * jcp_.oc_block);
                if (mask_flag) {
                    uni_vpxor(
                            vmm_scales_tmp_, vmm_scales_tmp_, vmm_scales_tmp_);
                    cvt2ps(data_type::f32, vmm_scales_tmp_, reg_wei_scales_,
                            scale_offset, get_tail_size());
                } else {
                    uni_vmovups(vmm_scales_tmp_,
                            ptr[reg_wei_scales_ + scale_offset]);
                }
            }
            if (is_vmm_scales_set) {
                uni_vmulps(vmm_scales_, vmm_scales_, vmm_scales_tmp_);
            } else {
                uni_vmovups(vmm_scales_, vmm_scales_tmp_);
            }
            is_vmm_scales_set = true;
        }
        if (jcp_.wei_adj_scale != 1.f) {
            mov(reg_scale_adjust_, float2int(1.f / jcp_.wei_adj_scale));
            auto vmm_scale_adjust = vmm_scales_tmp_;
            auto xmm_scale_adjust = Xmm(vmm_scale_adjust.getIdx());
            uni_vmovq(xmm_scale_adjust, reg_scale_adjust_);
            uni_vbroadcastss(vmm_scale_adjust, xmm_scale_adjust);
            if (is_vmm_scales_set) {
                uni_vmulps(vmm_scales_, vmm_scales_, vmm_scale_adjust);
            } else {
                uni_vmovups(vmm_scales_, vmm_scale_adjust);
            }
            is_vmm_scales_set = true;
        }

        // The order of this load is important. `vmm_bias` is used as a
        // temporary vector register for scales. Load bias data into it
        // after scales are processed.
        const auto vmm_bias_ = vmm_tmp_;
        if (jcp_.with_bias) {
            int bias_offset = jcp_.typesize_bia * ocb * jcp_.oc_block;
            cvt2ps(jcp_.bia_dt, vmm_bias_, reg_bias_, bias_offset,
                    mask_flag ? get_tail_size() : get_blocking_size());
        }

        if (jcp_.signed_input) {
            const int comp_offset = sizeof(int32_t) * ocb * jcp_.oc_block;
            cvt2ps(data_type::s32, vmm_comp_, reg_compensation_, comp_offset,
                    mask_flag ? get_tail_size() : get_blocking_size());
        }

        for (int ur = 0; ur < ur_w; ur++) {
            const Vmm vmm = vmm_out(ur, ocb);
            uni_vcvtdq2ps(vmm, vmm);
            if (jcp_.signed_input) uni_vaddps(vmm, vmm, vmm_comp_);

            if (is_vmm_scales_set) uni_vmulps(vmm, vmm, vmm_scales_);

            if (jcp_.with_bias) uni_vaddps(vmm, vmm, vmm_bias_);
        }
    }

    if (p_sum_scale && *p_sum_scale != 1.f)
        mov(reg_ptr_sum_scale_, reinterpret_cast<size_t>(p_sum_scale));
    if (p_sum_zp && *p_sum_zp != 0) {
        mov(reg_ptr_sum_zp_, reinterpret_cast<size_t>(p_sum_zp));
    }
    if (jcp_.with_eltwise || jcp_.with_binary || jcp_.with_sum)
        apply_postops(ur_w, last_oc_block, p_sum_scale, p_sum_zp);

    if (jcp_.with_dst_scales) {
        mov(reg_dst_scales_, ptr[param1_ + GET_OFF(dst_scales)]);
        uni_vbroadcastss(vmm_dst_scales_, ptr[reg_dst_scales_]);

        for_(int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++)
        for (int ur = 0; ur < ur_w; ur++) {
            const auto vmm = vmm_out(ur, ocb);
            uni_vmulps(vmm, vmm, vmm_dst_scales_);
        }
    }

    if (jcp_.dst_zero_point) {
        mov(reg_zp_dst_, ptr[param1_ + GET_OFF(dst_zero_point)]);
        const auto &vmm_zp_dst = vmm_tmp_;
        uni_vbroadcastss(vmm_zp_dst, ptr[reg_zp_dst_]);
        uni_vcvtdq2ps(vmm_zp_dst, vmm_zp_dst);

        for_(int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++)
        for (int ur = 0; ur < ur_w; ur++) {
            const auto vmm_dst = vmm_out(ur, ocb);
            uni_vaddps(vmm_dst, vmm_dst, vmm_zp_dst);
        }
    }

    // Properly saturate the accumulators for integer datatypes

    // No need to saturate on lower bound for signed integer types, as
    // the conversion to int would return INT_MIN, and then proper
    // saturation will happen when storing data
    if (jcp_.dst_dt == data_type::u8) {
        uni_vpxor(vmm_zero_, vmm_zero_, vmm_zero_);
        for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
            for (int ur = 0; ur < ur_w; ur++) {
                const Vmm vmm = vmm_out(ur, ocb);
                uni_vmaxps(vmm, vmm, vmm_zero_);
            }
        }
    }

    if (one_of(jcp_.dst_dt, data_type::u8, data_type::s8, data_type::s32)) {
        float saturation_ubound = types::max_value<float>(jcp_.dst_dt);
        const Xmm xmm_saturation(vmm_saturation_.getIdx());
        mov(reg_ptr_saturation_ubound_, float2int(saturation_ubound));
        uni_vmovq(xmm_saturation, reg_ptr_saturation_ubound_);
        uni_vbroadcastss(vmm_saturation_, xmm_saturation);

        for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
            for (int ur = 0; ur < ur_w; ur++) {
                const Vmm vmm = vmm_out(ur, ocb);
                uni_vminps(vmm, vmm, vmm_saturation_);
            }
        }
    }

    if (one_of(jcp_.dst_dt, data_type::u8, data_type::s8, data_type::s32)) {
        for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
            for (int ur = 0; ur < ur_w; ur++) {
                const Vmm vmm = vmm_out(ur, ocb);
                uni_vcvtps2dq(vmm, vmm);
            }
        }
    }

    /* write out register to output_addr */
    for (int ocb = 0; ocb < jcp_.nb_oc_blocking; ocb++) {
        const bool mask_flag = last_oc_block && ocb == jcp_.nb_oc_blocking - 1;
        for (int ur = 0; ur < ur_w; ur++) {
            const int aux_dst_off = jcp_.typesize_out
                    * (ur * jcp_.ngroups * jcp_.oc_without_padding
                            + ocb * jcp_.oc_block);
            const Vmm r_vmm = vmm_out(ur, ocb);
            store_data(jcp_.dst_dt, r_vmm, reg_dst_, aux_dst_off,
                    mask_flag ? get_tail_size() : get_blocking_size());
        }
    }
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::icb_loop(
        int ur_w, int l_overflow, int r_overflow, bool is_last_sp_block) {

    const int shift_src_icb = jcp_.typesize_in * jcp_.ic_block;
    const size_t shift_filt_icb = (size_t)jcp_.typesize_in * jcp_.kd * jcp_.kh
            * jcp_.kw * jcp_.ic_block * jcp_.oc_block;

    prepare_output(ur_w);

    Label skip_icb_loop, icb_loop_label;

    mov(reg_icb_, jcp_.nb_ic);
    mov(reg_oc_blocks_, ptr[param1_ + GET_OFF(oc_blocks)]);

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(jcp_)
            && jcp_.ndims > 3) {
        mov(reg_scratch_,
                qword[param1_ + GET_OFF(zp_src_pad_str_compensation)]);
        mov(zp_src_pad_comp_addr_, reg_scratch_);
    }

    L(icb_loop_label);
    {
        if (jcp_.ngroups % jcp_.ch_block != 0
                || jcp_.ic_without_padding != jcp_.ic) {
            Label common_ker, end_ker;
            if (jcp_.is_depthwise) {
                cmp(reg_oc_blocks_, jcp_.nb_ch - 1);
                jne(common_ker, T_NEAR);
            } else {
                cmp(reg_icb_, 1);
                jg(common_ker, T_NEAR);
            }

            kh_loop(ur_w, l_overflow, r_overflow, last_sp_block);
            jmp(end_ker, T_NEAR);

            L(common_ker);
            kh_loop(ur_w, l_overflow, r_overflow, no_last_block);

            L(end_ker);
        } else {
            kh_loop(ur_w, l_overflow, r_overflow, no_last_block);
        }

        add(reg_src_, shift_src_icb);
        safe_add(reg_filt_, shift_filt_icb, reg_ker_long_offt_);
        dec(reg_icb_);
        cmp(reg_icb_, 0);
        jg(icb_loop_label, T_NEAR);
    }

    /* come-back pointers */
    sub(reg_src_, jcp_.nb_ic * shift_src_icb);
    safe_sub(reg_filt_, jcp_.nb_ic * shift_filt_icb, reg_ker_long_offt_);
    L(skip_icb_loop);

    if (jcp_.ngroups % jcp_.ch_block != 0
            || jcp_.oc_without_padding != jcp_.oc) {
        Label common_store, end_store;
        if (jcp_.is_depthwise)
            cmp(reg_oc_blocks_, jcp_.nb_ch - 1);
        else
            cmp(reg_oc_blocks_, jcp_.nb_oc - jcp_.nb_oc_blocking);
        jne(common_store, T_NEAR);

        store_output(ur_w, true);
        jmp(end_store, T_NEAR);

        L(common_store);
        store_output(ur_w, false);

        L(end_store);

    } else {
        store_output(ur_w, false);
    }
}

template <cpu_isa_t isa, typename Vmm>
void jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<isa, Vmm>::generate() {
    preamble();

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(jcp_))
        sub(rsp, reserved_stack_size_);

    const auto vmm_one_128 = Xbyak::Xmm(vmm_one_.getIdx());
    mov(reg_scratch_, 0x10001);
    uni_vmovq(vmm_one_128, reg_scratch_);
    uni_vpbroadcastd(vmm_one_, vmm_one_128);

    mov(reg_src_, ptr[param1_ + GET_OFF(src)]);
    mov(reg_filt_, ptr[param1_ + GET_OFF(filt)]);
    mov(reg_dst_, ptr[param1_ + GET_OFF(dst)]);

    const int dst_shift = jcp_.typesize_out * jcp_.ur_w * jcp_.ngroups
            * jcp_.oc_without_padding;
    const int src_shift = jcp_.typesize_in * (jcp_.ur_w / jcp_.stride_w)
            * jcp_.ngroups * jcp_.ic_without_padding;

    const int l_overflow = max(0,
            ((jcp_.kw - 1) * (jcp_.dilate_w + 1) - jcp_.l_pad) / jcp_.stride_w);
    const int r_overflow = max(0,
            ((jcp_.kw - 1) * (jcp_.dilate_w + 1) - max(0, jcp_.r_pad))
                    / jcp_.stride_w);

    const int r_overflow1 = nstl::max(0,
            ((jcp_.kw - 1) * (jcp_.dilate_w + 1) - nstl::max(0, jcp_.r_pad)
                    - jcp_.ur_w_tail)
                    / jcp_.stride_w);
    int nur_w = jcp_.ow / jcp_.ur_w;
    if (r_overflow1 > 0) nur_w--;

    if (jcp_.ur_w == jcp_.ow) {
        icb_loop(jcp_.ur_w, l_overflow, r_overflow, true);
    } else if (nur_w == 0) {
        icb_loop(jcp_.ur_w, l_overflow, r_overflow1, jcp_.ur_w_tail == 0);
        add(reg_src_, src_shift);
        add(reg_dst_, dst_shift);
        if (jcp_.ur_w_tail != 0) icb_loop(jcp_.ur_w_tail, 0, r_overflow, true);
    } else {
        xor_(reg_nur_w_, reg_nur_w_);
        if (l_overflow > 0) {
            icb_loop(jcp_.ur_w, l_overflow, 0, false);
            add(reg_src_, src_shift);
            add(reg_dst_, dst_shift);
            inc(reg_nur_w_);
        }
        if ((l_overflow <= 0 && nur_w > 0) || (l_overflow > 0 && nur_w > 1)) {
            Label ow_loop_label;
            L(ow_loop_label);
            {
                icb_loop(jcp_.ur_w, 0, 0, false);
                add(reg_src_, src_shift);
                add(reg_dst_, dst_shift);
                inc(reg_nur_w_);
                cmp(reg_nur_w_, nur_w);
                jl(ow_loop_label, T_NEAR);
            }
        }
        if (r_overflow1 > 0) {
            icb_loop(jcp_.ur_w, 0, r_overflow1, jcp_.ur_w_tail == 0);
            add(reg_src_, src_shift);
            add(reg_dst_, dst_shift);
        }
        if (jcp_.ur_w_tail != 0) {
            icb_loop(jcp_.ur_w_tail, 0, r_overflow, true);
        }
    }

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(jcp_))
        add(rsp, reserved_stack_size_);

    postamble();

    if (jcp_.with_eltwise)
        postops_injector_->prepare_table(/* generate = */ true);
}

template <cpu_isa_t isa>
jit_uni_x8s8s32x_deconvolution_fwd_t<isa>::jit_uni_x8s8s32x_deconvolution_fwd_t(
        const pd_t *apd)
    : primitive_t(apd) {}

template <cpu_isa_t isa>
jit_uni_x8s8s32x_deconvolution_fwd_t<
        isa>::~jit_uni_x8s8s32x_deconvolution_fwd_t()
        = default;

template <cpu_isa_t isa>
status_t jit_uni_x8s8s32x_deconvolution_fwd_t<isa>::pd_t::init(
        engine_t *engine) {
    using namespace data_type;
    using skip_mask_t = primitive_attr_t::skip_mask_t;

    VDISPATCH_DECONVOLUTION(is_fwd(), VERBOSE_BAD_PROPKIND);
    VDISPATCH_DECONVOLUTION((desc()->alg_kind & alg_kind::deconvolution_direct),
            VERBOSE_BAD_ALGORITHM);
    VDISPATCH_DECONVOLUTION(utils::one_of(src_md(0)->data_type, s8, u8),
            VERBOSE_UNSUPPORTED_DT);
    VDISPATCH_DECONVOLUTION(
            weights_md(0)->data_type == s8, VERBOSE_UNSUPPORTED_DT);
    VDISPATCH_DECONVOLUTION(
            IMPLICATION(with_bias(),
                    utils::one_of(weights_md(1)->data_type, f32, s32, s8, u8)),
            VERBOSE_UNSUPPORTED_DT);
    VDISPATCH_DECONVOLUTION(
            utils::one_of(dst_md(0)->data_type, f32, s32, s8, u8),
            VERBOSE_UNSUPPORTED_DT);
    VDISPATCH_DECONVOLUTION(
            desc()->accum_data_type == s32, VERBOSE_UNSUPPORTED_DT);
    VDISPATCH_DECONVOLUTION(
            attr()->has_default_values(skip_mask_t::scales
                    | skip_mask_t::post_ops | skip_mask_t::zero_points),
            VERBOSE_UNSUPPORTED_ATTR);
    VDISPATCH_DECONVOLUTION(attr_scales_ok(), VERBOSE_UNSUPPORTED_SCALES_CFG);
    VDISPATCH_DECONVOLUTION(impl::is_dense_format_kind({src_md(0),
                                    weights_md(0), weights_md(1), dst_md(0)}),
            VERBOSE_UNSUPPORTED_SPARSE_CFG);

    CHECK(jit_uni_x8s8s32x_deconv_fwd_kernel_t<isa>::init_conf(jcp_, *desc(),
            src_md_, weights_md_, dst_md_, with_bias(), bias_md_, attr_,
            dnnl_get_max_threads()));

    auto scratchpad = scratchpad_registry().registrar();
    jit_uni_x8s8s32x_deconv_fwd_kernel_t<isa>::init_scratchpad(
            scratchpad, jcp_, *attr());

    return status::success;
}

template <cpu_isa_t isa>
status_t jit_uni_x8s8s32x_deconvolution_fwd_t<isa>::init(engine_t *engine) {
    CHECK(safe_ptr_assign(kernel_,
            new jit_uni_x8s8s32x_deconv_fwd_kernel_t<isa>(pd()->jcp_,
                    *pd()->attr(), memory_desc_wrapper(pd()->dst_md()))));

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(pd()->jcp_)) {
        CHECK(safe_ptr_assign(zp_src_pad_comp_kernel_,
                zp::create_deconv_zp_pad_str_comp_ker<isa>(pd()->jcp_)));
        const auto zp_kernel_status = zp_src_pad_comp_kernel_->create_kernel();
        if (zp_kernel_status != status::success) return zp_kernel_status;
    }

    return kernel_->create_kernel();
}

template <cpu_isa_t isa>
status_t jit_uni_x8s8s32x_deconvolution_fwd_t<isa>::execute(
        const exec_ctx_t &ctx) const {
    const auto &_pd = pd();
    const auto &ndims = _pd->ndims();
    if (ndims == 3)
        return execute_forward_1d(ctx);
    else if (ndims == 4)
        return execute_forward_2d(ctx);
    else if (ndims == 5)
        return execute_forward_3d(ctx);
    else
        return status::unimplemented;
    return status::success;
}

template <cpu_isa_t isa>
status_t jit_uni_x8s8s32x_deconvolution_fwd_t<isa>::execute_forward_1d(
        const exec_ctx_t &ctx) const {
    const auto src = CTX_IN_MEM(const char *, DNNL_ARG_SRC);
    const auto weights = CTX_IN_MEM(const int8_t *, DNNL_ARG_WEIGHTS);
    const auto bias = CTX_IN_MEM(const char *, DNNL_ARG_BIAS);
    auto dst = CTX_OUT_MEM(char *, DNNL_ARG_DST);

    const int32_t *src_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC);
    const int32_t *dst_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_DST);

    const auto &jcp = pd()->jcp_;

    const memory_desc_wrapper src_d(pd()->src_md());
    const memory_desc_wrapper dst_d(pd()->dst_md());
    const memory_desc_wrapper weights_d(pd()->weights_md(0));
    const memory_desc_wrapper bias_d(pd()->weights_md(1));

    const size_t dst_dt_size = types::data_type_size(dst_d.data_type());

    const auto post_ops_binary_rhs_arg_vec
            = binary_injector::prepare_binary_args(jcp.post_ops, ctx);
    const auto &scratchpad = ctx.get_scratchpad_grantor();
    int32_t *zp_src_comp_scratch = scratchpad.get<int32_t>(key_deconv_zp);

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(jcp))
        zp::compute_deconv_zp_pad_str_comp_ker(jcp, pd()->with_groups(),
                weights_d, weights, src_zero_points, zp_src_comp_scratch,
                zp_src_pad_comp_kernel_.get());

    const int oc_chunks = jcp.nb_oc / jcp.nb_oc_blocking;
    const int nb_groups = jcp.nb_ch;

    const void *src_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC);
    const void *wei_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS);
    const void *dst_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST);

    const size_t offset = weights_d.size() - weights_d.additional_buffer_size();
    auto w = const_cast<int8_t *>(weights);
    int32_t *compensation = (jcp.signed_input)
            ? reinterpret_cast<int32_t *>(&w[offset])
            : nullptr;
    const int32_t *zp_compensation = jcp.src_zero_point
            ? get_src_zp_comp_from_wei(
                      weights, weights_d, jcp.signed_input, jcp.ngroups, jcp.oc)
            : nullptr;

    parallel(jcp.nthr, [= COMPAT_THIS_CAPTURE](const int ithr, const int nthr) {
        int start {0}, end {0};
        const int work_amount = jcp.mb * nb_groups * oc_chunks;
        balance211(work_amount, nthr, ithr, start, end);

        auto p = jit_deconv_args_t();

        float *dst_scales_inv_ptr = nullptr;
        if (jcp.with_dst_scales) {
            const float *dst_scales_ptr
                    = static_cast<const float *>(dst_scales);
            dst_scales_inv_ptr
                    = ctx.get_scratchpad_grantor().template get<float>(
                              key_conv_dst_scales)
                    + ithr;
            dst_scales_inv_ptr[0] = 1.f / dst_scales_ptr[0];
        }

        int n {0}, g {0}, occ {0};
        if (jcp.loop_order == loop_ngc)
            nd_iterator_init(start, n, jcp.mb, g, nb_groups, occ, oc_chunks);
        else if (jcp.loop_order == loop_cgn)
            nd_iterator_init(start, occ, oc_chunks, g, nb_groups, n, jcp.mb);
        else
            assert(!"unsupported loop order");
        while (start < end) {

            const int ocb = occ * jcp.nb_oc_blocking;
            const int g_oc
                    = (g * jcp.ch_block * jcp.nb_oc + ocb) * jcp.oc_block;
            const int g_ic = g * jcp.ch_block * jcp.ic;

            p.dst = dst + dst_dt_size * dst_d.blk_off(n, g_oc);
            p.src = src + src_d.blk_off(n, g_ic);
            p.filt = weights + wht_blk_off(weights_d, g, ocb, 0);
            p.bias = jcp.with_bias
                    ? bias + (bias_d.blk_off(g_oc) * jcp.typesize_bia)
                    : nullptr;
            p.compensation = (jcp.signed_input) ? compensation + g_oc : nullptr;
            p.src_scales = src_scales;
            p.wei_scales = jcp.with_wei_scales
                    ? static_cast<const float *>(wei_scales)
                            + jcp.is_oc_scale * g_oc
                    : nullptr;
            p.dst_scales = dst_scales_inv_ptr;
            p.t_overflow = 0;
            p.b_overflow = 0;
            p.kh_padding = jcp.kh;
            p.oc_blocks = jcp.is_depthwise ? g : ocb;
            p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
            p.zp_compensation
                    = jcp.src_zero_point ? zp_compensation + g_oc : nullptr;
            p.zp_src_pad_str_compensation = zp_src_comp_scratch
                    ? zp_src_comp_scratch + g_oc
                    : nullptr;
            p.src_zero_point = src_zero_points;
            p.dst_zero_point = dst_zero_points;
            p.dst_orig = dst;

            (*kernel_)(&p);

            ++start;
            if (jcp.loop_order == loop_ngc)
                nd_iterator_step(n, jcp.mb, g, nb_groups, occ, oc_chunks);
            else if (jcp.loop_order == loop_cgn)
                nd_iterator_step(occ, oc_chunks, g, nb_groups, n, jcp.mb);
            else
                assert(!"unsupported loop order");
        }
    });
    return status::success;
}

template <cpu_isa_t isa>
status_t jit_uni_x8s8s32x_deconvolution_fwd_t<isa>::execute_forward_2d(
        const exec_ctx_t &ctx) const {
    const auto src = CTX_IN_MEM(const char *, DNNL_ARG_SRC);
    const auto weights = CTX_IN_MEM(const int8_t *, DNNL_ARG_WEIGHTS);
    const auto bias = CTX_IN_MEM(const char *, DNNL_ARG_BIAS);
    auto dst = CTX_OUT_MEM(char *, DNNL_ARG_DST);

    const int32_t *src_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC);
    const int32_t *dst_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_DST);

    const memory_desc_wrapper src_d(pd()->src_md());
    const memory_desc_wrapper dst_d(pd()->dst_md());
    const memory_desc_wrapper weights_d(pd()->weights_md(0));
    const memory_desc_wrapper bias_d(pd()->weights_md(1));

    const size_t dst_dt_size = types::data_type_size(dst_d.data_type());

    const auto &jcp = pd()->jcp_;
    const auto post_ops_binary_rhs_arg_vec
            = binary_injector::prepare_binary_args(jcp.post_ops, ctx);

    const auto &scratchpad = ctx.get_scratchpad_grantor();
    int32_t *zp_src_comp_scratch = scratchpad.get<int32_t>(key_deconv_zp);

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(jcp))
        zp::compute_deconv_zp_pad_str_comp_ker(jcp, pd()->with_groups(),
                weights_d, weights, src_zero_points, zp_src_comp_scratch,
                zp_src_pad_comp_kernel_.get());

    const int oc_chunks = jcp.nb_oc / jcp.nb_oc_blocking;
    const int nb_groups = jcp.nb_ch;

    const size_t src_h_stride = src_d.blk_off(0, 0, 1);
    const size_t dst_h_stride = dst_d.blk_off(0, 0, 1);
    const size_t wht_kh_stride = wht_blk_off(weights_d, 0, 0, 0, 1);

    const void *src_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC);
    const void *wei_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS);
    const void *dst_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST);

    const size_t offset = weights_d.size() - weights_d.additional_buffer_size();
    auto w = const_cast<int8_t *>(weights);
    int32_t *compensation = (jcp.signed_input)
            ? reinterpret_cast<int32_t *>(&w[offset])
            : nullptr;
    const int32_t *zp_compensation = jcp.src_zero_point
            ? get_src_zp_comp_from_wei(
                      weights, weights_d, jcp.signed_input, jcp.ngroups, jcp.oc)
            : nullptr;

    parallel(jcp.nthr, [= COMPAT_THIS_CAPTURE](const int ithr, const int nthr) {
        int start {0}, end {0};
        const int work_amount = jcp.mb * nb_groups * oc_chunks * jcp.oh;
        balance211(work_amount, nthr, ithr, start, end);

        auto p = jit_deconv_args_t();

        float *dst_scales_inv_ptr = nullptr;
        if (jcp.with_dst_scales) {
            const float *dst_scales_ptr
                    = static_cast<const float *>(dst_scales);
            dst_scales_inv_ptr
                    = ctx.get_scratchpad_grantor().template get<float>(
                              key_conv_dst_scales)
                    + ithr;
            dst_scales_inv_ptr[0] = 1.f / dst_scales_ptr[0];
        }

        /*loop order = cgn*/
        int n {0}, g {0}, occ {0}, oh_s {0};
        if (jcp.loop_order == loop_ngc)
            nd_iterator_init(start, n, jcp.mb, g, nb_groups, occ, oc_chunks,
                    oh_s, jcp.oh);
        else if (jcp.loop_order == loop_cgn)
            nd_iterator_init(start, occ, oc_chunks, g, nb_groups, n, jcp.mb,
                    oh_s, jcp.oh);
        else
            assert(!"unsupported loop order");
        while (start < end) {

            const int ocb = occ * jcp.nb_oc_blocking;
            const int g_oc
                    = (g * jcp.ch_block * jcp.nb_oc + ocb) * jcp.oc_block;
            const int g_ic = g * jcp.ch_block * jcp.ic;
            const int work_rem = end - start;
            const int oh_e
                    = oh_s + work_rem > jcp.oh ? jcp.oh : oh_s + work_rem;

            const auto dst_w = dst + dst_dt_size * dst_d.blk_off(n, g_oc);
            const auto src_w = src + src_d.blk_off(n, g_ic);
            const auto wht_w = weights + wht_blk_off(weights_d, g, ocb, 0);
            const auto bias_w = jcp.with_bias
                    ? bias + (bias_d.blk_off(g_oc) * jcp.typesize_bia)
                    : nullptr;
            const int32_t *compensation_w
                    = (jcp.signed_input) ? compensation + g_oc : nullptr;

            for (int oj = oh_s; oj < oh_e; oj++) {
                int ih_max = 0, kh_lo = 0, kh_len = 0;
                if (jcp.dilate_h != 0 && jcp.stride_h == 1) {
                    /* dilation */
                    const int dilate_h = jcp.dilate_h + 1;
                    // Note: use div_up to account for "holes" in filter
                    const int o_t_overflow = div_up(
                            max(0, (jcp.kh - 1) * dilate_h - oj - jcp.t_pad),
                            dilate_h);
                    const int o_b_overflow
                            = div_up(max(0,
                                             (jcp.kh - 1) * dilate_h + 1
                                                     - jcp.oh + oj - jcp.b_pad),
                                    dilate_h);
                    kh_len = jcp.kh - o_t_overflow - o_b_overflow;
                    kh_lo = o_b_overflow;
                    ih_max = oj + jcp.t_pad - o_b_overflow * dilate_h;
                } else {
                    const int o_t_overflow = max(
                            0, (jcp.kh - (oj + 1 + jcp.t_pad)) / jcp.stride_h);
                    const int o_b_overflow = max(0,
                            ((oj + jcp.kh) - (jcp.oh + jcp.b_pad))
                                    / jcp.stride_h);
                    const int overflow_kh_hi = jcp.kh - 1
                            - modulo(jcp.oh + jcp.b_pad - (oj + 1),
                                    jcp.stride_h);
                    const int overflow_kh_lo = (oj + jcp.t_pad) % jcp.stride_h;

                    kh_len = (overflow_kh_hi - overflow_kh_lo) / jcp.stride_h
                            + 1 - o_t_overflow - o_b_overflow;
                    kh_lo = overflow_kh_lo + o_b_overflow * jcp.stride_h;
                    ih_max = (oj + jcp.t_pad - kh_lo) / jcp.stride_h;
                }

                const int wei_stride
                        = (!jcp.signed_input && !jcp.src_zero_point)
                        ? kh_lo * wht_kh_stride
                        : 0;
                p.src = src_w + ih_max * src_h_stride;
                p.dst = dst_w + dst_dt_size * oj * dst_h_stride;
                p.filt = wht_w + wei_stride;
                p.bias = bias_w;
                p.compensation = compensation_w;
                p.t_overflow = jcp.dilate_h > 0
                        ? jcp.kh - kh_len - kh_lo
                        : max(0,
                                  jcp.kh
                                          - (kh_lo
                                                  + max(0, kh_len - 1)
                                                          * jcp.stride_h
                                                  + 1));
                p.b_overflow = kh_lo;
                p.kh_padding = kh_len;
                p.src_scales = src_scales;
                p.wei_scales = jcp.with_wei_scales
                        ? static_cast<const float *>(wei_scales)
                                + jcp.is_oc_scale * g_oc
                        : nullptr;
                p.dst_scales = dst_scales_inv_ptr;
                p.oc_blocks = jcp.is_depthwise ? g : ocb;
                p.post_ops_binary_rhs_arg_vec
                        = post_ops_binary_rhs_arg_vec.data();
                p.zp_compensation
                        = jcp.src_zero_point ? zp_compensation + g_oc : nullptr;
                p.zp_src_pad_str_compensation = jcp.src_zero_point
                        ? zp_src_comp_scratch + g_oc
                        : nullptr;
                p.src_zero_point = src_zero_points;
                p.dst_zero_point = dst_zero_points;
                p.dst_orig = dst;

                (*kernel_)(&p);
            }
            if (jcp.loop_order == loop_ngc)
                nd_iterator_jump(start, end, n, jcp.mb, g, nb_groups, occ,
                        oc_chunks, oh_s, jcp.oh);
            else if (jcp.loop_order == loop_cgn)
                nd_iterator_jump(start, end, occ, oc_chunks, g, nb_groups, n,
                        jcp.mb, oh_s, jcp.oh);
            else
                assert(!"unsupported loop order");
        }
    });
    return status::success;
}

template <cpu_isa_t isa>
status_t jit_uni_x8s8s32x_deconvolution_fwd_t<isa>::execute_forward_3d(
        const exec_ctx_t &ctx) const {
    const auto src = CTX_IN_MEM(const char *, DNNL_ARG_SRC);
    const auto weights = CTX_IN_MEM(const int8_t *, DNNL_ARG_WEIGHTS);
    const auto bias = CTX_IN_MEM(const char *, DNNL_ARG_BIAS);
    const auto dst = CTX_OUT_MEM(char *, DNNL_ARG_DST);

    const int32_t *src_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC);
    const int32_t *dst_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_DST);

    const memory_desc_wrapper src_d(pd()->src_md());
    const memory_desc_wrapper dst_d(pd()->dst_md());
    const memory_desc_wrapper weights_d(pd()->weights_md(0));
    const memory_desc_wrapper bias_d(pd()->weights_md(1));

    const size_t dst_dt_size = types::data_type_size(dst_d.data_type());

    const auto &jcp = pd()->jcp_;
    const auto post_ops_binary_rhs_arg_vec
            = binary_injector::prepare_binary_args(jcp.post_ops, ctx);

    const auto &scratchpad = ctx.get_scratchpad_grantor();
    int32_t *const zp_src_comp_scratch = scratchpad.get<int32_t>(key_deconv_zp);

    if (zp::should_calculate_deconv_zp_src_pad_str_comp(jcp))
        zp::compute_deconv_zp_pad_str_comp_ker(jcp, pd()->with_groups(),
                weights_d, weights, src_zero_points, zp_src_comp_scratch,
                zp_src_pad_comp_kernel_.get());

    const int oc_chunks = jcp.nb_oc / jcp.nb_oc_blocking;
    const int &nb_groups = jcp.nb_ch;

    const size_t src_d_stride = src_d.blk_off(0, 0, 1);
    const size_t src_h_stride = src_d.blk_off(0, 0, 0, 1);
    const size_t dst_d_stride = dst_d.blk_off(0, 0, 1);
    const size_t dst_h_stride = dst_d.blk_off(0, 0, 0, 1);
    const size_t wht_kd_stride = wht_blk_off(weights_d, 0, 0, 0, 1);
    const size_t wht_kh_stride = wht_blk_off(weights_d, 0, 0, 0, 0, 1);

    const void *src_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC);
    const void *wei_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS);
    const void *dst_scales
            = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST);

    const size_t offset = weights_d.size() - weights_d.additional_buffer_size();
    auto w = const_cast<int8_t *>(weights);
    int32_t *compensation = (jcp.signed_input)
            ? reinterpret_cast<int32_t *>(&w[offset])
            : nullptr;
    const int32_t *zp_compensation = jcp.src_zero_point
            ? get_src_zp_comp_from_wei(
                      weights, weights_d, jcp.signed_input, jcp.ngroups, jcp.oc)
            : nullptr;

    parallel(jcp.nthr, [= COMPAT_THIS_CAPTURE](const int ithr, const int nthr) {
        int start {0}, end {0};
        int work_amount = jcp.mb * nb_groups * oc_chunks * jcp.od * jcp.oh;
        balance211(work_amount, nthr, ithr, start, end);

        auto p = jit_deconv_args_t();

        float *dst_scales_inv_ptr = nullptr;
        if (jcp.with_dst_scales) {
            const float *dst_scales_ptr
                    = static_cast<const float *>(dst_scales);
            dst_scales_inv_ptr
                    = ctx.get_scratchpad_grantor().template get<float>(
                              key_conv_dst_scales)
                    + ithr;
            dst_scales_inv_ptr[0] = 1.f / dst_scales_ptr[0];
        }

        /*loop order = cgn*/
        int n {0}, g {0}, occ {0}, od_s {0}, oh_s {0};
        if (jcp.loop_order == loop_ngc)
            nd_iterator_init(start, n, jcp.mb, g, nb_groups, occ, oc_chunks,
                    od_s, jcp.od, oh_s, jcp.oh);
        else if (jcp.loop_order == loop_cgn)
            nd_iterator_init(start, occ, oc_chunks, g, nb_groups, n, jcp.mb,
                    od_s, jcp.od, oh_s, jcp.oh);
        else
            assert(!"unsupported loop order");
        while (start < end) {

            const int ocb = occ * jcp.nb_oc_blocking;
            const int g_oc
                    = (g * jcp.ch_block * jcp.nb_oc + ocb) * jcp.oc_block;
            const int g_ic = g * jcp.ch_block * jcp.ic;
            const int work_rem = end - start;
            const int oh_e
                    = oh_s + work_rem > jcp.oh ? jcp.oh : oh_s + work_rem;
            int input_d_s = 0, kd_len = 0, kd_lo = 0;
            int d_t_overflow, d_back_overflow;

            if (jcp.dilate_d != 0 && jcp.stride_d == 1) {
                /* dilation */
                int dilate_d = jcp.dilate_d + 1;
                // Note: use div_up to account for "holes" in filter
                d_t_overflow = div_up(
                        max(0, (jcp.kd - 1) * dilate_d - od_s - jcp.f_pad),
                        dilate_d);
                d_back_overflow
                        = div_up(max(0,
                                         (jcp.kd - 1) * dilate_d + 1 - jcp.od
                                                 + od_s - jcp.back_pad),
                                dilate_d);
                kd_len = jcp.kd - d_t_overflow - d_back_overflow;
                kd_lo = d_back_overflow;
                input_d_s = od_s + jcp.f_pad - d_back_overflow * dilate_d;
            } else {
                const int d_t_overflow = max(
                        0, (jcp.kd - (od_s + 1 + jcp.f_pad)) / jcp.stride_d);
                const int d_back_overflow = max(0,
                        ((od_s + jcp.kd) - (jcp.od + jcp.back_pad))
                                / jcp.stride_d);
                const int overflow_kd_hi = jcp.kd - 1
                        - modulo(jcp.od + jcp.back_pad - (od_s + 1),
                                jcp.stride_d);
                const int overflow_kd_lo = (od_s + jcp.f_pad) % jcp.stride_d;

                kd_len = (overflow_kd_hi - overflow_kd_lo) / jcp.stride_d + 1
                        - d_t_overflow - d_back_overflow;
                kd_lo = overflow_kd_lo + d_back_overflow * jcp.stride_d;
                input_d_s = (od_s + jcp.f_pad - kd_lo) / jcp.stride_d;
            }

            auto dst_w = dst
                    + dst_dt_size
                            * (dst_d.blk_off(n, g_oc) + od_s * dst_d_stride);
            const auto src_w
                    = src + src_d.blk_off(n, g_ic) + input_d_s * src_d_stride;
            const auto wht_w = weights + wht_blk_off(weights_d, g, ocb, 0)
                    + ((jcp.signed_input || jcp.src_zero_point) ? 0 : kd_lo)
                            * wht_kd_stride;
            const auto bias_w = jcp.with_bias
                    ? bias + (bias_d.blk_off(g_oc) * jcp.typesize_bia)
                    : nullptr;
            const int32_t *compensation_w
                    = (jcp.signed_input) ? compensation + g_oc : nullptr;

            for (int oj = oh_s; oj < oh_e; oj++) {
                int ih_max = 0, kh_lo = 0, kh_len = 0;
                if (jcp.dilate_h != 0 && jcp.stride_h == 1) {
                    /* dilation */
                    const int dilate_h = jcp.dilate_h + 1;
                    // Note: use div_up to account for "holes" in filter
                    const int o_t_overflow = div_up(
                            max(0, (jcp.kh - 1) * dilate_h - oj - jcp.t_pad),
                            dilate_h);
                    const int o_b_overflow
                            = div_up(max(0,
                                             (jcp.kh - 1) * dilate_h + 1
                                                     - jcp.oh + oj - jcp.b_pad),
                                    dilate_h);
                    kh_len = jcp.kh - o_t_overflow - o_b_overflow;
                    kh_lo = o_b_overflow;
                    ih_max = oj + jcp.t_pad - o_b_overflow * dilate_h;
                } else {
                    const int o_t_overflow = max(
                            0, (jcp.kh - (oj + 1 + jcp.t_pad)) / jcp.stride_h);
                    const int o_b_overflow = max(0,
                            ((oj + jcp.kh) - (jcp.oh + jcp.b_pad))
                                    / jcp.stride_h);
                    const int overflow_kh_hi = jcp.kh - 1
                            - modulo(jcp.oh + jcp.b_pad - (oj + 1),
                                    jcp.stride_h);
                    const int overflow_kh_lo = (oj + jcp.t_pad) % jcp.stride_h;

                    kh_len = (overflow_kh_hi - overflow_kh_lo) / jcp.stride_h
                            + 1 - o_t_overflow - o_b_overflow;
                    kh_lo = overflow_kh_lo + o_b_overflow * jcp.stride_h;
                    ih_max = (oj + jcp.t_pad - kh_lo) / jcp.stride_h;
                }

                const int wei_stride
                        = (!jcp.signed_input && !jcp.src_zero_point)
                        ? kh_lo * wht_kh_stride
                        : 0;
                p.src = src_w + ih_max * src_h_stride;
                p.dst = dst_w + dst_dt_size * oj * dst_h_stride;
                p.filt = wht_w + wei_stride;
                p.bias = bias_w;
                p.compensation = compensation_w;
                /* Note: Currently this kernel doesn't support dilations and
                strides together */
                p.t_overflow = jcp.dilate_h > 0
                        ? jcp.kh - kh_len - kh_lo
                        : max(0,
                                  jcp.kh
                                          - (kh_lo
                                                  + max(0, kh_len - 1)
                                                          * jcp.stride_h
                                                  + 1));
                p.b_overflow = kh_lo;
                p.f_overflow = jcp.dilate_d > 0
                        ? jcp.kd - kd_len - kd_lo
                        : max(0,
                                  jcp.kd
                                          - (kd_lo
                                                  + max(0, kd_len - 1)
                                                          * jcp.stride_d
                                                  + 1));
                p.back_overflow = kd_lo;
                p.kh_padding = kh_len;
                p.kd_padding = kd_len;
                p.src_scales = src_scales;
                p.wei_scales = jcp.with_wei_scales
                        ? static_cast<const float *>(wei_scales)
                                + jcp.is_oc_scale * g_oc
                        : nullptr;
                p.dst_scales = dst_scales_inv_ptr;
                p.oc_blocks = jcp.is_depthwise ? g : ocb;
                p.post_ops_binary_rhs_arg_vec
                        = post_ops_binary_rhs_arg_vec.data();
                p.zp_compensation
                        = jcp.src_zero_point ? zp_compensation + g_oc : nullptr;
                p.zp_src_pad_str_compensation = jcp.src_zero_point
                        ? zp_src_comp_scratch + g_oc
                        : nullptr;
                p.src_zero_point = src_zero_points;
                p.dst_zero_point = dst_zero_points;
                p.dst_orig = dst;

                (*kernel_)(&p);
            }

            if (jcp.loop_order == loop_ngc)
                nd_iterator_jump(start, end, n, jcp.mb, g, nb_groups, occ,
                        oc_chunks, od_s, jcp.od, oh_s, jcp.oh);
            else if (jcp.loop_order == loop_cgn)
                nd_iterator_jump(start, end, occ, oc_chunks, g, nb_groups, n,
                        jcp.mb, od_s, jcp.od, oh_s, jcp.oh);
            else
                assert(!"unsupported loop order");
        }
    });

    return status::success;
}

using namespace data_type;
template struct jit_uni_x8s8s32x_deconvolution_fwd_t<avx2>;
template struct jit_uni_x8s8s32x_deconvolution_fwd_t<sse41>;
template struct jit_uni_x8s8s32x_deconv_fwd_kernel_t<avx2>;
template struct jit_uni_x8s8s32x_deconv_fwd_kernel_t<sse41>;
template struct jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<avx2, Xbyak::Ymm>;
template struct jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<avx2, Xbyak::Xmm>;
template struct jit_uni_x8s8s32x_deconv_fwd_kernel_vmm_t<sse41, Xbyak::Xmm>;
} // namespace x64
} // namespace cpu
} // namespace impl
} // namespace dnnl