/*******************************************************************************
* Copyright 2019-2021 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 <assert.h>
#include <functional>

#include "common/c_types_map.hpp"
#include "common/dnnl_thread.hpp"
#include "common/math_utils.hpp"
#include "common/nstl.hpp"
#include "common/type_helpers.hpp"
#include "common/utils.hpp"

#include "cpu/x64/injectors/jit_uni_postops_injector.hpp"
#include "cpu/x64/jit_avx512_core_bf16cvt.hpp"
#include "cpu/x64/jit_generator.hpp"
#include "cpu/x64/jit_uni_binary.hpp"

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

#define PARAM_OFF(x) offsetof(call_params_t, x)

static bcast_set_t get_supported_bcast_strategies() {
    return {broadcasting_strategy_t::scalar, broadcasting_strategy_t::per_oc,
            broadcasting_strategy_t::per_oc_spatial};
}

template <data_type_t src_type>
bool jit_uni_binary_t<src_type>::post_ops_ok(
        const primitive_attr_t *attr, const memory_desc_wrapper &dst_d) {
    using namespace primitive_kind;

    const auto &p = attr->post_ops_;
    const auto is_eltwise = [&](int idx) {
        if (p.entry_[idx].is_eltwise()) {
            const auto alg = p.entry_[idx].eltwise.alg;
            return eltwise_injector::is_alg_supported(alg);
        }
        return false;
    };
    const auto is_binary = [&](int idx) { return p.entry_[idx].is_binary(); };
    const auto is_binary_bf16 = [&](int idx) {
        return is_binary(idx)
                && p.entry_[idx].binary.src1_desc.data_type == data_type::bf16;
    };
    const bool is_avx512_core = mayiuse(avx512_core);

    for (int i = 0; i < p.len(); i++) {
        if (p.contain(primitive_kind::sum, i)) {
            if (i > 0) return false;
        } else if (!(is_eltwise(i) || is_binary(i))
                || (!is_avx512_core && is_binary_bf16(i)))
            return false;
    }

    const int vlen = is_avx512_core ? cpu_isa_traits<avx512_core>::vlen
                                    : cpu_isa_traits<avx2>::vlen;
    const auto supported_strategies = get_supported_bcast_strategies();
    const bool postops_per_oc_broadcast_exists
            = binary_injector::any_binary_postop_rhs_per_oc_broadcast(
                    p, dst_d, supported_strategies);
    const int blksize = vlen / sizeof(float);

    const bool blocked_format = !dst_d.is_plain() && dst_d.is_blocking_desc();

    if (postops_per_oc_broadcast_exists && blocked_format) {
        /*
         * check blocking_desc consistency, currently when among postops exists
         * per_oc broadcast, binary kernel doesn't support situations when blocked
         * format size is smaller then vlen. example: sse41 vlen size is 4 and format
         * is nChw8c - not supported, avx2 vlen size is 8 and format is
         * nChw8c - supported.
         */
        const auto blocking_desc = dst_d.blocking_desc();
        if (blocking_desc.inner_nblks != 1
                || blocking_desc.inner_blks[0] != blksize
                || blocking_desc.inner_idxs[0] != 1)
            return false;
    }

    return binary_injector::binary_args_broadcast_supported(
                   p, dst_d, get_supported_bcast_strategies())
            && IMPLICATION(postops_per_oc_broadcast_exists,
                    binary_injector::all_binary_postop_rhs_per_oc_broadcast(p,
                            dst_d, supported_strategies,
                            [&dst_d](const memory_desc_wrapper &rhs_arg_md) {
                                return IMPLICATION(!mayiuse(avx2),
                                        dst_d.consistent_with(rhs_arg_md)
                                                || dst_d.is_plain());
                            }));
}

using namespace Xbyak;

enum class op_t : unsigned { none, c_blocked, n_spatial_c, n_c_spatial };

enum class bcast_t : unsigned {
    none, // tensor operation
    scalar,
    per_c,
    per_w
};

static op_t get_op_type(const memory_desc_wrapper &src0_d) {
    const auto &strides = src0_d.blocking_desc().strides;
    const auto ndims = src0_d.ndims();

    if (!src0_d.is_plain())
        return op_t::c_blocked;
    else if (strides[1] == 1)
        return op_t::n_spatial_c;
    else if (strides[0] >= strides[1]
            && IMPLICATION(ndims >= 3, strides[1] >= strides[2]))
        return op_t::n_c_spatial;
    return op_t::none;
}

static bcast_t get_bcast_type(
        const memory_desc_wrapper &src1_d, const dims_t &bcast_dims) {
    const auto ndims = src1_d.ndims();

    if (src1_d.nelems() == 1)
        return bcast_t::scalar;
    else if (ndims >= 3 && bcast_dims[1] == 1)
        return bcast_t::per_w;
    else
        return bcast_t::per_c;
}

struct binary_kernel_t : public jit_generator {
    struct call_params_t {
        // keep all sizes at 8 bytes -- jit code expects this
        const void *src0, *src1, *dst;
        const float *scales_src0, *scales_src1;
        size_t spat_offt_count;
        const void *post_ops_binary_rhs_arg_vec;
        size_t oc_l_off;
    };

    binary_kernel_t(int vlen) : vlen_(vlen), simd_w_(vlen / sizeof(float)) {}
    ~binary_kernel_t() override = default;

    void operator()(binary_kernel_t::call_params_t *p) {
        jit_generator::operator()(p);
    }

    int simd_w() const noexcept { return simd_w_; }
    int vlen() const noexcept { return vlen_; }
    op_t op_type() const noexcept { return op_type_; }
    bcast_t bcast_type() const noexcept { return bcast_type_; }

protected:
    const int vlen_;
    const size_t simd_w_;

    op_t op_type_ = op_t::none;
    bcast_t bcast_type_ = bcast_t::none;
};

