/*******************************************************************************
* Copyright 2020-2025 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 <cstdlib>
#include <functional>

#include "cpu/x64/injectors/jit_uni_postops_injector.hpp"
#include "cpu/x64/injectors/jit_uni_binary_injector.hpp"
#include "cpu/x64/jit_gemm_x8s8s32x_conv_zp_src_pad_comp.hpp"
#include "cpu/x64/jit_gemm_x8s8s32x_convolution_utils.hpp"
#include "cpu/x64/jit_generator.hpp"
#include "cpu/x64/jit_primitive_conf.hpp"

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

namespace gemm_x8s8s32x_convolution_utils {
using namespace dnnl::impl::cpu::gemm_x8s8s32x_convolution_utils;

template <cpu_isa_t isa>
struct jit_pp_ker_t : pp_ker_t, public jit_generator_t {
    DECLARE_CPU_JIT_AUX_FUNCTIONS(
            gemm_x8s8s32x_convolution_utils::jit_pp_ker_t);

    jit_pp_ker_t(const convolution_pd_t *pd, const conv_gemm_conf_t &jcp)
            : pp_ker_t(pd, jcp)
            , jit_generator_t(jit_name())
            , do_eltwise_(false)
            , do_sum_(false)
            , sum_scale_(0)
            , sum_data_type_(dnnl_f32)
            , default_OC_loop_unroll_(4)
            , max_OC_loop_unroll_(isa == avx512_core ? 12 : 6)
            , idx_compute_vreg_start_(0)
            , idx_compute_vreg_max_(isa == avx512_core ? 31 : 15)
            , compute_vregs_per_iter_(1)
    {
        if (utils::one_of(isa, avx2, sse41)) {
            idx_compute_vreg_start_ += 2;   //  Vmm(0), Vmm(1) - for masks
        }
        if (do_scale_) {
            vreg_scale = Vmm(idx_compute_vreg_start_++);
        }
        dst_data_type_size_ = types::data_type_size(dst_data_type_);
        if (dst_data_type_ == data_type::u8 || utils::one_of(isa, avx2, sse41)) {
            vreg_zero = Vmm(idx_compute_vreg_start_++);
        }
        bool only_eltwise_or_sum = true;
        bool with_binary = false;
        for (int idx = 0; idx < post_ops_.len(); ++idx) {
            const auto &e = post_ops_.entry_[idx];
            if (e.is_eltwise(true)) {
                do_eltwise_ = true;
            } else if (e.is_binary()) {
                with_binary = true;
                only_eltwise_or_sum = false;
            } else if (e.is_sum()) {
                do_sum_ = true;
                sum_scale_ = e.sum.scale;
                sum_data_type_ = e.sum.dt;
            } else {
                only_eltwise_or_sum = false;
            }
        }
        if (with_binary) {
#define PARAM_OFF(field) offsetof(ker_args_t, field)
            static constexpr bool preserve_gpr = true;
            static constexpr bool preserve_vmm = true;
            static constexpr size_t helper_vmm_idx = 1;
            static constexpr size_t tail_size = 0;
            static constexpr bool use_exact_tail_scalar_bcast = false;
            const binary_injector::rhs_arg_static_params_t rhs_sp {
                helper_vmm_idx, r13, r14, r15, preserve_gpr,
                preserve_vmm, PARAM_OFF(post_ops_binary_rhs_arg_vec),
                PARAM_OFF(dst_orig), memory_desc_wrapper(pd->dst_md()),
                tail_size, kreg_rem_mask_short, use_exact_tail_scalar_bcast};
#undef PARAM_OFF
            const binary_injector::static_params_t bsp {this->reg_param_bak, rhs_sp};
            jit_binary_injector_ = utils::make_unique<
                    binary_injector::jit_uni_binary_injector_t<isa>>(
                    this, bsp);
        }
        if (post_ops_.len() > 0 && !only_eltwise_or_sum) {
            vreg_d_weights = Vmm(idx_compute_vreg_max_--);
            vreg_d_bias = Vmm(idx_compute_vreg_max_--);
        }

        do_signed_scaling_ = jcp_.signed_input;
        if (do_signed_scaling_)
            vreg_signed_scale = Vmm(idx_compute_vreg_start_++);

        if (do_bias_) {
            bias_data_type_size_ = types::data_type_size(bias_data_type_);
            compute_vregs_per_iter_++;
        }


        if (jcp.with_dst_scale) {
            vreg_dst_scale = Vmm(idx_compute_vreg_start_++);
            // compute_vregs_per_iter_++;
        }

        if (jcp.zp.dst_exists) {
            vreg_zp_dst_common = Vmm(idx_compute_vreg_start_++);
            // compute_vregs_per_iter_++;
        }

        if (do_sum_) {
            vreg_sum_scale = Vmm(idx_compute_vreg_start_++);
            compute_vregs_per_iter_++;
        }

        for (int i = 0; i < post_ops_.len(); i++) {
            auto &post_op = post_ops_.entry_[i];
            if (post_op.is_eltwise()) {
                jit_eltwise_injectors_.push_back(new jit_uni_eltwise_injector_t<isa>(
                        this, post_op.eltwise, data_type::f32, true, eltwise_reserved, mask_post_op_reserved));
            } else if (post_op.is_depthwise()) {
                jit_depthwise_injectors_.push_back(new jit_uni_depthwise_injector_f32<isa>(
                        this, post_op, mask_post_op_reserved));
            }
        }

        int max_unroll = (idx_compute_vreg_max_ - idx_compute_vreg_start_ + 1) / compute_vregs_per_iter_;
        max_OC_loop_unroll_ = nstl::min(max_OC_loop_unroll_, max_unroll);
        default_OC_loop_unroll_ = nstl::min(default_OC_loop_unroll_, max_unroll);
    }
    ~jit_pp_ker_t() {
        for (auto inj : jit_eltwise_injectors_)
            delete inj;
        jit_eltwise_injectors_.clear();
        for (auto inj : jit_depthwise_injectors_)
            delete inj;
        jit_depthwise_injectors_.clear();
    }

    status_t create_kernel() override {
        return jit_generator_t::create_kernel();
    }
    void operator()(void *void_dst, acc_data_t *acc, const char *bias,
            const float *scales, float dst_scale, float sum_scale,
            float signed_scale, int g, size_t start, size_t end,
            const zero_point_call_params_t &zp,
            const void *post_ops_binary_rhs_arg_vec, const void *dst_orig,
            const exec_ctx_t & /* ctx */, const memory_desc_t & /* dst_md */,
            const single_gemm_conv_chunk_desc_t &chunk_desc) const override {
        if (end <= start) return;

        char *dst = (char *)void_dst;

        ker_args_t args;
        size_t oc_offset = start % OC_;
        size_t os_offset = start / OC_;
        args.acc = acc + start;
        args.dst = dst
                   + (os_offset * dst_os_stride_ + oc_offset)
                     * dst_data_type_size_;
        const ptrdiff_t g_oc_offset = g * jcp_.oc;
        const ptrdiff_t g_oc_offset_prologue = g_oc_offset + oc_offset;
        args.dst_orig = dst_orig;
        args.bias = bias + (g * jcp_.oc + oc_offset) * bias_data_type_size_;
        args.zp_src = zp.src + (jcp_.zp.src_is_common ? 0 : g_oc_offset_prologue);
        args.zp_src_comp
            = zp.src_comp ? zp.src_comp + g_oc_offset_prologue : nullptr;
        args.zp_dst = zp.dst;
        args.scales = scales + jcp_.scale_idx_mult * (g * jcp_.oc + oc_offset);
        args.dst_scale = dst_scale;
        args.sum_scale = sum_scale_;
        args.signed_scale = signed_scale;
        args.len = end - start;
        args.oc_offset = oc_offset;
        args.g_offset = g * jcp_.oc;
        args.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec;
        args.dst_orig = dst_orig;

        if (zp_pad_comp_helper_) {
            const auto hw
                = std::div(static_cast<dim_t>(os_offset), chunk_desc.w_size_);
            args.h = hw.quot + chunk_desc.h_off_;
            args.w = hw.rem + chunk_desc.w_off_;
            args.w_size = chunk_desc.w_size_ + chunk_desc.w_off_;
            args.w_off = chunk_desc.w_off_;
            args.zp_src_pad_comp = zp.src_pad_comp;
            const auto zp_src_pad_com_d
                = zp_pad_comp_helper_->calculate_zp_src_pad_com_d(
                    chunk_desc.d_off_);
            args.zp_src_pad_com_d_offset = zp_src_pad_com_d.offset;
            args.should_apply_zp_src_pad_comp_d
                = zp_src_pad_com_d.should_apply_pad_comp_d;
        }

        jit_generator_t::operator()(&args);
    }

