path: root/thirdparty/oidn/mkl-dnn/src/cpu/jit_avx512_common_lrn.cpp

/*******************************************************************************
* Copyright 2017-2018 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 "c_types_map.hpp"
#include "mkldnn_thread.hpp"
#include "type_helpers.hpp"
#include "utils.hpp"

#include "jit_avx512_common_lrn.hpp"

#include "jit_generator.hpp"

#define FWD_RBC 4
#define BWD_RBC 3

#define XMM_SIZE (4*sizeof(float))
#define ZMM_SIZE (vlen)
#define BUFFER_BLOCK (XMM_SIZE + ZMM_SIZE + XMM_SIZE)
#define BUFFER_NEXT_OFFSET (XMM_SIZE + ZMM_SIZE)
#define SRC_PREV_OFFSET (vlen - XMM_SIZE)

#define IRB_LOOP(statement) for(int irb = 0; irb < loop_size; irb++) { \
    statement;\
}

namespace mkldnn {
namespace impl {
namespace cpu {

using namespace mkldnn::impl::status;
using namespace mkldnn::impl::utils;

using namespace Xbyak;

enum params { vsize = 16, vlen = 64};

typedef struct {
    const float *src;
    float *dst, *ws0, *ws1;
} jit_args_fwd_t;

typedef struct {
    const float *src, *diff_dst, *ws0, *ws1;
    float *diff_src;
} jit_args_bwd_t;

struct nChw16c_across {
/*  version:
 *  -1: channels 0..15,
 *   1: channels C-16 .. C-1,
 *   0: other channels
 *   3: channels only for this kernel(without prev and next)
 */
    int H, W, version;
    nChw16c_across(int h, int w, int v) : H(h), W(w), version(v) {}
};