template <cpu_isa_t isa>
struct jit_uni_binary_kernel_t : public binary_kernel_t {
    DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_uni_binary_kernel_t)

    using Vmm = typename cpu_isa_traits<isa>::Vmm;
    const AddressFrame &vmmword
            = (isa == sse41) ? xword : ((isa == avx2) ? yword : zword);

    const binary_pd_t *pd_;
    bool is_bf16_;
    const bool is_avx512 = utils::one_of(isa, avx512_core, avx512_core_bf16);
    const Reg64 &reg_param_ = abi_param1;
    const Reg64 &reg_src0_ = r8;
    const Reg64 &reg_src1_ = r9;
    const Reg64 &reg_dst_ = r10;
    const Reg64 &reg_offt_src0_ = r11;
    const Reg64 &reg_offt_src0_count_ = r12;
    const Reg64 &reg_offt_src1_ = rax;
    const Reg64 &reg_reverse_spat_offt_ = r13;
    const Reg64 &reg_tmp_ = r14;
    const Reg64 &reg_elt_inj_table_ = r15;
    const Reg64 &reg_off_rhs_postops_ = rdx;
    const Reg64 reg_scales_src0_ = rbx;
    const Reg64 reg_scales_src1_ = rbp;
    const Opmask &tail_opmask_ = k2;
    const Opmask &cmp_mask = k4;
    const Xmm xsum_scale_ = Xmm(15);
    const Vmm vreg_one_ = Vmm(is_avx512 ? 25 : 13);
    const Vmm vbcast_src1_ = Vmm(is_avx512 ? 30 : 14);
    const Vmm vsum_scale_ = Vmm(is_avx512 ? 31 : 15);
    const Vmm vreg_scales_src0_ = Vmm(is_avx512 ? 17 : 9);
    const Vmm vreg_scales_src1_ = Vmm(is_avx512 ? 18 : 10);

    size_t unroll_regs_ = is_avx512 ? 8 : 4;
    size_t tail_size_ = 0;
    size_t data_type_size_ = 0;
    bool do_scale_src0_ = false;
    bool do_scale_src1_ = false;
    bool do_sum_ = false;
    bool with_eltwise_ = false;
    float sum_scale_ = 0.f;
    size_t offt_src0_ = 0;
    size_t offt_src1_ = 0;
    bool use_stride_src1_ = false;
    bool broadcast_src1_value_ = false;
    bool use_stride_rhs_postops_ = false;

    static constexpr cpu_isa_t inject_isa
            = isa == avx512_core_bf16 ? avx512_core : isa;
    bool is_tail_kernel_ = false;
    std::unique_ptr<injector::jit_uni_postops_injector_t<inject_isa>>
            postops_injector_;
    const Opmask &elt_inj_opmask_ = k1;

    void init() {
        const memory_desc_wrapper src0_d(pd_->src_md(0));
        const memory_desc_wrapper src1_d(pd_->src_md(1));
        const auto &bcast_dims = pd_->broadcast_dims();

        bcast_type_ = pd_->is_tensor_op() ? bcast_t::none
                                          : get_bcast_type(src1_d, bcast_dims);
        op_type_ = get_op_type(src0_d);
        assert(op_type_ != op_t::none);
        is_bf16_ = src0_d.data_type() == data_type::bf16;
        data_type_size_ = is_bf16_ ? sizeof(bfloat16_t) : sizeof(float);

        const auto &po = pd_->attr()->post_ops_;
        const bool postops_per_oc_broadcast_exists
                = binary_injector::any_binary_postop_rhs_per_oc_broadcast(
                        po, src0_d, get_supported_bcast_strategies());
        broadcast_src1_value_ = (op_type_ == op_t::n_c_spatial
                                        && bcast_type_ == bcast_t::per_c)
                || (utils::one_of(op_type_, op_t::n_spatial_c, op_t::c_blocked)
                        && bcast_type_ == bcast_t::per_w)
                || bcast_type_ == bcast_t::scalar;
        use_stride_src1_ = !broadcast_src1_value_
                && (bcast_type_ == bcast_t::none
                        || (op_type_ == op_t::n_spatial_c
                                && bcast_type_ == bcast_t::per_c)
                        || (op_type_ == op_t::n_c_spatial
                                && bcast_type_ == bcast_t::per_w));
        use_stride_rhs_postops_ = postops_per_oc_broadcast_exists
                && op_type_ == op_t::n_spatial_c;

        tail_size_ = get_tail_size(src0_d, postops_per_oc_broadcast_exists);

        do_scale_src0_ = !pd_->attr()
                                  ->scales_.get(DNNL_ARG_SRC_0)
                                  .has_default_values();
        do_scale_src1_ = !pd_->attr()
                                  ->scales_.get(DNNL_ARG_SRC_1)
                                  .has_default_values();

        offt_src0_ = vlen_ / (is_bf16_ ? 2 : 1);
        offt_src1_ = use_stride_src1_ ? offt_src0_ : 0;
        do_sum_ = po.contain(primitive_kind::sum, 0)
                && po.entry_[0].sum.scale != 0.f;
        sum_scale_ = do_sum_ ? po.entry_[0].sum.scale : 0.f;
        with_eltwise_ = po.find(primitive_kind::eltwise) != -1;
        const bool with_binary = po.find(primitive_kind::binary) != -1;
        const bool with_postops = with_binary || with_eltwise_;

        if (with_postops) init_post_ops_injector();
    }

    std::size_t get_tail_size(const memory_desc_wrapper &src0_d,
            bool postops_per_oc_broadcast_exists) const {
        const auto &dims = src0_d.dims();
        const auto &ndims = src0_d.ndims();

        dim_t nelems = 0;

        if (op_type_ == op_t::c_blocked
                && (is_tail_kernel_ || bcast_type_ == bcast_t::per_w))
            nelems = dims[1];
        else if (bcast_type_ == bcast_t::none
                && !postops_per_oc_broadcast_exists)
            nelems = src0_d.nelems(true);
        else {
            if (op_type_ == op_t::n_spatial_c)
                nelems = dims[1];
            else if (op_type_ == op_t::n_c_spatial && ndims >= 3)
                nelems = bcast_type_ == bcast_t::per_w
                        ? dims[ndims - 1]
                        : utils::array_product(dims + 2, ndims - 2);
        }
        // it's float due to for bfloat16 we still load 16 elements, not 32.
        return nelems % simd_w_;
    }

    void init_post_ops_injector() {
        const memory_desc_wrapper src0_d(pd_->src_md(0));
        const auto &po = pd_->attr()->post_ops_;

        const eltwise_injector::static_params_t esp(true /*save_state*/,
                reg_elt_inj_table_, elt_inj_opmask_, true /*is_fwd*/,
                false /*use_dst*/);
        const binary_injector::rhs_arg_static_params_t rhs_arg_bsp {10,
                reg_tmp_, reg_elt_inj_table_, true /*preserve gpr*/,
                true /*preserve vmm*/, PARAM_OFF(post_ops_binary_rhs_arg_vec),
                src0_d, tail_size_, tail_opmask_,
                false /*use_exact_tail_scalar_bcast*/};
        const binary_injector::static_params_t bsp(
                this->param1, get_supported_bcast_strategies(), rhs_arg_bsp);

        postops_injector_ = utils::make_unique<
                injector::jit_uni_postops_injector_t<inject_isa>>(
                this, po, bsp, esp);
    }

    void apply_postops(int unroll, bool tail) {
        binary_injector::rhs_arg_dynamic_params_t rhs_arg_params;
        for (int vmm_idx = 1; vmm_idx < unroll + 1; vmm_idx++) {
            if (utils::one_of(op_type_, op_t::c_blocked, op_t::n_c_spatial)) {
                rhs_arg_params.vmm_idx_to_oc_elem_off_addr.emplace(
                        vmm_idx, ptr[param1 + PARAM_OFF(oc_l_off)]);
            } else if (op_type_ == op_t::n_spatial_c) {
                rhs_arg_params.vmm_idx_to_oc_off_oprnd.emplace(
                        vmm_idx, reg_off_rhs_postops_);
                rhs_arg_params.vmm_idx_to_oc_elem_off_val.emplace(
                        vmm_idx, (vmm_idx - 1) * static_cast<int>(simd_w_));
            }
            if (tail) rhs_arg_params.vmm_tail_idx_.emplace(vmm_idx);
        }
        postops_injector_->compute_vector_range(1, unroll + 1, rhs_arg_params);
    }

    void load_kernel_params() {
        mov(reg_tmp_, float2int(sum_scale_));
        uni_vmovq(xsum_scale_, reg_tmp_);
        uni_vbroadcastss(vsum_scale_, xsum_scale_);
        mov(reg_offt_src0_count_, ptr[reg_param_ + PARAM_OFF(spat_offt_count)]);
        mov(reg_src0_, ptr[reg_param_ + PARAM_OFF(src0)]);
        mov(reg_src1_, ptr[reg_param_ + PARAM_OFF(src1)]);
        mov(reg_dst_, ptr[reg_param_ + PARAM_OFF(dst)]);
        if (do_scale_src0_)
            mov(reg_scales_src0_, ptr[reg_param_ + PARAM_OFF(scales_src0)]);
        if (do_scale_src1_)
            mov(reg_scales_src1_, ptr[reg_param_ + PARAM_OFF(scales_src1)]);
    }

    Address src0_ptr(size_t offt = 0) {
        return vmmword[reg_src0_ + reg_offt_src0_ + offt];
    }

    Address src1_ptr(size_t offt = 0) {
        return vmmword[reg_src1_ + reg_offt_src1_ + offt];
    }

    Address dst_ptr(size_t offt = 0) {
        return vmmword[reg_dst_ + reg_offt_src0_ + offt];
    }

    unsigned int cmp_predicate(alg_kind_t alg) {
        using namespace alg_kind;
        switch (alg) {
            case binary_ge: return _cmp_nlt_us;
            case binary_gt: return _cmp_nle_us;
            case binary_le: return _cmp_le_os;
            case binary_lt: return _cmp_lt_os;
            case binary_eq: return _cmp_eq_oq;
            case binary_ne: return _cmp_neq_uq;
            default: assert(!"not supported operation!"); return -1;
        }
    }

    void perform_op(const Vmm &v0, const Vmm &v1, const Vmm &s_src0,
            const Vmm &s_src1) {
        using namespace alg_kind;
        const auto alg = pd_->desc()->alg_kind;
        const bool cmp_op = utils::one_of(alg, alg_kind::binary_ge,
                alg_kind::binary_gt, alg_kind::binary_le, alg_kind::binary_lt,
                alg_kind::binary_eq, alg_kind::binary_ne);
        if (do_scale_src0_) uni_vmulps(v0, v0, s_src0);
        if (do_scale_src1_ && offt_src1_ != 0 && !broadcast_src1_value_)
            uni_vmulps(v1, v1, s_src1);

        if (alg == binary_add)
            uni_vaddps(v0, v0, v1);
        else if (alg == binary_mul)
            uni_vmulps(v0, v0, v1);
        else if (alg == binary_max)
            uni_vmaxps(v0, v0, v1);
        else if (alg == binary_min)
            uni_vminps(v0, v0, v1);
        else if (alg == binary_div)
            uni_vdivps(v0, v0, v1);
        else if (alg == binary_sub)
            uni_vsubps(v0, v0, v1);
        else if (cmp_op) {
            const unsigned int predicate = cmp_predicate(alg);
            if (is_avx512) {
                vcmpps(cmp_mask, v0, v1, predicate);
                vmovups(v0 | cmp_mask | T_z, vreg_one_);
            } else {
                uni_vcmpps(v0, v0, v1, predicate);
                uni_vminps(v0, v0, vreg_one_);
            }
        } else
            assert(!"not supported operation!");
    }

    virtual void prepare_isa_subkernel() = 0;
    virtual void compute_bcast(bool tail) = 0;
    virtual void compute_dst(int unroll, bool tail) = 0;

    void forward() {
        Label unroll_loop, unroll_loop_tail, nelems_tail, end;

        if (do_scale_src0_)
            uni_vbroadcastss(vreg_scales_src0_, dword[reg_scales_src0_]);

        // reverse spat_offt to dispatch between labels
        mov(reg_reverse_spat_offt_, reg_offt_src0_count_);
        xor_(reg_offt_src0_,
                reg_offt_src0_); // offt_src0 to get addr of src0/dst
        xor_(reg_offt_src1_, reg_offt_src1_); // offt_src1 to get addr of src1
        if (use_stride_rhs_postops_)
            xor_(reg_off_rhs_postops_, reg_off_rhs_postops_);
        const size_t vec_size = simd_w_ * data_type_size_;
        const auto alg = pd_->desc()->alg_kind;

        if (utils::one_of(alg, alg_kind::binary_ge, alg_kind::binary_gt,
                    alg_kind::binary_le, alg_kind::binary_lt,
                    alg_kind::binary_eq, alg_kind::binary_ne)) {
            Xmm xreg_one = Xmm(vreg_one_.getIdx());
            mov(reg_tmp_, float2int(1));
            uni_vmovq(xreg_one, reg_tmp_);
            uni_vbroadcastss(vreg_one_, xreg_one);
        }

        compute_bcast(false); // bcast/load vreg just one time per a kernel call

        // used in c_blocked strategy for last blocked if tail exists
        const bool treat_each_compute_step_as_tail
                = is_tail_kernel_ && tail_size_;

        if (do_scale_src1_) {
            uni_vbroadcastss(vreg_scales_src1_, dword[reg_scales_src1_]);
            if (broadcast_src1_value_ || offt_src1_ == 0)
                uni_vmulps(vbcast_src1_, vbcast_src1_, vreg_scales_src1_);
        }

        L(unroll_loop);
        {
            const size_t offt = unroll_regs_ * vec_size;
            const size_t offt_elems = unroll_regs_ * simd_w_;
            cmp(reg_reverse_spat_offt_, offt);
            jl(unroll_loop_tail, T_NEAR);

            compute_dst(unroll_regs_, treat_each_compute_step_as_tail);
            sub(reg_reverse_spat_offt_, offt);
            add(reg_offt_src0_, offt);
            if (use_stride_src1_) add(reg_offt_src1_, offt);
            if (use_stride_rhs_postops_) add(reg_off_rhs_postops_, offt_elems);
            jmp(unroll_loop);
        }

        L(unroll_loop_tail);
        {
            cmp(reg_reverse_spat_offt_, vec_size);
            jl(nelems_tail, T_NEAR);

            compute_dst(1, treat_each_compute_step_as_tail);
            sub(reg_reverse_spat_offt_, vec_size);
            add(reg_offt_src0_, vec_size);
            if (use_stride_src1_) add(reg_offt_src1_, vec_size);
            if (use_stride_rhs_postops_) add(reg_off_rhs_postops_, simd_w_);

            jmp(unroll_loop_tail);
        }

        L(nelems_tail);
        {
            cmp(reg_reverse_spat_offt_, 1);
            jl(end, T_NEAR);

            compute_dst(1, true);
        }

        L(end);
    }

    void generate() override {
        preamble();
        load_kernel_params();
        prepare_isa_subkernel();
        forward();
        postamble();

        if (with_eltwise_ && postops_injector_)
            postops_injector_->prepare_table();
    }

    jit_uni_binary_kernel_t(const binary_pd_t *pd, bool tail_kernel = false)
        : binary_kernel_t(cpu_isa_traits<isa>::vlen)
        , pd_(pd)
        , is_tail_kernel_(tail_kernel) {
        init();
    }
    ~jit_uni_binary_kernel_t() override = default;
};