private:
    void generate() override;

    struct ker_args_t {
        char *dst;
        const acc_data_t *acc;
        const char *bias;
        const float *scales;
        float dst_scale;
        float sum_scale;
        float signed_scale;
        size_t len;
        size_t oc_offset;
        const int32_t *zp_src;
        const int32_t *zp_dst;
        const int32_t *zp_src_comp;
        const int32_t *zp_src_pad_comp;
        size_t g_oc_offset_prologue;
        size_t g_offset;
        const void *post_ops_binary_rhs_arg_vec;
        const void *dst_orig;
        dim_t h;
        dim_t w;
        dim_t w_size;
        dim_t w_off;
        dim_t zp_src_pad_com_d_offset;
        bool should_apply_zp_src_pad_comp_d;
    };

    nstl::vector<jit_uni_eltwise_injector_t<isa> *> jit_eltwise_injectors_;
    nstl::vector<jit_uni_depthwise_injector_f32<isa> *> jit_depthwise_injectors_;
    std::unique_ptr<binary_injector::jit_uni_binary_injector_t<isa>>
            jit_binary_injector_;

    size_t number_of_reserved_zmm_regs_;
    using Vmm = typename cpu_isa_traits_t<isa>::Vmm;
    static const size_t vlen = cpu_isa_traits_t<isa>::vlen / sizeof(float);

    Xbyak::Reg64 reg_param = abi_param1;
    Xbyak::Reg64 reg_param_bak = r11;
    Xbyak::Reg64 reg_dst = rdx;
    Xbyak::Reg64 reg_acc = rax;
    Xbyak::Reg64 reg_bias = rbx;
    Xbyak::Reg64 reg_scales = rsi;
    Xbyak::Reg64 reg_g_offset = rbp;

    Xbyak::Reg64 reg_len = r8;
    Xbyak::Reg64 reg_tmp = rcx; // intentional for shifting purposes
    Xbyak::Reg64 reg_oc_offset = r9;
    Xbyak::Reg64 reg_rem_mask_short = r10;
    Xbyak::Opmask kreg_rem_mask_short = k1;

    Vmm vreg_zero, vreg_scale, vreg_sum_scale, vreg_signed_scale, vreg_comp;

    //  sse41/avx2
    Xbyak::Reg64 reg_ptr_maskmovdqu_dst = rdi; // sse41: store destination - must be rdi
    Xbyak::Label l_table;
    Xbyak::Reg64 reg_table = r12;
    Xbyak::Reg64 reg_shift_table = r13;
    Vmm vreg_mask = Vmm(0); //  sse41: mask for blendvps must be in xmm0
    Vmm vreg_store_mask = Vmm(1);

    //  post_ops
    Xbyak::Opmask mask_post_op_reserved = k2;
    Xbyak::Reg64 eltwise_reserved = rax;
    Xbyak::Reg64 reg_d_weights = r14;
    Xbyak::Reg64 reg_d_bias = r15;
    Vmm vreg_d_weights, vreg_d_bias;

    size_t dst_data_type_size_ = 0;
    size_t bias_data_type_size_ = 0;

    bool do_eltwise_;
    bool do_sum_;
    float sum_scale_;
    data_type_t sum_data_type_;
    bool do_signed_scaling_;

    int default_OC_loop_unroll_;
    int max_OC_loop_unroll_;
    int idx_compute_vreg_start_;
    int idx_compute_vreg_max_;
    int compute_vregs_per_iter_;

    Vmm vreg_dst_scale;
    Vmm vreg_zp_dst_common;

    std::unique_ptr<jit_gemm_x8s8s32x_zp_pad_comp_helper> zp_pad_comp_helper_;

    int idx_vreg_dst(int iter) {
        int idx = idx_compute_vreg_start_ + iter * compute_vregs_per_iter_ + 0;
        assert(idx <= idx_compute_vreg_max_);
        return idx;
    }
    int idx_vreg_bias(int iter) {
        int idx = idx_compute_vreg_start_ + iter * compute_vregs_per_iter_ + 1;
        assert(idx <= idx_compute_vreg_max_);
        return idx;
    }
    int idx_vreg_prev_dst(int iter) {
        int idx = idx_compute_vreg_start_ + iter * compute_vregs_per_iter_ + 2;
        assert(idx <= idx_compute_vreg_max_);
        return idx;
    }

    Vmm vreg_dst(int idx) { return Vmm(idx_vreg_dst(idx)); };
    Xbyak::Ymm ymm_dst(int idx) { return Xbyak::Ymm(idx_vreg_dst(idx)); };
    Xbyak::Xmm xmm_dst(int idx) { return Xbyak::Xmm(idx_vreg_dst(idx)); };
    Vmm vreg_bias(int idx) { return Vmm(idx_vreg_bias(idx)); };
    Vmm vreg_prev_dst(int idx) { return Vmm(idx_vreg_prev_dst(idx)); };

    Vmm get_masked_vreg_dst(int idx, bool apply_mask) {
        Vmm vreg_dst_ = vreg_dst(idx);
        if (apply_mask)
            vreg_dst_ = vreg_dst_ | kreg_rem_mask_short;
        // else
        //     vreg_dst_ = vreg_dst_ | kreg_rem_mask_vlen_;
        return vreg_dst_;
    }
};