struct jit_avx512_common_lrn_fwd_t::jit_avx512_common_lrn_kernel_f32:
       public jit_generator {
    int HW, W;
    bool is_first;
    bool is_last;
    bool is_single;

    Reg64 src = rax;
    Reg64 dst = r8;
    Reg64 scratch0 = rdx;
    Reg64 scratch1 = rsi;
    Reg64 imm_addr64 = rbx;

    Zmm zalpha = zmm0;
    Xmm xalpha = xmm0;
    Zmm zk = zmm1;
    Xmm xk = xmm1;

    Reg64 param = abi_param1;
    Reg64 t = rsp;
    Reg64 hw = r9;

    int xsrc_prev = 2;
    int zsrc = 7;
    int xsrc_next = 3;
    int zc = 7;

    int za = 2;
    int zb = 3;
    int zd = 5;
    int ze = 6;
    int zsum = 4;
    int zdst = 2;
    int zbase = 3;
    int zsum2 = 5;

    prop_kind_t pk;
    int use_h_parallelism;

    float alpha, k;

    DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_avx512_common_lrn_kernel_f32)

    void (*ker)(jit_args_fwd_t *);
    void operator()(jit_args_fwd_t *arg) { ker(arg); }

    enum {
        prf0_offt = 1*FWD_RBC,
        prf2_offt = 8*FWD_RBC
    };

    inline void compute_loop(int loop_size_param)
    {
        // loop_size - param for IRB_LOOP macro
        int loop_size = FWD_RBC;

        auto xreg = [=](int irb, int i) {
            return Xmm(irb*3 + i);
        };

        auto zreg = [=](int irb, int i) {
            return Zmm(irb*7 + i);
        };

        if (!is_first && !is_single) {
            IRB_LOOP(mic_prefetcht0(ptr[src + (irb + prf0_offt - HW)*vlen]));
            IRB_LOOP(mic_prefetcht2(ptr[src + (irb + prf2_offt - HW)*vlen]));
        }
        IRB_LOOP(mic_prefetcht0(EVEX_compress_addr(src, (irb + prf0_offt)*vlen)));
        IRB_LOOP(mic_prefetcht2(EVEX_compress_addr(src, (irb + prf2_offt)*vlen)));
        if (!is_last && !is_single) {
            IRB_LOOP(mic_prefetcht0(ptr[src + (irb + prf0_offt + HW)*vlen]));
            IRB_LOOP(mic_prefetcht2(ptr[src + (irb + prf2_offt + HW)*vlen]));
        }
        if (pk != prop_kind::forward_inference) {
            IRB_LOOP(mic_prefetcht0(EVEX_compress_addr(scratch0,
                       (irb + prf0_offt)*vlen)));
            IRB_LOOP(mic_prefetcht2(EVEX_compress_addr(scratch0,
                       (irb + prf2_offt)*vlen)));
        }
        IRB_LOOP(mic_prefetcht0(EVEX_compress_addr(dst, (irb + prf0_offt)*vlen)));
        IRB_LOOP(mic_prefetcht2(EVEX_compress_addr(dst, (irb + prf2_offt)*vlen)));
        if (pk != prop_kind::forward_inference) {
            IRB_LOOP(mic_prefetcht0(EVEX_compress_addr(scratch1,
                         (irb + prf0_offt) * vlen)));
            IRB_LOOP(mic_prefetcht2(EVEX_compress_addr(scratch1,
                         (irb + prf2_offt) * vlen)));
        }

        loop_size = loop_size_param;
        if (loop_size == 0)
            return;
        if (!is_first && !is_single) {
            IRB_LOOP(vmovups(xreg(irb, xsrc_prev),
                        ptr[src + (irb - HW) * vlen + SRC_PREV_OFFSET]));
        }
        IRB_LOOP(vmovups(zreg(irb, zsrc), EVEX_compress_addr(src,irb*vlen)));
        if (!is_last && !is_single) {
            IRB_LOOP(vmovups(xreg(irb, xsrc_next),
                        ptr[src + (irb + HW) * vlen]));
        }

        if (!is_first && !is_single) {
            IRB_LOOP(vmovups(ptr[t + irb*BUFFER_BLOCK],
                        xreg(irb, xsrc_prev)));
        }
        IRB_LOOP(vmovups(EVEX_compress_addr(t, irb*BUFFER_BLOCK + XMM_SIZE),
                    zreg(irb, zsrc)));
        if (!is_last && !is_single) {
            IRB_LOOP(vmovups(ptr[t + irb*BUFFER_BLOCK + BUFFER_NEXT_OFFSET],
                    xreg(irb, xsrc_next)));
        }

        IRB_LOOP(vmovups(zreg(irb, za), EVEX_compress_addr(t, irb*BUFFER_BLOCK
                        + XMM_SIZE - 2*sizeof(float))));
        IRB_LOOP(vmovups(zreg(irb, zb), EVEX_compress_addr(t, irb*BUFFER_BLOCK
                        + XMM_SIZE - sizeof(float))));
        IRB_LOOP(vmovups(zreg(irb, zd), EVEX_compress_addr(t, irb*BUFFER_BLOCK
                        + XMM_SIZE + sizeof(float))));
        IRB_LOOP(vmovups(zreg(irb, ze), EVEX_compress_addr(t, irb*BUFFER_BLOCK
                        + XMM_SIZE + 2*sizeof(float))));

        assert(zc == zsrc);
        IRB_LOOP(vmulps(zreg(irb, zsum), zreg(irb, zc), zreg(irb, zc)));

        IRB_LOOP(vfmadd231ps(zreg(irb, zsum), zreg(irb, za), zreg(irb, za)));
        IRB_LOOP(vfmadd231ps(zreg(irb, zsum), zreg(irb, zb), zreg(irb, zb)));
        IRB_LOOP(vfmadd231ps(zreg(irb, zsum), zreg(irb, zd), zreg(irb, zd)));
        IRB_LOOP(vfmadd231ps(zreg(irb, zsum), zreg(irb, ze), zreg(irb, ze)));

        IRB_LOOP(vfmadd132ps(zreg(irb, zsum), zk, zalpha));

        IRB_LOOP(vmovaps(zreg(irb, zbase), zreg(irb, zsum)));

        IRB_LOOP(vmulps(zreg(irb, zsum2), zreg(irb, zsum), zreg(irb, zsum)));
        IRB_LOOP(vmulps(zreg(irb, zsum), zreg(irb, zsum), zreg(irb, zsum2)));

        IRB_LOOP(vsqrtps(zreg(irb, zsum), zreg(irb, zsum)));
        IRB_LOOP(vsqrtps(zreg(irb, zsum), zreg(irb, zsum)));

        if (pk != prop_kind::forward_inference) {
            IRB_LOOP(vmovups(EVEX_compress_addr(scratch0, irb*vlen),
                        zreg(irb, zsum)));
        }
        IRB_LOOP(vdivps(zreg(irb, zdst), zreg(irb, zsrc), zreg(irb, zsum)));
        IRB_LOOP(vmovups(EVEX_compress_addr(dst, irb*vlen), zreg(irb, zdst)));
        if (pk != prop_kind::forward_inference) {
            /* ws1 = zdst / zbase = zsrc / (zbase^1.75) */
            IRB_LOOP(vdivps(zreg(irb, zsum), zreg(irb, zdst), zreg(irb, zbase)));
            IRB_LOOP(vmovups(EVEX_compress_addr(scratch1, irb*vlen),
                        zreg(irb, zsum)));
        }
    }

    jit_avx512_common_lrn_kernel_f32(
        const struct nChw16c_across &J,
        prop_kind_t prop_kind,
        int use_h_parallel,
        float A,
        float K,
        void *code_ptr = nullptr,
        size_t code_size = 2 * Xbyak::DEFAULT_MAX_CODE_SIZE)
        : jit_generator(code_ptr, code_size)
        , pk(prop_kind)
        , use_h_parallelism(use_h_parallel)
        , alpha(A)
        , k(K)
    {
        this->preamble();

        mov(src, ptr[param + 0]);
        mov(dst, ptr[param + 8]);
        if (pk != prop_kind::forward_inference)
        {
            mov(scratch0, ptr[param + 16]);
            mov(scratch1, ptr[param + 24]);
        }
        is_first = J.version == -1 || J.version == -2;
        is_last  = J.version == +1 || J.version == -2;
        is_single = J.version == 3;

        W = J.W;
        HW = J.W*J.H;
        int LSB = use_h_parallelism ? W : HW;

        sub(t, FWD_RBC*BUFFER_BLOCK);
        mov(imm_addr64, float2int(this->alpha));
        movq(xalpha, imm_addr64);
        vbroadcastss(zalpha, xalpha);

        mov(imm_addr64, float2int(this->k));
        movq(xk, imm_addr64);
        vbroadcastss(zk, xk);

        if (is_first || is_single) {
            vxorps(xmm2, xmm2, xmm2);
            for(int irb = 0; irb < FWD_RBC; irb++) {
                vmovups(ptr[t + irb*BUFFER_BLOCK], xmm2);
            }
        }
        if (is_last || is_single) {
            vxorps(xmm2, xmm2, xmm2);
            for(int irb = 0; irb < FWD_RBC; irb++) {
                vmovups(ptr[t + irb*BUFFER_BLOCK + BUFFER_NEXT_OFFSET],
                    xmm2);
            }
        }

        int LSREST = LSB % FWD_RBC;
        int LS = LSB - LSREST;

        Label lrn_loop;

        if (LS > 0) {
            mov(hw, LS);

            L(lrn_loop);
            {
                compute_loop(FWD_RBC);

                add(src, FWD_RBC*vlen);
                add(dst, FWD_RBC*vlen);
                if (pk != prop_kind::forward_inference)
                {
                    add(scratch0, FWD_RBC*vlen);
                    add(scratch1, FWD_RBC*vlen);
                }

                for(int irb = 0; irb < FWD_RBC; irb++)
                    dec(hw);
                cmp(hw, 0);
                jne(lrn_loop, T_NEAR);
            }
        }

        compute_loop(LSREST);

        add(t, FWD_RBC*BUFFER_BLOCK);
        this->postamble();

        ker = reinterpret_cast<decltype(ker)>(const_cast<uint8_t*>(
                    this->getCode()));
    }
};