template <cpu_isa_t isa, data_type_t src_type>
struct jit_uni_binary_subkernel_t;

template <data_type_t src_type>
struct jit_uni_binary_subkernel_t<avx512_core_bf16, src_type>
    : public jit_uni_binary_kernel_t<avx512_core_bf16> {
    const Opmask &bf16_bcast_opmask = k3;

    void prepare_tail_mask() {
        if (!tail_size_) return;

        const int mask_f32 = (1 << tail_size_) - 1;
        const Reg32 regw_tmp = reg_tmp_.cvt32();
        mov(regw_tmp, mask_f32);
        kmovd(tail_opmask_, regw_tmp);
    }

    void prepare_bf16_bcast_mask() {
        if (is_bf16_ && bcast_type_ != bcast_t::none) {
            const Reg32 regw_tmp = reg_tmp_.cvt32();
            mov(regw_tmp, 1);
            kmovd(bf16_bcast_opmask, regw_tmp);
        }
    }

    void prepare_isa_subkernel() override {
        prepare_tail_mask();
        prepare_bf16_bcast_mask();
    }

    void load_no_tail(const Vmm &dst, const Address &src, data_type_t dt) {
        switch (dt) {
            case data_type::f32: uni_vmovups(dst, src); break;
            case data_type::bf16:
                vpmovzxwd(dst, src);
                vpslld(dst, dst, 0x10);
                break;
            default: assert(!"unreachable");
        }
    }

    void load_tail(const Vmm &dst, const Opmask &opmask, const Address &src,
            data_type_t dt) {
        switch (dt) {
            case data_type::f32: uni_vmovups_tail(dst, opmask, src); break;
            case data_type::bf16:
                vpmovzxwd(dst | opmask, src);
                vpslld(dst, dst, 0x10);
                break;
            default: assert(!"unreachable");
        }
    }

    void store_no_tail(const Address &dst, const Vmm &src, data_type_t dt) {
        const Ymm ymm_src = Ymm(src.getIdx());
        switch (dt) {
            case data_type::f32: uni_vmovups(dst, src); break;
            case data_type::bf16:
                vcvtneps2bf16(ymm_src, src);
                vmovdqu16(dst, ymm_src);
                break;
            default: assert(!"unreachable");
        }
    }

    void store_tail(const Address &dst, const Opmask &opmask, const Vmm &src,
            data_type_t dt) {
        const Ymm ymm_src = Ymm(src.getIdx());
        switch (dt) {
            case data_type::f32: uni_vmovups_tail(dst, opmask, src); break;
            case data_type::bf16:
                vcvtneps2bf16(ymm_src, src);
                vmovdqu16(dst | opmask, ymm_src);
                break;
            default: assert(!"unreachable");
        }
    }

    void load(const Vmm &dst, const Address &src, data_type_t dt, bool tail) {
        if (!tail)
            load_no_tail(dst, src, dt);
        else
            load_tail(dst, tail_opmask_, src, dt);
    }

    void store(const Address &dst, const Vmm &src, data_type_t dt, bool tail) {
        if (!tail)
            store_no_tail(dst, src, dt);
        else
            store_tail(dst, tail_opmask_, src, dt);
    }

    void bcast(const Vmm &dst, const Address &src, data_type_t dt) {
        switch (dt) {
            case data_type::f32: uni_vbroadcastss(dst, src); break;
            case data_type::bf16:
                vpmovzxwd(dst | bf16_bcast_opmask, src);
                vpslld(dst, dst, 0x10);
                uni_vbroadcastss(dst, Xmm(dst.getIdx()));
                break;
            default: assert(!"unreachable");
        }
    }

    void compute_bcast(bool tail) override {
        if (broadcast_src1_value_)
            bcast(vbcast_src1_, src1_ptr(), src_type);
        else if (offt_src1_ == 0)
            load(vbcast_src1_, src1_ptr(), src_type, tail);
    }

    void compute_dst(int unroll, bool tail) override {
        for (int i = 0; i < unroll; i++) {
            const Vmm vreg_tmp_src0 = Vmm(i + 1);
            const Vmm vreg_tmp = Vmm(unroll + i + 1);
            const Vmm vreg_tmp_src1 = offt_src1_ ? vreg_tmp : vbcast_src1_;
            load(vreg_tmp_src0, src0_ptr(i * offt_src0_), src_type, tail);

            if (offt_src1_) {
                load(vreg_tmp_src1, src1_ptr(i * offt_src1_), src_type, tail);
            }
            perform_op(vreg_tmp_src0, vreg_tmp_src1, vreg_scales_src0_,
                    vreg_scales_src1_);
            if (do_sum_) {
                load(vreg_tmp, dst_ptr(i * offt_src0_), src_type, tail);
                uni_vfmadd231ps(vreg_tmp_src0, vreg_tmp, vsum_scale_);
            }
        }

        if (postops_injector_) apply_postops(unroll, tail);

        for (int i = 0; i < unroll; i++) {
            const Vmm vreg_tmp_src0 = Vmm(i + 1);
            store(dst_ptr(i * offt_src0_), vreg_tmp_src0, src_type, tail);
        }
    }

    jit_uni_binary_subkernel_t(const binary_pd_t *pd, bool tail_kernel)
        : jit_uni_binary_kernel_t(pd, tail_kernel) {}
};