template <cpu_isa_t isa>
void jit_pp_ker_t<isa>::generate() {
    using namespace Xbyak;
    using namespace utils;

    preamble();

    const auto &p = post_ops_;
    std::size_t post_ops_pointers_count = 0;
    for (int i = 0; i < p.len(); i++) {
        if (p.entry_[i].is_depthwise() || p.entry_[i].is_quantization()) {
            post_ops_pointers_count++;
        }
    }

#define PARAM_OFF(x) offsetof(ker_args_t, x)
    if (post_ops_pointers_count != 0) {
        sub(rsp, post_ops_pointers_count * sizeof(float *));

        auto aux_reg0 = reg_dst;
        auto aux_reg1 = reg_acc;

        mov(aux_reg0, ptr[reg_param + PARAM_OFF(post_ops_binary_rhs_arg_vec)]);
        for (size_t i = 0; i < post_ops_pointers_count; i++) {
            mov(aux_reg1, ptr[aux_reg0 + i * sizeof(float *)]);
            mov(ptr[rsp + i * sizeof(float *)], aux_reg1);
        }
    }

    mov(reg_param_bak, reg_param);
    mov(reg_dst, ptr[reg_param + PARAM_OFF(dst)]);
    mov(reg_acc, ptr[reg_param + PARAM_OFF(acc)]);
    mov(reg_bias, ptr[reg_param + PARAM_OFF(bias)]);
    mov(reg_scales, ptr[reg_param + PARAM_OFF(scales)]);
    mov(reg_len, ptr[reg_param + PARAM_OFF(len)]);
    mov(reg_oc_offset, ptr[reg_param + PARAM_OFF(oc_offset)]);
    mov(reg_g_offset, ptr[reg_param + PARAM_OFF(g_offset)]);

    if (jcp_.zp.dst_exists) {
        mov(reg_tmp, ptr[reg_param + PARAM_OFF(zp_dst)]);
        vcvtdq2ps(vreg_zp_dst_common, ptr_b[reg_tmp]);
    }

    if (jcp_.with_dst_scale)
        vbroadcastss(vreg_dst_scale, ptr[reg_param + PARAM_OFF(dst_scale)]);
    if (do_sum_)
        uni_vbroadcastss(vreg_sum_scale, ptr[reg_param + PARAM_OFF(sum_scale)]);
    if (do_signed_scaling_)
        uni_vbroadcastss(vreg_signed_scale, ptr[reg_param + PARAM_OFF(signed_scale)]);
    if (do_scale_ && jcp_.scale_idx_mult == 0)
        uni_vbroadcastss(vreg_scale, dword[reg_scales]);
#undef PARAM_OFF

    if (do_eltwise_ || dst_data_type_ == data_type::u8 || utils::one_of(isa, avx2, sse41))
        uni_vpxor(vreg_zero, vreg_zero, vreg_zero);

    if (utils::one_of(isa, avx2, sse41))
        mov(reg_table, l_table);

    auto apply_post_ops = [&](size_t offset, int idx, bool apply_mask) {
        std::size_t post_ops_data_offset = 0;
        int eltwise_inj_idx = 0;
        int depthwise_inj_idx = 0;
        int binary_inj_idx = 0;
        for (int i = 0; i < post_ops_.len(); i++) {
            auto &post_op = post_ops_.entry_[i];
            if (post_op.is_sum()) {
                auto dst_addr = ptr[reg_dst + offset * dst_data_type_size_];
                auto vreg_prev_dst_ = vreg_prev_dst(idx);
                switch (sum_data_type_) {
                    case data_type::f32:
                    case data_type::s32:
                        uni_vmovups(vreg_prev_dst_, dst_addr);
                        break;
                    case data_type::s8:
                        uni_vpmovsxbd(vreg_prev_dst_, dst_addr);
                        break;
                    case data_type::u8:
                        uni_vpmovzxbd(vreg_prev_dst_, dst_addr);
                        break;
                    default:
                        assert(!"unsupported data type");
                }
                if (sum_data_type_ != data_type::f32)
                    uni_vcvtdq2ps(vreg_prev_dst(idx), vreg_prev_dst(idx));

                uni_vfmadd231ps(vreg_dst(idx), vreg_prev_dst(idx), vreg_sum_scale);
            } else if (post_op.is_binary()) {
                binary_injector::rhs_arg_dynamic_params_t rhs_arg_params;
                auto dst_addr = ptr[reg_dst + offset * dst_data_type_size_];
                rhs_arg_params.vmm_idx_to_out_addr.emplace(idx, dst_addr);
                rhs_arg_params.vmm_idx_to_out_elem_off_val.emplace(
                        idx, 0 * sizeof(float));
                if (mayiuse(avx512_core) && apply_mask) 
                    rhs_arg_params.vmm_tail_idx_.emplace(idx);
                jit_binary_injector_->compute_vector(
                        idx, binary_inj_idx, post_op, rhs_arg_params);

                binary_inj_idx++;
            } else if (post_op.is_eltwise()) {
                jit_eltwise_injectors_[eltwise_inj_idx]->compute_vector_range(vreg_dst(idx).getIdx(),
                                                                              vreg_dst(idx).getIdx() + 1);
                eltwise_inj_idx++;
            } else if (post_op.is_depthwise()) {
                add(reg_oc_offset, reg_g_offset);

                const Xbyak::RegExp depthwise_arg_base = rsp + post_ops_data_offset;
                mov(reg_d_weights, ptr[depthwise_arg_base]);
                lea(reg_d_weights, ptr[reg_d_weights + reg_oc_offset * sizeof(float) + offset]);

                jit_depthwise_injectors_[depthwise_inj_idx]->compute_vector_range(vreg_dst(idx).getIdx(), vreg_dst(idx).getIdx() + 1, reg_d_weights, reg_d_weights);

                sub(reg_oc_offset, reg_g_offset);

                post_ops_data_offset += jit_depthwise_injectors_[depthwise_inj_idx]->memoryStep();
                depthwise_inj_idx++;
                binary_inj_idx++;
            } else if (post_op.is_quantization()) {
                add(reg_oc_offset, reg_g_offset);
                bool do_dequantization = post_op.quantization.alg == alg_kind::quantization_quantize_dequantize;
                bool do_rounding = do_dequantization || dst_data_type_ == dnnl_f32 || i != post_ops_.len() - 1;

                const Xbyak::RegExp quantization_arg_base = rsp + post_ops_data_offset;
                size_t crop_low_off = post_op.quantization.offset[post_op.quantization.crop_low] * sizeof(float);
                if (post_op.quantization.per_channel[post_op.quantization.crop_low]) {
                    mov(reg_d_weights, ptr[quantization_arg_base]);
                    uni_vmovups(vreg_d_weights, ptr[reg_d_weights + reg_oc_offset * sizeof(float) + offset * sizeof(float) + crop_low_off]);
                } else {
                    mov(reg_d_weights, ptr[quantization_arg_base]);
                    uni_vbroadcastss(vreg_d_weights, ptr[reg_d_weights + crop_low_off]);
                }

                size_t crop_high_off = post_op.quantization.offset[post_op.quantization.crop_high] * sizeof(float);
                if (post_op.quantization.per_channel[post_op.quantization.crop_high]) {
                    mov(reg_d_bias, ptr[quantization_arg_base]);
                    uni_vmovups(vreg_d_bias, ptr[reg_d_bias + reg_oc_offset * sizeof(float) + offset * sizeof(float) + crop_high_off]);
                } else {
                    mov(reg_d_bias, ptr[quantization_arg_base]);
                    uni_vbroadcastss(vreg_d_bias, ptr[reg_d_bias + crop_high_off]);
                }

                uni_vmaxps(vreg_dst(idx), vreg_dst(idx), vreg_d_weights);
                uni_vminps(vreg_dst(idx), vreg_dst(idx), vreg_d_bias);

                size_t inp_scale_off = post_op.quantization.offset[post_op.quantization.inp_scale] * sizeof(float);
                if (post_op.quantization.per_channel[post_op.quantization.inp_scale]) {
                    mov(reg_d_weights, ptr[quantization_arg_base]);
                    uni_vmovups(vreg_d_weights, ptr[reg_d_weights + reg_oc_offset * sizeof(float) + offset * sizeof(float) + inp_scale_off]);
                } else {
                    mov(reg_d_weights, ptr[quantization_arg_base ]);
                    uni_vbroadcastss(vreg_d_weights, ptr[reg_d_weights + inp_scale_off]);
                }

                size_t inp_shift_off = post_op.quantization.offset[post_op.quantization.inp_shift] * sizeof(float);
                if (post_op.quantization.per_channel[post_op.quantization.inp_shift]) {
                    mov(reg_d_bias, ptr[quantization_arg_base]);
                    uni_vmovups(vreg_d_bias, ptr[reg_d_bias + reg_oc_offset * sizeof(float) + offset * sizeof(float) + inp_shift_off]);
                } else {
                    mov(reg_d_bias, ptr[quantization_arg_base]);
                    uni_vbroadcastss(vreg_d_bias, ptr[reg_d_bias + inp_shift_off]);
                }

                uni_vfmadd213ps(vreg_dst(idx), vreg_d_weights, vreg_d_bias);

                if (do_rounding)
                    uni_vroundps(vreg_dst(idx), vreg_dst(idx), 0);

                if (do_dequantization) {
                    size_t output_scale_off = post_op.quantization.offset[post_op.quantization.output_scale] * sizeof(float);
                    if (post_op.quantization.per_channel[post_op.quantization.output_scale]) {
                        mov(reg_d_weights, ptr[quantization_arg_base ]);
                        uni_vmovups(vreg_d_weights, ptr[reg_d_weights + reg_oc_offset * sizeof(float) + offset * sizeof(float) + output_scale_off]);
                    } else {
                        mov(reg_d_weights, ptr[quantization_arg_base]);
                        uni_vbroadcastss(vreg_d_weights, ptr[reg_d_weights + output_scale_off]);
                    }

                    size_t output_shift_off = post_op.quantization.offset[post_op.quantization.output_shift] * sizeof(float);
                    if (post_op.quantization.per_channel[post_op.quantization.output_shift]) {
                        mov(reg_d_bias, ptr[quantization_arg_base]);
                        uni_vmovups(vreg_d_bias, ptr[reg_d_bias + reg_oc_offset * sizeof(float) + offset * sizeof(float) + output_shift_off]);
                    } else {
                        mov(reg_d_bias, ptr[quantization_arg_base]);
                        uni_vbroadcastss(vreg_d_bias, ptr[reg_d_bias + output_shift_off]);
                    }

                    uni_vfmadd213ps(vreg_dst(idx), vreg_d_weights, vreg_d_bias);
                }
                sub(reg_oc_offset, reg_g_offset);

                post_ops_data_offset += sizeof(float*);
                binary_inj_idx++;
            }
        }
    };

    // Load accumulated value, convert to float,
    // multiply weight scale, bias (if any), and simple operations (if any);
    // then convert to destination type and store
    auto compute = [&](size_t offset, int idx, bool apply_mask) {
        auto acc_addr = ptr[reg_acc + offset * sizeof(acc_data_t)];

        if (do_scale_ && jcp_.scale_idx_mult > 0) {
            assert(jcp_.scale_idx_mult == 1);
            auto scale_addr = ptr[reg_scales + offset * sizeof(float)];
            auto vreg_scale_ = vreg_scale;
            if (isa == avx512_core) {
                if (apply_mask)
                    vreg_scale_ = vreg_scale_ | kreg_rem_mask_short;
                uni_vmovups(vreg_scale_, scale_addr);
            } else {
                if (apply_mask)
                    if (isa != sse41) {
                        uni_vblendvps(vreg_scale, vreg_zero, scale_addr, vreg_mask);
                    } else {
                        uni_vmovups(vreg_scale, vreg_zero);
                        uni_vblendvps(vreg_scale, vreg_scale, scale_addr, vreg_mask);
                    }
                else
                    uni_vmovups(vreg_scale, scale_addr);
            }
        }

        auto vreg_dst_ = vreg_dst(idx);
        if (isa == avx512_core) {
            if (apply_mask)
                vreg_dst_ = vreg_dst_ | kreg_rem_mask_short;
            uni_vcvtdq2ps(vreg_dst_, acc_addr);
        } else {
            if (apply_mask) {
                if (isa != sse41) {
                    uni_vblendvps(vreg_dst_, vreg_zero, acc_addr, vreg_mask);
                } else {
                    uni_vmovups(vreg_dst_, acc_addr);
                }
                uni_vcvtdq2ps(vreg_dst_, vreg_dst_);
            } else {
                if (isa == sse41) {
                    uni_vmovups(vreg_dst_, acc_addr);
                    uni_vcvtdq2ps(vreg_dst_, vreg_dst_);
                } else {
                    uni_vcvtdq2ps(vreg_dst_, acc_addr);
                }
            }
        }

        if (do_signed_scaling_)
            uni_vmulps(vreg_dst(idx), vreg_dst(idx), vreg_signed_scale);

        if (do_scale_)
            uni_vmulps(vreg_dst(idx), vreg_dst(idx), vreg_scale);
        if (do_bias_) {
            auto bias_addr = ptr[reg_bias + offset * bias_data_type_size_];
            auto vreg_bias_ = vreg_bias(idx);
            if (isa == avx512_core && apply_mask)
                vreg_bias_ = vreg_bias_ | kreg_rem_mask_short;

            switch (bias_data_type_) {
                case data_type::s8:
                    uni_vpmovsxbd(vreg_bias_, bias_addr);
                    break;
                case data_type::u8:
                    uni_vpmovzxbd(vreg_bias_, bias_addr);
                    break;
                case data_type::s32:
                case data_type::f32:
                    uni_vmovups(vreg_bias_, bias_addr);
                    break;
                default:
                    assert(!"unimplemented");
            }
            if (bias_data_type_ != data_type::f32)
                uni_vcvtdq2ps(vreg_bias(idx), vreg_bias(idx));
            uni_vaddps(vreg_dst(idx), vreg_dst(idx), vreg_bias(idx));
        }

        apply_post_ops(offset, idx, apply_mask);

        if (jcp_.with_dst_scale) {
            uni_vmulps(vreg_dst_, vreg_dst(idx), vreg_dst_scale);
        }

        if (jcp_.zp.dst_exists) {
            vaddps(vreg_dst_, vreg_dst(idx), vreg_zp_dst_common);
        }

        if (dst_data_type_ != data_type::f32) {
            if (isa == avx512_core) {
                auto rmode_control = T_rn_sae;
                vcvtps2dq(vreg_dst(idx) | rmode_control, vreg_dst(idx));
            } else {
                uni_vcvtps2dq(vreg_dst(idx), vreg_dst(idx));
            }
        }

        if (dst_data_type_ == data_type::u8)
            uni_vpmaxsd(vreg_dst(idx), vreg_dst(idx), vreg_zero);

        auto dst_addr = ptr[reg_dst + offset * dst_data_type_size_];
        switch (dst_data_type_) {
            case data_type::s8:
                if (isa == avx512_core) {
                    vpmovsdb(dst_addr, vreg_dst_);
                } else {
                    uni_vpackssdw(vreg_dst_, vreg_dst_, vreg_dst_);
                    if (isa != sse41)
                        vpermq(ymm_dst(idx), ymm_dst(idx), 0x08);
                    uni_vpacksswb(vreg_dst_, vreg_dst_, vreg_dst_);
                    if (apply_mask) {
                        lea(reg_ptr_maskmovdqu_dst, dst_addr);
                        maskmovdqu(vreg_dst_, vreg_store_mask);
                    } else {
                        if (isa != sse41) {
                            vmovq(dst_addr, xmm_dst(idx));
                        } else {
                            movd(dst_addr, xmm_dst(idx));
                        }
                    }
                }
                break;
            case data_type::u8:
                if (isa == avx512_core) {
                    vpmovusdb(dst_addr, vreg_dst_);
                } else {
                    uni_vpackusdw(vreg_dst_, vreg_dst_, vreg_dst_);
                    if (isa != sse41)
                        vpermq(ymm_dst(idx), ymm_dst(idx), 0x08);
                    uni_vpackuswb(vreg_dst_, vreg_dst_, vreg_dst_);
                    if (apply_mask) {
                        lea(reg_ptr_maskmovdqu_dst, dst_addr);
                        maskmovdqu(vreg_dst_, vreg_store_mask);
                    } else {
                        if (isa != sse41) {
                            vmovq(dst_addr, xmm_dst(idx));
                        } else {
                            movd(dst_addr, xmm_dst(idx));
                        }
                    }
                }
                break;
            case data_type::f32:
            case data_type::s32:
                if (isa == avx512_core) {
                    uni_vmovups(dst_addr, vreg_dst_);
                } else {
                    if (apply_mask) {
                        if (isa != sse41) {
                            vmaskmovps(dst_addr, vreg_mask, vreg_dst_);
                        } else {
                            lea(reg_ptr_maskmovdqu_dst, dst_addr);
                            maskmovdqu(vreg_dst_, vreg_mask);
                        }
                    } else {
                        uni_vmovups(dst_addr, vreg_dst_);
                    }
                }
                break;
            default:
                assert(!"unimplemented");
        }
    };

    // Advance all pointers by an immediate
    auto advance_ptrs_imm = [&](size_t offset) {
        add(reg_dst, offset * dst_data_type_size_);
        add(reg_acc, offset * sizeof(acc_data_t));
        if (jcp_.scale_idx_mult) {
            assert(jcp_.scale_idx_mult == 1);
            add(reg_scales, offset * sizeof(float));
        }
        if (do_bias_)
            add(reg_bias, offset * bias_data_type_size_);
    };

    // Advance all pointers by a value stored in a register
    auto advance_ptrs_reg = [&](Reg64 offset) {
        lea(reg_dst, ptr[reg_dst + offset * dst_data_type_size_]);
        lea(reg_acc, ptr[reg_acc + offset * sizeof(acc_data_t)]);
        if (jcp_.scale_idx_mult) {
            assert(jcp_.scale_idx_mult == 1);
            lea(reg_scales, ptr[reg_scales + offset * sizeof(float)]);
        }
        if (do_bias_)
            lea(reg_bias, ptr[reg_bias + offset * bias_data_type_size_]);
    };

    // Rewind pointers that point to data that is indexed by output channel
    // (bias or per-oc scaling factors)
    auto rewind_ptrs = [&]() {
        if (do_bias_)
            sub(reg_bias, OC_ * bias_data_type_size_);
        if (jcp_.scale_idx_mult) {
            assert(jcp_.scale_idx_mult == 1);
            sub(reg_scales, OC_ * sizeof(float));
        }
        add(reg_dst, (dst_os_stride_ - OC_) * dst_data_type_size_);
    };

    //                    <--------- OC --------------->
    //
    // ^  ................+..............+-------------+.......................
    // |  .               : not accessed |Prologue loop|                      .
    // |  .               +--------------+-------------+                      .
    //    .               |                            |                      .
    // O  .               |  Main loop (unrolled)      |                      .
    // S  .               |                            |                      .
    //    .               +--------------+-------------+                      .
    // |  .               | Epilogue loop|not accessed :                      .
    // v  ................+--------------+.............+.......................

    bool do_post_ops = post_ops_.len() != 0;

    Label prologue_end;
    cmp(reg_oc_offset, 0);
    je(prologue_end, T_NEAR);

    // Prologue loop
    {
        mov(reg_tmp, OC_);
        sub(reg_tmp, reg_oc_offset);
        cmp(reg_tmp, reg_len);
        cmovg(reg_tmp, reg_len);
        sub(reg_len, reg_tmp);

        Label prologue_loop, prologue_loop_tail, prologue_loop_end;
        cmp(reg_tmp, vlen);
        jl(prologue_loop_tail, T_NEAR);
        L(prologue_loop);
        {
            compute(0, 0, false);
            advance_ptrs_imm(vlen);
            if (do_post_ops)
                add(reg_oc_offset, vlen);
            sub(reg_tmp, vlen);
            cmp(reg_tmp, vlen);
            jge(prologue_loop, T_NEAR);
        }

        L(prologue_loop_tail);
        if (isa == avx512_core) {
            mov(reg_rem_mask_short, 1);
            // cl == reg_tmp because reg_tmp <= vlen here
            shl(reg_rem_mask_short, cl);
            sub(reg_rem_mask_short, 1);
            jz(prologue_loop_end, T_NEAR);

            kmovq(kreg_rem_mask_short, reg_rem_mask_short);
        } else {
            mov(reg_shift_table, vlen);
            sub(reg_shift_table, reg_tmp);
            uni_vmovups(vreg_mask, ptr[reg_table + reg_shift_table * sizeof(float)]);
            if (dst_data_type_ == data_type::s8 || dst_data_type_ == data_type::u8) {
                mov(reg_shift_table, vlen * sizeof(float));
                sub(reg_shift_table, reg_tmp);
                uni_vmovups(vreg_store_mask, ptr[reg_table + reg_shift_table]);
            }
        }
        compute(0, 0, true);
        advance_ptrs_reg(reg_tmp);

        L(prologue_loop_end);
        rewind_ptrs();
    }
    L(prologue_end);

    // Main loop
    Label main_loop_end;
    {
        cmp(reg_len, OC_);
        jl(main_loop_end, T_NEAR);

        size_t OC_loop, OC_tail;
        if (OC_ < max_OC_loop_unroll_ * vlen) {
            // Fully unroll small loops
            OC_loop = 0;
            OC_tail = OC_;
        } else {
            OC_loop = vlen * default_OC_loop_unroll_;
            OC_tail = OC_ % OC_loop;
        }

        assert(!!OC_loop || !!OC_tail);

        if (OC_tail % vlen) {
            int vlen_tail = OC_tail % vlen;
            if (isa == avx512_core) {
                unsigned tail_mask = (1 << vlen_tail) - 1;
                mov(reg_tmp, tail_mask);
                kmovq(kreg_rem_mask_short, reg_tmp);
            } else {
                mov(reg_shift_table, vlen - vlen_tail);
                uni_vmovups(vreg_mask, ptr[reg_table + reg_shift_table * sizeof(float)]);
                if (dst_data_type_ == data_type::s8 || dst_data_type_ == data_type::u8) {
                    mov(reg_shift_table, vlen * sizeof(float));
                    sub(reg_shift_table, vlen_tail);
                    uni_vmovups(vreg_store_mask, ptr[reg_table + reg_shift_table]);
                }
            }
        }

        Label main_loop;
        L(main_loop);
        {
            if (do_post_ops)
                mov(reg_oc_offset, 0);

            if (OC_loop) {
                mov(reg_tmp, rnd_dn(OC_, OC_loop));
                Label oc_loop;
                L(oc_loop);
                {
                    for (size_t offset = 0; offset < OC_loop; offset += vlen)
                        compute(offset, offset / vlen, false);
                    advance_ptrs_imm(OC_loop);
                    if (do_post_ops)
                        add(reg_oc_offset, OC_loop);
                    sub(reg_tmp, OC_loop);
                    jnz(oc_loop);
                }
            }

            if (OC_tail) {
                for (size_t offset = 0; offset < OC_tail; offset += vlen) {
                    bool use_mask = (offset + vlen) > OC_tail;
                    compute(offset, offset / vlen, use_mask);
                }
                advance_ptrs_imm(OC_tail);
            }

            rewind_ptrs();
            sub(reg_len, OC_);
            cmp(reg_len, OC_);
            jge(main_loop, T_NEAR);
        }
    }
    L(main_loop_end);

    // Epilogue loop
    Label epilogue_end;
    {
        cmp(reg_len, 0);
        je(epilogue_end, T_NEAR);

        Label epilogue_loop, epilogue_loop_tail;
        if (do_post_ops)
            mov(reg_oc_offset, 0);
        cmp(reg_len, vlen);
        jl(epilogue_loop_tail, T_NEAR);
        L(epilogue_loop);
        {
            compute(0, 0, false);
            sub(reg_len, vlen);
            advance_ptrs_imm(vlen);
            if (do_post_ops)
                add(reg_oc_offset, vlen);
            cmp(reg_len, vlen);
            jge(epilogue_loop, T_NEAR);
        }

        L(epilogue_loop_tail);
        mov(reg_tmp, reg_len); // reg_tmp is rcx, and we need cl for the shift
        if (isa == avx512_core) {
            mov(reg_rem_mask_short, 1);
            shl(reg_rem_mask_short, cl); // reg_tmp == rcx and reg_tail < vlen
            sub(reg_rem_mask_short, 1);
            jz(epilogue_end, T_NEAR);
            kmovq(kreg_rem_mask_short, reg_rem_mask_short);
        } else {
            mov(reg_shift_table, vlen);
            sub(reg_shift_table, reg_tmp);
            uni_vmovups(vreg_mask, ptr[reg_table + reg_shift_table * sizeof(float)]);
            if (dst_data_type_ == data_type::s8 || dst_data_type_ == data_type::u8) {
                mov(reg_shift_table, vlen * sizeof(float));
                sub(reg_shift_table, reg_tmp);
                uni_vmovups(vreg_store_mask, ptr[reg_table + reg_shift_table]);
            }
        }
        compute(0, 0, true);
    }

    L(epilogue_end);

    if (post_ops_pointers_count != 0) {
        add(rsp, post_ops_pointers_count * sizeof(float *));
    }

    postamble();

    for (auto &inj : jit_eltwise_injectors_)
        inj->prepare_table(true);

    if (utils::one_of(isa, avx2, sse41)) {
        align(64);
        L(l_table);
        for (size_t i = 0; i < vlen; i++) dd(0xFFFFFFFF);
        for (size_t i = 0; i < vlen; i++) dd(0x00000000);
    }
}