status_t jit_avx512_common_lrn_fwd_t::pd_t::init() {
    using namespace prop_kind;
    using namespace alg_kind;

    const memory_desc_wrapper data_d(src_md());
    bool ok = true
        && mayiuse(avx512_common)
        && is_fwd()
        && !has_zero_dim_memory()
        && everyone_is(data_type::f32, data_d.data_type())
        && data_d.ndims() == 4
        && data_d.dims()[1] % vsize == 0
        && attr()->has_default_values();
    if (!ok) return unimplemented;

    if (desc()->prop_kind == forward_training) {
        dims_t ws_dims = { MB(), C(), H(), 2*W() };
        mkldnn_memory_desc_init_by_tag(&ws_md_, 4, ws_dims, data_type::f32,
                format_tag::nChw16c);
    }

    bool args_ok_across = true
        && desc()->alg_kind == lrn_across_channels
        && desc()->local_size == 5
        && desc()->lrn_beta == 0.75
        && data_d.matches_tag(format_tag::nChw16c);

    return args_ok_across ? success : unimplemented;
}

jit_avx512_common_lrn_fwd_t::jit_avx512_common_lrn_fwd_t(const pd_t *apd)
    : cpu_primitive_t(apd)
    , use_h_parallelism(0), ker_(nullptr), ker_first_(nullptr)
    , ker_last_(nullptr) {
    using namespace alg_kind;
    const int C = pd()->C();
    const int H = pd()->H();
    const int W = pd()->W();
    const int ls = pd()->desc()->local_size;
    const float alpha = pd()->desc()->lrn_alpha / ls;
    const float k = pd()->desc()->lrn_k;

    auto pk = pd()->desc()->prop_kind;

    use_h_parallelism = H > 28 ? 1 : 0;

    if (C / vsize == 1) {
        ker_ = new jit_avx512_common_lrn_kernel_f32(nChw16c_across(H, W, 3), pk,
            use_h_parallelism, alpha, k);
    } else {
        ker_ = new jit_avx512_common_lrn_kernel_f32(nChw16c_across(H, W, 0), pk,
            use_h_parallelism, alpha, k);
        ker_first_ = new jit_avx512_common_lrn_kernel_f32(
            nChw16c_across(H, W, -1), pk, use_h_parallelism, alpha, k);
        ker_last_ = new jit_avx512_common_lrn_kernel_f32(
            nChw16c_across(H, W, +1), pk, use_h_parallelism, alpha, k);
    }
}