template <data_type_t src_type>
struct jit_uni_binary_subkernel_t<avx512_core, src_type>
    : public jit_uni_binary_kernel_t<avx512_core> {
    const Opmask &bf16_bcast_opmask = k3;

    // FP32->BF16 emulation
    std::unique_ptr<bf16_emulation_t> bf16_emu_ {nullptr};
    const Reg64 &reg_bf16_tmp = reg_tmp_;
    const Vmm bf16_emu_reserved_1 = Vmm(26);
    const Vmm bf16_emu_reserved_2 = Vmm(27);
    const Vmm bf16_emu_reserved_3 = Vmm(28);
    const Vmm bf16_emu_reserved_4 = Vmm(29);

    void prepare_tail_mask() {
        if (!tail_size_) return;

        const int mask_f32 = (1 << tail_size_) - 1;
        const Reg32 regw_tmp = reg_tmp_.cvt32();
        mov(regw_tmp, mask_f32);
        kmovd(tail_opmask_, regw_tmp);
    }

    void prepare_bf16_emulator() {
        if (is_bf16_) { // init emulation of bfloat16 operations
            bf16_emu_.reset(new bf16_emulation_t(this, bf16_emu_reserved_1,
                    bf16_emu_reserved_2, bf16_emu_reserved_3, reg_bf16_tmp,
                    bf16_emu_reserved_4, bf16_emu_reserved_4));
            (*bf16_emu_).init_vcvtneps2bf16();
        }
    }

    void prepare_bf16_bcast_mask() {
        if (is_bf16_ && bcast_type_ != bcast_t::none) {
            const Reg32 regw_tmp = reg_tmp_.cvt32();
            mov(regw_tmp, 1);
            kmovd(bf16_bcast_opmask, regw_tmp);
        }
    }

    void prepare_isa_subkernel() override {
        prepare_tail_mask();
        prepare_bf16_emulator();
        prepare_bf16_bcast_mask();
    }

    void load_no_tail(const Vmm &dst, const Address &src, data_type_t dt) {
        switch (dt) {
            case data_type::f32: uni_vmovups(dst, src); break;
            case data_type::bf16:
                vpmovzxwd(dst, src);
                vpslld(dst, dst, 0x10);
                break;
            default: assert(!"unreachable");
        }
    }

    void load_tail(const Vmm &dst, const Opmask &opmask, const Address &src,
            data_type_t dt) {
        switch (dt) {
            case data_type::f32: uni_vmovups_tail(dst, opmask, src); break;
            case data_type::bf16:
                vpmovzxwd(dst | opmask, src);
                vpslld(dst, dst, 0x10);
                break;
            default: assert(!"unreachable");
        }
    }

    void store_no_tail(const Address &dst, const Vmm &src, data_type_t dt) {
        const Ymm ymm_src = Ymm(src.getIdx());
        switch (dt) {
            case data_type::f32: uni_vmovups(dst, src); break;
            case data_type::bf16:
                (*bf16_emu_).vcvtneps2bf16(ymm_src, src);
                vmovdqu16(dst, ymm_src);
                break;
            default: assert(!"unreachable");
        }
    }

    void store_tail(const Address &dst, const Opmask &opmask, const Vmm &src,
            data_type_t dt) {
        const Ymm ymm_src = Ymm(src.getIdx());
        switch (dt) {
            case data_type::f32: uni_vmovups_tail(dst, opmask, src); break;
            case data_type::bf16:
                (*bf16_emu_).vcvtneps2bf16(ymm_src, src);
                vmovdqu16(dst | opmask, ymm_src);
                break;
            default: assert(!"unreachable");
        }
    }

    void load(const Vmm &dst, const Address &src, data_type_t dt, bool tail) {
        if (!tail)
            load_no_tail(dst, src, dt);
        else
            load_tail(dst, tail_opmask_, src, dt);
    }

    void store(const Address &dst, const Vmm &src, data_type_t dt, bool tail) {
        if (!tail)
            store_no_tail(dst, src, dt);
        else
            store_tail(dst, tail_opmask_, src, dt);
    }

    void bcast(const Vmm &dst, const Address &src, data_type_t dt) {
        switch (dt) {
            case data_type::f32: uni_vbroadcastss(dst, src); break;
            case data_type::bf16:
                vpmovzxwd(dst | bf16_bcast_opmask, src);
                vpslld(dst, dst, 0x10);
                uni_vbroadcastss(dst, Xmm(dst.getIdx()));
                break;
            default: assert(!"unreachable");
        }
    }

    void compute_bcast(bool tail) override {
        if (broadcast_src1_value_)
            bcast(vbcast_src1_, src1_ptr(), src_type);
        else if (offt_src1_ == 0)
            load(vbcast_src1_, src1_ptr(), src_type, tail);
    }

    void compute_dst(int unroll, bool tail) override {
        for (int i = 0; i < unroll; i++) {
            const Vmm vreg_tmp_src0 = Vmm(i + 1);
            const Vmm vreg_tmp = Vmm(unroll + i + 1);
            const Vmm vreg_tmp_src1 = offt_src1_ ? vreg_tmp : vbcast_src1_;
            load(vreg_tmp_src0, src0_ptr(i * offt_src0_), src_type, tail);
            if (offt_src1_) {
                load(vreg_tmp_src1, src1_ptr(i * offt_src1_), src_type, tail);
            }
            perform_op(vreg_tmp_src0, vreg_tmp_src1, vreg_scales_src0_,
                    vreg_scales_src1_);
            if (do_sum_) {
                load(vreg_tmp, dst_ptr(i * offt_src0_), src_type, tail);
                uni_vfmadd231ps(vreg_tmp_src0, vreg_tmp, vsum_scale_);
            }
        }

        if (postops_injector_) apply_postops(unroll, tail);

        for (int i = 0; i < unroll; i++) {
            const Vmm vreg_tmp_src0 = Vmm(i + 1);
            store(dst_ptr(i * offt_src0_), vreg_tmp_src0, src_type, tail);
        }
    }

    jit_uni_binary_subkernel_t(const binary_pd_t *pd, bool tail_kernel)
        : jit_uni_binary_kernel_t(pd, tail_kernel) {}
};