bool mayiuse_jit_pp_kernel(data_type_t dst_dt) noexcept {
    const auto is_bf16_dst_dt = dst_dt == data_type::bf16;
    return mayiuse(avx512_core) && !is_bf16_dst_dt;

}

pp_ker_t *jit_pp_ker_create(
        const convolution_pd_t *pd, const conv_gemm_conf_t &jcp) {
    if (mayiuse(avx512_core)) {
        return new jit_pp_ker_t<avx512_core>(pd, jcp);
    } else if (mayiuse(avx2)) {
        return new jit_pp_ker_t<avx2>(pd, jcp);
    } else if (mayiuse(sse41)) {
        return new jit_pp_ker_t<sse41>(pd, jcp);
    }
    return nullptr;
}

bool post_ops_ok(const post_ops_t &post_ops, const memory_desc_wrapper *dst_d) {
    using namespace x64::injector;
    static constexpr bool sum_at_pos_0_only = true;
    static constexpr bool sum_requires_scale_one = false;
    return mayiuse_jit_pp_kernel(dst_d->data_type())
            && dnnl::impl::cpu::x64::injector::post_ops_ok(
                    {avx512_core, {binary, eltwise, sum}, post_ops, dst_d,
                            sum_at_pos_0_only, sum_requires_scale_one});
}

} // namespace gemm_x8s8s32x_convolution_utils
} // namespace x64
} // namespace cpu
} // namespace impl
} // namespace dnnl