jit_avx512_common_lrn_fwd_t::~jit_avx512_common_lrn_fwd_t()
{ delete ker_; delete ker_first_; delete ker_last_; }

void jit_avx512_common_lrn_fwd_t::execute_forward(const exec_ctx_t &ctx) const
{
    auto src = CTX_IN_MEM(const data_t *, MKLDNN_ARG_SRC);
    auto dst = CTX_OUT_MEM(data_t *, MKLDNN_ARG_DST);
    auto ws = CTX_OUT_MEM(data_t *, MKLDNN_ARG_WORKSPACE);

    const int N = pd()->MB();
    const int C = pd()->C();
    const int H = pd()->H();
    const int W = pd()->W();

    parallel(0, [&](const int ithr, const int nthr) {
        size_t start{0}, end{0};
        const int C16 = C / vsize;
        const size_t work_amount = use_h_parallelism ? N*C16*H : N*C16;

        balance211(work_amount, nthr, ithr, start, end);
        if (use_h_parallelism) {
            int n{0}, c16{0}, h{0};
            nd_iterator_init(start, n, N, c16, C16, h, H);
            for (size_t iwork = start; iwork < end; ++iwork) {
                auto offset = n*C*H*W + c16*H*W*vsize
                    + h*W*vsize;
                auto ws_offset0 = n*C*H*2*W + c16*H*2*W*vsize
                    + h*2*W*vsize;
                auto ws_offset1 = ws_offset0 + W*vsize;

                jit_args_fwd_t args;
                args.src = &src[offset];
                args.dst = &dst[offset];
                args.ws0 = &ws[ws_offset0];
                args.ws1 = &ws[ws_offset1];

                if (C16 == 1)
                    (*ker_)(&args);
                else if (c16 == 0)
                    (*ker_first_)(&args);
                else if (c16 == C16 - 1)
                    (*ker_last_)(&args);
                else
                    (*ker_)(&args);
                nd_iterator_step(n, N, c16, C16, h, H);
            }
        } else {
            int n{0}, c16{0};
            nd_iterator_init(start, n, N, c16, C16);
            for (size_t iwork = start; iwork < end; ++iwork) {
                auto offset = n*C*H*W + c16*H*W*vsize;
                auto ws_offset0 = n*C*H*2*W + c16*H*2*W*vsize;
                auto ws_offset1 = ws_offset0 + H*W*vsize;

                jit_args_fwd_t args;
                args.src = &src[offset];
                args.dst = &dst[offset];
                args.ws0 = &ws[ws_offset0];
                args.ws1 = &ws[ws_offset1];

                if (C16 == 1)
                    (*ker_)(&args);
                else if (c16 == 0)
                    (*ker_first_)(&args);
                else if (c16 == C16 - 1)
                    (*ker_last_)(&args);
                else
                    (*ker_)(&args);

                nd_iterator_step(n, N, c16, C16);
            }
        }
    });
}