template <data_type_t src_type>
struct jit_uni_binary_subkernel_t<avx2, src_type>
    : public jit_uni_binary_kernel_t<avx2> {
    const Vmm tail_vmask_ = Vmm(0);

    void prepare_tail_mask() {
        if (!tail_size_) return;

        static const uint32_t mask_f32[14]
                = {0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff,
                        0xffffffff, 0xffffffff, 0, 0, 0, 0, 0, 0, 0};

        mov(reg_tmp_, reinterpret_cast<size_t>(&mask_f32[7 - tail_size_]));
        vmovups(tail_vmask_, ptr[reg_tmp_]);
    }

    void prepare_isa_subkernel() override { prepare_tail_mask(); }

    void load(const Vmm &dst, const Address &src, bool tail) {
        if (!tail)
            uni_vmovups(dst, src);
        else
            uni_vmovups_tail(dst, tail_vmask_, src);
    }

    void store(const Address &dst, const Vmm &src, bool tail) {
        if (!tail)
            uni_vmovups(dst, src);
        else
            uni_vmovups_tail(dst, tail_vmask_, src);
    }

    void compute_bcast(bool tail) override {
        if (broadcast_src1_value_)
            uni_vbroadcastss(vbcast_src1_, src1_ptr());
        else if (offt_src1_ == 0)
            load(vbcast_src1_, src1_ptr(), tail);
    }

    void compute_dst(int unroll, bool tail) override {
        for (int i = 0; i < unroll; i++) {
            const Vmm vreg_tmp_src0 = Vmm(i + 1);
            const Vmm vreg_tmp = Vmm(unroll + i + 1);
            const Vmm vreg_tmp_src1 = offt_src1_ ? vreg_tmp : vbcast_src1_;
            load(vreg_tmp_src0, src0_ptr(i * offt_src0_), tail);
            if (offt_src1_) load(vreg_tmp_src1, src1_ptr(i * offt_src1_), tail);

            perform_op(vreg_tmp_src0, vreg_tmp_src1, vreg_scales_src0_,
                    vreg_scales_src1_);

            if (do_sum_) {
                load(vreg_tmp, dst_ptr(i * offt_src0_), tail);
                uni_vfmadd231ps(vreg_tmp_src0, vreg_tmp, vsum_scale_);
            }
        }

        if (postops_injector_) apply_postops(unroll, tail);

        for (int i = 0; i < unroll; i++) {
            const Vmm vreg_tmp_src0 = Vmm(i + 1);
            store(dst_ptr(i * offt_src0_), vreg_tmp_src0, tail);
        }
    }

    jit_uni_binary_subkernel_t(const binary_pd_t *pd, bool tail_kernel)
        : jit_uni_binary_kernel_t(pd, tail_kernel) {}
};