struct jit_avx512_common_lrn_bwd_t::jit_avx512_common_lrn_kernel_f32:
    public jit_generator {
    int HW, W;
    bool is_first;
    bool is_last;
    bool is_single;

    Reg64 src = rax;
    Reg64 diffsrc = r8;
    Reg64 diffdst = r9;
    Reg64 workspace0 = rdx;
    Reg64 workspace1 = rsi;
    Reg64 imm_addr64 = rbx;

    Zmm znalphabeta = zmm0;
    Xmm xnalphabeta = xmm0;

    Reg64 param = abi_param1;
    Reg64 t = rsp;
    Reg64 hw = r10;

    int xws1_prev = 1;
    int xdiffdst_prev = 2;
    int zws1 = 1;

    int zsrc = 1;
    int zdiffdst = 5;
    int zdiffsrc = 6;

    int xws1_next = 1;
    int xdiffdst_next = 3;

    int za = 1;
    int zb = 2;
    int zd = 3;
    int ze = 4;
    int zws0 = 2;

    float nalphabeta;

    int use_h_parallelism;

    DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_avx512_common_lrn_kernel_f32)

    void (*ker)(jit_args_bwd_t *);
    void operator()(jit_args_bwd_t *arg) { ker(arg); }

    enum {
        prf0_offt = 1*BWD_RBC,
        prf2_offt = 8*BWD_RBC
    };

    inline void compute_loop(int loop_size_param, int prefetchL1,
            int prefetchL2)
    {
        // loop_size - param for IRB_LOOP macro
        int loop_size = loop_size_param;

        auto xreg = [=](int irb, int i) {
            return Xmm(irb*6 + i);
        };

        auto zreg = [=](int irb, int i) {
            return Zmm(irb*6 + i);
        };

// ---- prefetching -------------------------------------------
        if (!is_first && !is_single) {
            if (prefetchL1)
                IRB_LOOP(mic_prefetcht0(ptr[workspace1 + (irb + prf0_offt
                        - 2 * HW) * vlen]));
            if (prefetchL1)
                IRB_LOOP(mic_prefetcht0(ptr[diffdst    + (irb + prf0_offt
                        - HW) * vlen]));
        }

        if (prefetchL1)
            IRB_LOOP(mic_prefetcht0(ptr[src + (irb + prf0_offt)*vlen]));
        if (prefetchL2)
            IRB_LOOP(mic_prefetcht2(ptr[src + (irb + prf2_offt)*vlen]));

        if (prefetchL1)
            IRB_LOOP(mic_prefetcht0(ptr[workspace1 + (irb + prf0_offt)*vlen]));

        if (prefetchL1)
            IRB_LOOP(mic_prefetcht0(ptr[diffdst + (irb + prf0_offt)*vlen]));

        if (!is_last && !is_single) {
            if (prefetchL1)
                IRB_LOOP(mic_prefetcht0(ptr[workspace1 + (irb + prf0_offt
                        + 2 * HW) * vlen]));
            if (prefetchL2)
                IRB_LOOP(mic_prefetcht2(ptr[workspace1 + (irb + prf2_offt
                        + 2 * HW) * vlen]));

            if (prefetchL1)
                IRB_LOOP(mic_prefetcht0(ptr[diffdst +  (irb + prf0_offt
                          + HW) * vlen]));
            if (prefetchL2)
                IRB_LOOP(mic_prefetcht2(ptr[diffdst +  (irb + prf2_offt
                        + HW) * vlen]));
        }
        if (prefetchL1)
            IRB_LOOP(mic_prefetcht0(ptr[workspace0 + (irb + prf0_offt)*vlen]));
        if (prefetchL2)
            IRB_LOOP(mic_prefetcht2(ptr[workspace0 + (irb + prf2_offt)*vlen]));
// -----------------------------------------------------------

        if (loop_size_param == 0)
            return;

        if (!is_first && !is_single) {
            IRB_LOOP(vmovups(xreg(irb, xws1_prev), ptr[workspace1 + (irb
                    - 2 * HW) * vlen + SRC_PREV_OFFSET]));
            IRB_LOOP(vmovups(xreg(irb, xdiffdst_prev), ptr[diffdst + (irb
                    - HW) * vlen + SRC_PREV_OFFSET]));
            IRB_LOOP(vmulps(xreg(irb, xdiffdst_prev), xreg(irb, xdiffdst_prev),
                    xreg(irb, xws1_prev)));
        }

        IRB_LOOP(vmovups(zreg(irb, zws1),
                EVEX_compress_addr(workspace1, irb*vlen)));
        IRB_LOOP(vmovups(zreg(irb, zdiffdst),
                EVEX_compress_addr(diffdst, irb*vlen)));
        IRB_LOOP(vmulps(zreg(irb, zdiffsrc), zreg(irb, zdiffdst),
                zreg(irb, zws1)));

        if (!is_last && !is_single) {
            IRB_LOOP(vmovups(xreg(irb, xws1_next), ptr[workspace1 + (irb
                    + 2 * HW) * vlen]));
            IRB_LOOP(vmovups(xreg(irb, xdiffdst_next), ptr[diffdst +  (irb
                    + HW) * vlen]));
            IRB_LOOP(vmulps(xreg(irb, xdiffdst_next), xreg(irb, xdiffdst_next),
                    xreg(irb, xws1_next)));
        }

        if (!is_first && !is_single) {
            IRB_LOOP(vmovups(ptr[t + irb*BUFFER_BLOCK],
                    xreg(irb, xdiffdst_prev)));
        }
        IRB_LOOP(vmovups(EVEX_compress_addr(t, irb*BUFFER_BLOCK + XMM_SIZE),
                 zreg(irb, zdiffsrc)));
        if (!is_last && !is_single) {
            IRB_LOOP(vmovups(ptr[t + irb*BUFFER_BLOCK + BUFFER_NEXT_OFFSET],
                 xreg(irb, xdiffdst_next)));
        }

        IRB_LOOP(vmovups(zreg(irb, za), EVEX_compress_addr(t, irb*BUFFER_BLOCK
                + XMM_SIZE - 2*sizeof(float))));
        IRB_LOOP(vmovups(zreg(irb, zb), EVEX_compress_addr(t, irb*BUFFER_BLOCK
                + XMM_SIZE - 1*sizeof(float))));
        IRB_LOOP(vmovups(zreg(irb, zd), EVEX_compress_addr(t, irb*BUFFER_BLOCK
                + XMM_SIZE + 1*sizeof(float))));
        IRB_LOOP(vmovups(zreg(irb, ze), EVEX_compress_addr(t, irb*BUFFER_BLOCK
                + XMM_SIZE + 2*sizeof(float))));
        IRB_LOOP(vaddps(zreg(irb, zdiffsrc), zreg(irb, zdiffsrc),
                zreg(irb, za)));
        assert(zsrc == za);
        IRB_LOOP(vmovups(zreg(irb, zsrc), EVEX_compress_addr(src, irb*vlen)));
        IRB_LOOP(vaddps(zreg(irb, zdiffsrc), zreg(irb, zdiffsrc),
                zreg(irb, zb)));
        IRB_LOOP(vaddps(zreg(irb, zdiffsrc), zreg(irb, zdiffsrc),
                zreg(irb, zd)));
        IRB_LOOP(vaddps(zreg(irb, zdiffsrc), zreg(irb, zdiffsrc),
                zreg(irb, ze)));
        IRB_LOOP(vmulps(zreg(irb, zsrc), zreg(irb, zsrc), znalphabeta));

        IRB_LOOP(vmovups(zreg(irb, zws0),
                 EVEX_compress_addr(workspace0, irb*vlen)));
        IRB_LOOP(vdivps(zreg(irb, zdiffdst), zreg(irb, zdiffdst),
                 zreg(irb, zws0)));
        IRB_LOOP(vfmadd213ps(zreg(irb, zdiffsrc), zreg(irb, zsrc),
                 zreg(irb, zdiffdst)));

        Label unaligned_store, end_store;
        test(diffsrc, vlen - 1);
        jnz(unaligned_store, T_NEAR);
        IRB_LOOP(uni_vmovntps(EVEX_compress_addr(diffsrc, irb*vlen),
                 zreg(irb, zdiffsrc)));
        jmp(end_store, T_NEAR);
        L(unaligned_store); {
            IRB_LOOP(uni_vmovups(EVEX_compress_addr(diffsrc, irb*vlen),
                     zreg(irb, zdiffsrc)));
        }
        L(end_store);
    }

    jit_avx512_common_lrn_kernel_f32(
        const struct nChw16c_across &J,
        float A,
        float B,
        int use_h_parallel,
        void *code_ptr = nullptr,
        size_t code_size = 1 * Xbyak::DEFAULT_MAX_CODE_SIZE)
        : jit_generator(code_ptr, code_size)
        , nalphabeta(-2*A*B)
        , use_h_parallelism(use_h_parallel)
    {
        this->preamble();

        mov(src, ptr[param + 0]);
        mov(diffdst, ptr[param + 8]);
        mov(workspace0, ptr[param + 16]);
        mov(workspace1, ptr[param + 24]);
        mov(diffsrc, ptr[param + 32]);

        W = J.W;
        HW = J.H*J.W;
        int LSB = this->use_h_parallelism ? W : HW;

        sub(t, BWD_RBC*BUFFER_BLOCK);
        mov(imm_addr64, float2int(this->nalphabeta));
        movq(xnalphabeta, imm_addr64);
        vbroadcastss(znalphabeta, xnalphabeta);

        is_first = J.version == -1 || J.version == -2;
        is_last  = J.version == +1 || J.version == +2;
        is_single = J.version == 3;

        if (is_first || is_single) {
            vxorps(xmm1, xmm1, xmm1);
            for(int irb = 0; irb < BWD_RBC; irb++) {
                vmovups(ptr[t + irb*BUFFER_BLOCK], xmm1);
            }
        }
        if (is_last || is_single) {
            vxorps(xmm1, xmm1, xmm1);
            for(int irb = 0; irb < BWD_RBC; irb++) {
                vmovups(ptr[t + irb*BUFFER_BLOCK + BUFFER_NEXT_OFFSET], xmm1);
            }
        }

        int LSREST = LSB % BWD_RBC;
        int LS = LSB - LSREST;

        Label lrn_loop;

        if (LS > 0) {
            mov(hw, LS);

            L(lrn_loop);
            {
                compute_loop(BWD_RBC, 1, 1);

                add(src, BWD_RBC*vlen);
                add(diffsrc, BWD_RBC*vlen);
                add(diffdst, BWD_RBC*vlen);
                add(workspace0, BWD_RBC*vlen);
                add(workspace1, BWD_RBC*vlen);

                for(int irb = 0; irb < BWD_RBC; irb++)
                    dec(hw);
                cmp(hw, 0);
                jne(lrn_loop, T_NEAR);
            }
        }

        compute_loop(LSREST, 1, this->use_h_parallelism ? 0 : 1);

        add(t, BWD_RBC*BUFFER_BLOCK);
        this->postamble();

        ker = reinterpret_cast<decltype(ker)>(const_cast<uint8_t*>(
                    this->getCode()));
    }

};

status_t jit_avx512_common_lrn_bwd_t::pd_t::init() {
    using namespace alg_kind;

    const memory_desc_wrapper data_d(src_md());
    bool ok = true
        && mayiuse(avx512_common)
        && !is_fwd()
        && utils::everyone_is(data_type::f32, data_d.data_type())
        && !has_zero_dim_memory()
        && data_d.ndims() == 4
        && data_d.dims()[1] % vsize == 0
        && attr()->has_default_values();
    if (!ok) return unimplemented;

    dims_t ws_dims = { MB(), C(), H(), 2*W() };
    mkldnn_memory_desc_init_by_tag(&ws_md_, 4, ws_dims, data_type::f32,
            format_tag::nChw16c);

    if (!compare_ws(hint_fwd_pd_)) return unimplemented;

    bool args_ok_across = true
        && desc()->alg_kind == lrn_across_channels
        && desc()->local_size == 5
        && desc()->lrn_beta == 0.75
        && data_d.matches_tag(format_tag::nChw16c);

    return args_ok_across ? success : unimplemented;
}

jit_avx512_common_lrn_bwd_t::jit_avx512_common_lrn_bwd_t(const pd_t *apd)
    : cpu_primitive_t(apd)
    , use_h_parallelism(0),  ker_(nullptr), ker_first_(nullptr)
    , ker_last_(nullptr) {
    const int C = pd()->C();
    const int H = pd()->H();
    const int W = pd()->W();
    const int ls = pd()->desc()->local_size;
    const float alpha = pd()->desc()->lrn_alpha / ls;
    const float beta = pd()->desc()->lrn_beta;

    use_h_parallelism = H > 28 ? 1 : 0;

    if (C / vsize == 1) {
        ker_ = new jit_avx512_common_lrn_kernel_f32(nChw16c_across(H, W, 3),
        alpha, beta, use_h_parallelism);
    } else {
        ker_ = new jit_avx512_common_lrn_kernel_f32(nChw16c_across(H, W, 0),
            alpha, beta, use_h_parallelism);
        ker_first_ = new jit_avx512_common_lrn_kernel_f32(
            nChw16c_across(H, W, -1), alpha, beta, use_h_parallelism);
        ker_last_ = new jit_avx512_common_lrn_kernel_f32(
            nChw16c_across(H, W, +1), alpha, beta, use_h_parallelism);
    }
}