template <data_type_t src_type>
struct jit_uni_binary_subkernel_t<sse41, src_type>
    : public jit_uni_binary_kernel_t<sse41> {

    void prepare_isa_subkernel() override {}

    Address get_address(const int off, const int arg_num) {
        switch (arg_num) {
            case DNNL_ARG_SRC_0: return src0_ptr(off);
            case DNNL_ARG_SRC_1: return src1_ptr(off);
            case DNNL_ARG_DST: return dst_ptr(off);
            default: assert(!"unsupported arg_num"); break;
        }
        return Address(0);
    }

    void load(const Vmm &dst, const int off, const int arg_num, bool tail) {
        if (!tail)
            movups(dst, get_address(off, arg_num));
        else
            for (size_t i = 0; i < tail_size_; i++)
                pinsrd(dst, get_address(i * data_type_size_ + off, arg_num), i);
    }

    void store(const Vmm &src, const int off, bool tail) {
        if (!tail)
            movups(get_address(off, DNNL_ARG_DST), src);
        else
            for (size_t i = 0; i < tail_size_; i++)
                pextrd(get_address(i * data_type_size_ + off, DNNL_ARG_DST),
                        src, i);
    }

    void compute_bcast(bool tail) override {
        if (broadcast_src1_value_)
            uni_vbroadcastss(vbcast_src1_, src1_ptr());
        else if (offt_src1_ == 0)
            load(vbcast_src1_, 0, DNNL_ARG_SRC_1, tail);
    }

    void compute_dst(int unroll, bool tail) override {
        for (int i = 0; i < unroll; i++) {
            const Vmm vreg_tmp_src0 = Vmm(i + 1);
            const Vmm vreg_tmp = Vmm(unroll + i + 1);
            const Vmm vreg_tmp_src1 = offt_src1_ ? vreg_tmp : vbcast_src1_;
            load(vreg_tmp_src0, i * offt_src0_, DNNL_ARG_SRC_0, tail);
            if (offt_src1_)
                load(vreg_tmp_src1, i * offt_src1_, DNNL_ARG_SRC_1, tail);

            perform_op(vreg_tmp_src0, vreg_tmp_src1, vreg_scales_src0_,
                    vreg_scales_src1_);
            if (do_sum_) {
                load(vreg_tmp, i * offt_src0_, DNNL_ARG_DST, tail);
                mulps(vreg_tmp, vsum_scale_);
                addps(vreg_tmp_src0, vreg_tmp);
            }
        }

        if (postops_injector_) apply_postops(unroll, tail);

        for (int i = 0; i < unroll; i++) {
            const Vmm vreg_tmp_src0 = Vmm(i + 1);
            store(vreg_tmp_src0, i * offt_src0_, tail);
        }
    }

    jit_uni_binary_subkernel_t(const binary_pd_t *pd, bool tail_kernel)
        : jit_uni_binary_kernel_t(pd, tail_kernel) {}
};

#undef PARAM_OFF

template <data_type_t src_type>
binary_kernel_t *create_binary_kernel(const binary_pd_t *pd, bool tail_kernel) {
    if (mayiuse(avx512_core_bf16)) {
        using subkernel_t
                = jit_uni_binary_subkernel_t<avx512_core_bf16, src_type>;
        return new subkernel_t(pd, tail_kernel);
    } else if (mayiuse(avx512_core)) {
        using subkernel_t = jit_uni_binary_subkernel_t<avx512_core, src_type>;
        return new subkernel_t(pd, tail_kernel);
    } else if (mayiuse(avx2)) {
        using subkernel_t = jit_uni_binary_subkernel_t<avx2, src_type>;
        return new subkernel_t(pd, tail_kernel);
    } else {
        using subkernel_t = jit_uni_binary_subkernel_t<sse41, src_type>;
        return new subkernel_t(pd, tail_kernel);
    }
}

template <data_type_t src_type>
jit_uni_binary_t<src_type>::jit_uni_binary_t(const pd_t *apd)
    : primitive_t(apd) {}

template <data_type_t src_type>
jit_uni_binary_t<src_type>::~jit_uni_binary_t() = default;

template <data_type_t src_type>
status_t jit_uni_binary_t<src_type>::init(engine_t *engine) {
    CHECK(safe_ptr_assign(kernel_,
            create_binary_kernel<src_type>(pd(), false /*tail_kernel*/)));

    const memory_desc_wrapper src0_d(pd_->src_md(0));
    const auto &simd_w = kernel_->simd_w();
    const auto oc = src0_d.ndims() >= 2 ? src0_d.dims()[1] : 1;

    if (op_t::c_blocked == get_op_type(src0_d) && oc % simd_w) {
        CHECK(safe_ptr_assign(kernel_tail_,
                create_binary_kernel<src_type>(pd(), true /*tail_kernel*/)));
        CHECK(kernel_tail_->create_kernel());
    }

    return kernel_->create_kernel();
}

template <data_type_t src_type>
void jit_uni_binary_t<src_type>::execute_no_bcast_strategy(const data_t *src0,
        const data_t *src1, data_t *dst, const float *scale0,
        const float *scale1,
        const std::vector<const void *> &post_ops_binary_rhs_arg_vec,
        const bcast_t bcast_type) const {
    const auto kernel = kernel_.get();
    const auto &simd_w = kernel_->simd_w();

    const memory_desc_wrapper src0_d(pd()->src_md(0));
    const dim_t nelems0 = src0_d.nelems(true);
    const dim_t nelems0_simd = nelems0 / simd_w;
    const dim_t nelems0_tail = nelems0 % simd_w;
    const bool has_tail = nelems0_tail > 0;

    const bool point_broadcast = bcast_type == bcast_t::scalar;

    // Compute strategy:
    // Compute number of vectors, divide it equally between all threads.
    // Last one will also handle a tail if present.
    parallel(0, [&](const int ithr, const int nthr) {
        dim_t start = 0, end = 0;
        balance211(nelems0_simd + has_tail, nthr, ithr, start, end);
        if (start >= end) return;

        const bool ithr_does_tail = has_tail && end == nelems0_simd + has_tail;
        const dim_t n_simd_to_do = (end - start - ithr_does_tail) * simd_w;
        const dim_t tail_to_do = ithr_does_tail * nelems0_tail;

        binary_kernel_t::call_params_t p;
        p.spat_offt_count = (n_simd_to_do + tail_to_do) * sizeof(data_t);
        p.src0 = src0 + start * simd_w;
        p.src1 = src1 + (point_broadcast ? 0 : (start * simd_w));
        p.dst = dst + start * simd_w;
        p.scales_src0 = scale0;
        p.scales_src1 = scale1;
        p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
        (*kernel)(&p);
    });
}