jit_avx512_common_lrn_bwd_t::~jit_avx512_common_lrn_bwd_t()
{ delete ker_; delete ker_first_; delete ker_last_; }

void jit_avx512_common_lrn_bwd_t::execute_backward(const exec_ctx_t &ctx) const
{
    auto src = CTX_IN_MEM(const data_t *, MKLDNN_ARG_SRC);
    auto diff_dst = CTX_IN_MEM(const data_t *, MKLDNN_ARG_DIFF_DST);
    auto ws = CTX_IN_MEM(const data_t *, MKLDNN_ARG_WORKSPACE);
    auto diff_src = CTX_OUT_MEM(data_t *, MKLDNN_ARG_DIFF_SRC);

    const int N = pd()->MB();
    const int C = pd()->C();
    const int H = pd()->H();
    const int W = pd()->W();

    parallel(0, [&](const int ithr, const int nthr) {
        size_t start{0}, end{0};
        const int C16 = C / vsize;
        const size_t work_amount = use_h_parallelism ? N*C16*H : N*C16;

        balance211(work_amount, nthr, ithr, start, end);
        if (use_h_parallelism) {
            int n{0}, c16{0}, h{0};
            nd_iterator_init(start, n, N,  h, H, c16, C16);
            for (size_t iwork = start; iwork < end; ++iwork) {
                auto offset = n*C*H*W + c16*H*W*vsize
                    + h*W*vsize;
                auto ws_offset0 = n*C*H*2*W + c16*H*2*W*vsize
                    + h*2*W*vsize;
                auto ws_offset1 = ws_offset0 + W*vsize;

                jit_args_bwd_t args;
                args.src = &src[offset];
                args.diff_dst = &diff_dst[offset];
                args.ws0 = &ws[ws_offset0];
                args.ws1 = &ws[ws_offset1];
                args.diff_src = &diff_src[offset];

                if (C16 == 1)
                    (*ker_)(&args);
                else if (c16 == 0)
                    (*ker_first_)(&args);
                else if (c16 == C16 - 1)
                    (*ker_last_)(&args);
                else
                    (*ker_)(&args);
                nd_iterator_step(n, N, h, H, c16, C16);
            }
        } else {
            int n{0}, c16{0};
            nd_iterator_init(start, n, N, c16, C16);
            for (size_t iwork = start; iwork < end; ++iwork) {
                auto offset = n*C*H*W + c16*H*W*vsize;
                auto ws_offset0 = n*C*H*2*W + c16*H*2*W*vsize;
                auto ws_offset1 = ws_offset0 + H*W*vsize;

                jit_args_bwd_t args;
                args.src = &src[offset];
                args.diff_dst = &diff_dst[offset];
                args.ws0 = &ws[ws_offset0];
                args.ws1 = &ws[ws_offset1];
                args.diff_src = &diff_src[offset];

                if (C16 == 1)
                    (*ker_)(&args);
                else if (c16 == 0)
                    (*ker_first_)(&args);
                else if (c16 == C16 - 1)
                    (*ker_last_)(&args);
                else
                    (*ker_)(&args);

                nd_iterator_step(n, N, c16, C16);
            }
        }
    });
}

}
}
}