template <data_type_t src_type>
void jit_uni_binary_t<src_type>::execute_bcast_per_c_strategy(
        const data_t *src0, const data_t *src1, data_t *dst,
        const float *scale0, const float *scale1,
        const std::vector<const void *> &post_ops_binary_rhs_arg_vec,
        const op_t op_type, const bcast_t bcast_type,
        const bool blocked_oc_tail) const {
    const auto kernel = kernel_.get();
    const auto kernel_tail = kernel_tail_.get();
    const auto &simd_w = kernel_->simd_w();

    const memory_desc_wrapper src0_d(pd()->src_md(0));
    const memory_desc_wrapper src1_d(pd()->src_md(1));
    const auto ndims = src0_d.ndims();
    const auto &dims = src0_d.dims();
    const dim_t MB = dims[0];
    const dim_t C = ndims >= 2 ? dims[1] : 1;
    const dim_t SP = ndims >= 3 ? utils::array_product(dims + 2, ndims - 2) : 1;

    const auto &bcast_dims = pd()->broadcast_dims();
    const bool no_broadcast = bcast_type == bcast_t::none;
    const bool point_broadcast = bcast_type == bcast_t::scalar;

    const dim_t nelems_slice_src0
            = utils::array_product(src0_d.padded_dims() + 1, ndims - 1);
    const dim_t nelems_slice_src1 = no_broadcast
            ? nelems_slice_src0
            : ((bcast_dims[0] == 0) ? utils::array_product(
                       src1_d.padded_dims() + 1, ndims - 1)
                                    : 0);

    if (op_type == op_t::c_blocked) {
        const dim_t C_blocks = std::ceil(src0_d.padded_dims()[1] / simd_w);
        // Compute strategy:
        // Each block is individual - parallel over MB and C_blocks safely.

        const std::function<void(binary_kernel_t::call_params_t *, dim_t)>
                kernel_blocked_no_tail = [&](binary_kernel_t::call_params_t *p,
                                                 dim_t C_blk) { (*kernel)(p); };
        const std::function<void(binary_kernel_t::call_params_t *, dim_t)>
                kernel_blocked_tail
                = [&](binary_kernel_t::call_params_t *p, dim_t C_blk) {
                      if (C_blk == (C_blocks - 1))
                          (*kernel_tail)(p);
                      else
                          (*kernel)(p);
                  };
        const auto &kernel_blocked = blocked_oc_tail ? kernel_blocked_tail
                                                     : kernel_blocked_no_tail;

        parallel_nd(MB, C_blocks, [&](dim_t mb, dim_t C_blk) {
            binary_kernel_t::call_params_t p;
            p.spat_offt_count = SP * simd_w * sizeof(data_t);
            const dim_t off = mb * nelems_slice_src0 + C_blk * SP * simd_w;
            p.dst = dst + off;
            p.src0 = src0 + off;
            const dim_t src1_off = point_broadcast
                    ? mb * nelems_slice_src1
                    : (no_broadcast ? off
                                    : mb * nelems_slice_src1 + C_blk * simd_w);
            p.src1 = src1 + src1_off;
            p.oc_l_off = C_blk * simd_w;
            p.scales_src0 = scale0;
            p.scales_src1 = scale1;
            p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
            kernel_blocked(&p, C_blk);
        });
    } else if (op_type == op_t::n_spatial_c) {
        // Compute strategy:
        // Each line of channels is individual, parallel over MB and spatial.

        parallel_nd(MB, SP, [&](dim_t mb, dim_t sp) {
            binary_kernel_t::call_params_t p;
            p.spat_offt_count = C * sizeof(data_t);
            const auto off = mb * nelems_slice_src0 + sp * C;
            p.dst = dst + off;
            p.src0 = src0 + off;
            const auto src1_off = no_broadcast ? off : mb * nelems_slice_src1;
            p.src1 = src1 + src1_off;
            p.oc_l_off = 0;
            p.scales_src0 = scale0;
            p.scales_src1 = scale1;
            p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
            (*kernel)(&p);
        });
    } else if (op_type == op_t::n_c_spatial) {
        // Compute strategy:
        // Each line of spatial is individual, parallel over MB and C.

        parallel_nd(MB, C, [&](dim_t mb, dim_t c) {
            binary_kernel_t::call_params_t p;
            p.spat_offt_count = SP * sizeof(data_t);
            const auto off = mb * nelems_slice_src0 + c * SP;
            p.dst = dst + off;
            p.src0 = src0 + off;
            const dim_t src1_off = point_broadcast
                    ? mb * nelems_slice_src1
                    : (no_broadcast ? off : mb * nelems_slice_src1 + c);
            p.src1 = src1 + src1_off;
            p.oc_l_off = c;
            p.scales_src0 = scale0;
            p.scales_src1 = scale1;
            p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
            (*kernel)(&p);
        });
    }
}

template <data_type_t src_type>
void jit_uni_binary_t<src_type>::execute_bcast_per_w_strategy(
        const data_t *src0, const data_t *src1, data_t *dst,
        const float *scale0, const float *scale1,
        const std::vector<const void *> &post_ops_binary_rhs_arg_vec,
        const op_t op_type, const bool blocked_oc_tail) const {
    const auto kernel = kernel_.get();
    const auto kernel_tail = kernel_tail_.get();
    const auto &simd_w = kernel_->simd_w();

    const memory_desc_wrapper src0_d(pd()->src_md(0));
    const auto ndims = src0_d.ndims();
    const auto &dims = src0_d.dims();
    const dim_t MB = dims[0];
    const dim_t W = ndims >= 3 ? dims[ndims - 1] : 1;
    const dim_t C = ndims >= 2 ? dims[1] : 1;
    const dim_t SP = ndims >= 3 ? utils::array_product(dims + 2, ndims - 2) : 1;
    // spatial dimensions without the last one
    const dim_t N = SP / W;

    const auto &bcast_dims = pd()->broadcast_dims();

    const dim_t nelems_slice_src0
            = utils::array_product(src0_d.padded_dims() + 1, ndims - 1);

    if (op_type == op_t::c_blocked) {
        const dim_t C_blocks = std::ceil(src0_d.padded_dims()[1] / simd_w);
        // Compute strategy:
        // Each line of channels is individual, parallel over MB, C_blocks
        // and spatial (width and other spatial dims separately).

        const std::function<void(binary_kernel_t::call_params_t *, dim_t)>
                kernel_blocked_no_tail = [&](binary_kernel_t::call_params_t *p,
                                                 dim_t C_blk) { (*kernel)(p); };
        const std::function<void(binary_kernel_t::call_params_t *, dim_t)>
                kernel_blocked_tail
                = [&](binary_kernel_t::call_params_t *p, dim_t C_blk) {
                      if (C_blk == (C_blocks - 1))
                          (*kernel_tail)(p);
                      else
                          (*kernel)(p);
                  };
        const auto &kernel_blocked = blocked_oc_tail ? kernel_blocked_tail
                                                     : kernel_blocked_no_tail;

        parallel_nd(MB, C_blocks, N, W,
                [&](dim_t mb, dim_t C_blk, dim_t n, dim_t w) {
                    binary_kernel_t::call_params_t p;
                    p.spat_offt_count = simd_w * sizeof(data_t);
                    const auto off = mb * nelems_slice_src0
                            + simd_w * (C_blk * SP + n * W + w);
                    p.dst = dst + off;
                    p.src0 = src0 + off;
                    // check if mb is broadcast
                    const dim_t src1_off = bcast_dims[0] == 1
                            ? w * simd_w
                            : (mb * W + w) * simd_w;
                    p.src1 = src1 + src1_off;
                    p.oc_l_off = C_blk * simd_w;
                    p.scales_src0 = scale0;
                    p.scales_src1 = scale1;
                    p.post_ops_binary_rhs_arg_vec
                            = post_ops_binary_rhs_arg_vec.data();
                    kernel_blocked(&p, C_blk);
                });
    } else if (op_type == op_t::n_spatial_c) {
        // Compute strategy:
        // Each line of channels is individual, parallel over MB and spatial
        // (width and other spatial dims separately).

        parallel_nd(MB, N, W, [&](dim_t mb, dim_t n, dim_t w) {
            binary_kernel_t::call_params_t p;
            p.spat_offt_count = C * sizeof(data_t);
            const auto off = mb * nelems_slice_src0 + n * W * C + w * C;
            p.dst = dst + off;
            p.src0 = src0 + off;
            const dim_t src1_off = bcast_dims[0] == 1 ? w : mb * W + w;
            p.src1 = src1 + src1_off;
            p.oc_l_off = 0;
            p.scales_src0 = scale0;
            p.scales_src1 = scale1;
            p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
            (*kernel)(&p);
        });
    } else if (op_type == op_t::n_c_spatial) {
        // Compute strategy:
        // Each line of width is individual, parallel over MB, C and spatial
        // without W. Use a kernel which broadcasts c_i value
        // into a vector register.

        parallel_nd(MB, C, N, [&](dim_t mb, dim_t c, dim_t n) {
            binary_kernel_t::call_params_t p;
            p.spat_offt_count = W * sizeof(data_t);
            const auto off = mb * nelems_slice_src0 + c * N * W + n * W;
            p.dst = dst + off;
            p.src0 = src0 + off;
            const dim_t src1_off = bcast_dims[0] == 1 ? 0 : mb * W;
            p.src1 = src1 + src1_off;
            p.oc_l_off = c;
            p.scales_src0 = scale0;
            p.scales_src1 = scale1;
            p.post_ops_binary_rhs_arg_vec = post_ops_binary_rhs_arg_vec.data();
            (*kernel)(&p);
        });
    }
}

template <data_type_t src_type>
status_t jit_uni_binary_t<src_type>::execute(const exec_ctx_t &ctx) const {
    status_t status = status::success;
    const auto src0 = CTX_IN_MEM(const data_t *, DNNL_ARG_SRC_0);
    const auto src1 = CTX_IN_MEM(const data_t *, DNNL_ARG_SRC_1);
    auto dst = CTX_OUT_CLEAN_MEM(data_t *, DNNL_ARG_DST, status);
    CHECK(status);
    const auto &post_ops = pd()->attr()->post_ops_;
    const auto &post_ops_binary_rhs_arg_vec
            = binary_injector::prepare_binary_args(post_ops, ctx);
    const auto scales_src0 = pd()->attr()->scales_.get(DNNL_ARG_SRC_0).scales_;
    const auto scales_src1 = pd()->attr()->scales_.get(DNNL_ARG_SRC_1).scales_;

    const memory_desc_wrapper src0_d(pd()->src_md(0));
    const memory_desc_wrapper src1_d(pd()->src_md(1));
    const auto ndims = src0_d.ndims();
    const auto &dims = src0_d.dims();
    const dim_t C = ndims >= 2 ? dims[1] : 1;

    const bool postops_per_oc_broadcast_exists
            = binary_injector::any_binary_postop_rhs_per_oc_broadcast(
                    post_ops, src0_d, get_supported_bcast_strategies());
    const auto &bcast_dims = pd()->broadcast_dims();
    const auto bcast_type = pd()->is_tensor_op()
            ? bcast_t::none
            : get_bcast_type(src1_d, bcast_dims);
    const bool point_broadcast = bcast_type == bcast_t::scalar;
    const auto op_type = get_op_type(src0_d);
    const bool with_postops = !post_ops.entry_.empty();
    const auto &simd_w = kernel_->simd_w();
    const bool has_oc_tail = C % simd_w;
    const bool point_broadcast_no_oc_tail = point_broadcast && !has_oc_tail;
    const auto alg = pd()->desc()->alg_kind;
    // Use strategy with kernel_tail for GreaterEqual op with oc_tail and
    // blocked format due to overwriting the vector tail by vcmpps.
    const bool vector_overwrite = utils::one_of(alg, alg_kind::binary_ge,
            alg_kind::binary_gt, alg_kind::binary_le, alg_kind::binary_lt,
            alg_kind::binary_eq, alg_kind::binary_ne);
    const bool blocked_oc_tail = op_type == op_t::c_blocked && has_oc_tail
            && (with_postops || point_broadcast || bcast_type == bcast_t::per_w
                    || vector_overwrite);

    if ((bcast_type == bcast_t::none || point_broadcast_no_oc_tail)
            && !postops_per_oc_broadcast_exists && !blocked_oc_tail)
        execute_no_bcast_strategy(src0, src1, dst, scales_src0, scales_src1,
                post_ops_binary_rhs_arg_vec, bcast_type);
    else if (bcast_type == bcast_t::per_w)
        execute_bcast_per_w_strategy(src0, src1, dst, scales_src0, scales_src1,
                post_ops_binary_rhs_arg_vec, op_type, blocked_oc_tail);
    else
        execute_bcast_per_c_strategy(src0, src1, dst, scales_src0, scales_src1,
                post_ops_binary_rhs_arg_vec, op_type, bcast_type,
                blocked_oc_tail);

    return status::success;
}

using namespace data_type;

template struct jit_uni_binary_t<f32>;
template struct jit_uni_binary_t<bf16>;

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