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

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

#include "cpu/cpu_primitive.hpp"

#include "cpu/x64/amx_tile_configure.hpp"
#include "cpu/x64/cpu_barrier.hpp"
#include "cpu/x64/injectors/jit_uni_binary_injector.hpp"
#include "cpu/x64/jit_brgemm_inner_product.hpp"
#include "cpu/x64/jit_transpose_utils.hpp"

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

using namespace dnnl::impl::data_type;
using namespace dnnl::impl::format_tag;
using namespace dnnl::impl::memory_tracking::names;
using namespace dnnl::impl::status;
using namespace dnnl::impl::utils;

using namespace nstl;

#define get_blk_off(d, dt, ...) \
    (types::data_type_size((dt)) * (d).blk_off(__VA_ARGS__))

namespace {
template <typename ker_type>
void copy_data_chunk(ker_type &ker, char *tr_data, const char *data,
        int os_work, bool is_last_blk) {
    auto ctx = jit_brgemm_copy_to_coarse_t::ctx_t();
    ctx.data = (void *)data;
    ctx.tr_data = (void *)tr_data;
    ctx.os_work = os_work;
    ctx.last_row_blk = is_last_blk ? 1 : 0;
    (*ker)(&ctx);
}
} // namespace

template <cpu_isa_t isa>
status_t brgemm_inner_product_fwd_t<isa>::execute_forward(
        const exec_ctx_t &ctx) const {
    auto src = CTX_IN_MEM(const char *, DNNL_ARG_SRC);
    auto weights = CTX_IN_MEM(const char *, DNNL_ARG_WEIGHTS);
    auto bias = CTX_IN_MEM(const char *, DNNL_ARG_BIAS);
    auto dst = CTX_OUT_MEM(char *, DNNL_ARG_DST);
    const auto post_ops_binary_rhs_arg_vec
            = binary_injector::prepare_binary_args(
                    pd()->attr()->post_ops_, ctx);

    memory_tracking::grantor_t scratchpad = ctx.get_scratchpad_grantor();
    const memory_desc_wrapper src_d(pd()->src_md());
    const memory_desc_wrapper dst_d(pd()->dst_md());
    const memory_desc_wrapper weights_d(pd()->weights_md(0));

    DEFINE_SCALES_BUFFER(oscales);

    const auto &jbgp = pd()->jbgp_;
    const bool is_f32 = everyone_is(f32, jbgp.src_dt, jbgp.wei_dt, jbgp.dst_dt);

    const size_t src_dt_size = types::data_type_size(jbgp.src_dt);
    const size_t bia_dt_size
            = jbgp.with_bias ? types::data_type_size(jbgp.bia_dt) : 0;
    const size_t acc_dt_size = types::data_type_size(jbgp.acc_dt);
    const size_t dst_dt_size = types::data_type_size(jbgp.dst_dt);

    auto addr_batch_global = scratchpad.template get<brgemm_batch_element_t>(
            key_brgemm_primitive_batch);
    auto a_buffer_global = (jbgp.use_buffer_a)
            ? scratchpad.template get<char>(key_brgemm_primitive_buffer_a)
            : nullptr;
    auto c_buffer_global = (jbgp.use_buffer)
            ? scratchpad.template get<char>(key_brgemm_primitive_buffer)
            : nullptr;
    static constexpr bool is_amx = (isa == avx512_core_bf16_amx_int8
            || isa == avx512_core_bf16_amx_bf16);
    auto wsp_tile_base = is_amx
            ? ctx.get_scratchpad_grantor().template get<char>(
                    key_conv_amx_tile_buffer)
            : nullptr;

    const int ic_chunks = div_up(jbgp.nb_ic, jbgp.nb_ic_blocking);

    const bool are_post_ops_applicable = one_of(true, jbgp.with_sum,
            jbgp.with_bias, jbgp.with_scales, jbgp.with_eltwise,
            jbgp.with_binary, jbgp.acc_dt != jbgp.dst_dt, jbgp.signed_input);

    size_t offset = types::data_type_size(jbgp.wei_dt)
            * (weights_d.size() - weights_d.additional_buffer_size());
    auto compensation = (jbgp.signed_input)
            ? reinterpret_cast<const int32_t *>(&weights[offset])
            : nullptr;

    bool is_os_tail = (jbgp.mb < jbgp.os_block);
    bool is_oc_tail = (jbgp.oc < jbgp.oc_block);
    int base_brg_ker_idx = brgemm_inner_product_utils::
            get_brg_kernel_index( // TODO: Can be calculated on initialization stage
                    jbgp, false, false, is_os_tail, is_oc_tail, false);

    const auto ker = [&](int ithr_oc_mb, int nthr_oc_mb, int ithr_ic, int n,
                             int ocb, int icc, bool do_init, int buffer_a_osb,
                             bool copy_buffer_a) {
        const int ithr = nthr_oc_mb * ithr_ic + ithr_oc_mb;
        auto addr_batch = addr_batch_global + ithr * jbgp.adjusted_batch_size;

        const size_t a_buffer_osb_stride
                = src_dt_size * jbgp.LDA * jbgp.os_block;
        const size_t a_buffer_per_thr
                = a_buffer_osb_stride * jbgp.nb_os_blocking;
        auto a_buffer = (jbgp.use_buffer_a)
                ? a_buffer_global + ithr * a_buffer_per_thr
                        + buffer_a_osb * a_buffer_osb_stride
                : nullptr;

        const int oc = ocb * jbgp.oc_block;
        const size_t dst_off = get_blk_off(dst_d, jbgp.dst_dt, n, oc);

        const bool use_c_buffer = (jbgp.with_sum)
                || (jbgp.use_buffer && (jbgp.nthr_ic_b == 1 || ithr_ic > 0));

        char *c_buffer = nullptr;
        if (use_c_buffer) {
            const size_t c_buf_thr_idx = jbgp.nthr_ic_b <= 1
                    ? ithr
                    : (jbgp.acc_dt != jbgp.dst_dt || jbgp.with_sum
                                    ? ithr_ic
                                    : ithr_ic - 1);
            const size_t c_buf_num_rows = jbgp.nthr_ic_b > 1 ? jbgp.mb : jbgp.M;
            const size_t c_buffer_shift
                    = c_buf_thr_idx * c_buf_num_rows * jbgp.LDC;
            const size_t c_buffer_off = acc_dt_size * c_buffer_shift
                    + (jbgp.nthr_ic_b > 1 ? acc_dt_size * dst_off / dst_dt_size
                                          : 0);
            c_buffer = c_buffer_global + c_buffer_off;
        }

        char *wsp_tile = is_amx ? wsp_tile_base
                        + ithr
                                * jit_brgemm_primitive_conf_t::
                                        tile_wsp_per_thread
                                : nullptr;
        int icb = icc * jbgp.nb_ic_blocking;
        int ic = icb * jbgp.ic_block;

        bool kernel_init = do_init;

        bool is_os_tail = (jbgp.mb - n < jbgp.os_block);
        bool is_oc_tail = (jbgp.oc - oc < jbgp.oc_block);
        bool is_last_ic_chunk = icc == ic_chunks - 1;
        bool is_ic_tail = is_last_ic_chunk && jbgp.K_tail > 0;
        const int remaining_ic_blks
                = (jbgp.use_buffer_a ? utils::rnd_up(jbgp.ic, jbgp.ic_block)
                                     : jbgp.ic)
                - ic;
        const int gemm_batch = nstl::min(
                jbgp.gemm_batch_size, remaining_ic_blks / jbgp.ic_block);

        auto is_bs_tail = (gemm_batch != jbgp.gemm_batch_size);
        int brg_ker_idx = brgemm_inner_product_utils::get_brg_kernel_index(
                jbgp, is_bs_tail, kernel_init, is_os_tail, is_oc_tail, false);
        auto brg_kernel = brg_kernels_[brg_ker_idx].get();

        if (copy_buffer_a) {
            auto src_ptr = src + get_blk_off(src_d, jbgp.src_dt, n, ic);
            copy_data_chunk(copy_src_kernel_, a_buffer, src_ptr,
                    is_os_tail ? jbgp.mb - n : jbgp.os_block, is_last_ic_chunk);
        }
        if (gemm_batch > 0 && brg_kernel != nullptr) {
            if (is_amx && (is_os_tail || is_oc_tail))
                amx_tile_configure(&brg_kernel_palettes_[brg_ker_idx][0]);
            for (int b = 0; b < gemm_batch; b++) {
                auto A_ptr = jbgp.use_buffer_a
                        ? (a_buffer + src_dt_size * b * jbgp.ic_block)
                        : (src
                                + get_blk_off(src_d, jbgp.src_dt, n,
                                        ic + b * jbgp.ic_block));
                addr_batch[b].ptr.A = A_ptr;
                addr_batch[b].ptr.B = weights
                        + get_blk_off(weights_d, jbgp.wei_dt, ocb, icb + b);
            }

            auto ptr_D = dst + dst_off;
            auto ptr_C = use_c_buffer ? c_buffer : ptr_D;

            if (jbgp.nthr_ic_b == 1 && are_post_ops_applicable
                    && is_last_ic_chunk && !is_ic_tail) {
                void *scratch = is_amx
                        ? static_cast<void *>(wsp_tile)
                        : (jbgp.signed_input ? static_cast<void *>(
                                   const_cast<int *>(&compensation[oc]))
                                             : nullptr);
                auto ptr_bias
                        = jbgp.with_bias ? bias + bia_dt_size * oc : nullptr;
                const brgemm_post_ops_data_t post_ops_data {
                        static_cast<const void *>(ptr_bias),
                        &oscales[jbgp.is_oc_scale * oc],
                        post_ops_binary_rhs_arg_vec.data(),
                        static_cast<size_t>(oc), 0, dst};

                brgemm_kernel_execute_postops(brg_kernel, gemm_batch,
                        addr_batch, (void *)ptr_C, (void *)ptr_D, post_ops_data,
                        scratch);
            } else {
                brgemm_kernel_execute(brg_kernel, gemm_batch, addr_batch,
                        (void *)ptr_C, is_amx ? (void *)wsp_tile : nullptr);
            }

            if (is_amx && (is_os_tail || is_oc_tail))
                amx_tile_configure(&brg_kernel_palettes_[base_brg_ker_idx][0]);
        }

        if (is_ic_tail) {
            assert(!jbgp.use_buffer_a);
            int ic_block = gemm_batch * jbgp.K / jbgp.ic_block;
            addr_batch[0].ptr.A = src
                    + get_blk_off(src_d, jbgp.src_dt, n,
                            ic + ic_block * jbgp.ic_block);
            addr_batch[0].ptr.B = weights
                    + get_blk_off(weights_d, jbgp.wei_dt, ocb, icb + ic_block);

            auto use_init_ker = (kernel_init && gemm_batch == 0);
            int brg_ker_idx = brgemm_inner_product_utils::get_brg_kernel_index(
                    jbgp, false, use_init_ker, is_os_tail, is_oc_tail, true);
            auto brg_kernel_ic_tail = brg_kernels_[brg_ker_idx].get();
            if (is_amx)
                amx_tile_configure(&brg_kernel_palettes_[brg_ker_idx][0]);
            auto ptr_D = dst + dst_off;
            auto ptr_C = use_c_buffer ? c_buffer : ptr_D;
            if (jbgp.nthr_ic_b == 1 && are_post_ops_applicable) {
                void *scratch = is_amx
                        ? static_cast<void *>(wsp_tile)
                        : (jbgp.signed_input ? static_cast<void *>(
                                   const_cast<int *>(&compensation[oc]))
                                             : nullptr);
                auto ptr_bias
                        = jbgp.with_bias ? bias + bia_dt_size * oc : nullptr;
                const brgemm_post_ops_data_t post_ops_data {
                        static_cast<const void *>(ptr_bias),
                        &oscales[jbgp.is_oc_scale * oc],
                        post_ops_binary_rhs_arg_vec.data(),
                        static_cast<size_t>(oc), 0, dst};

                brgemm_kernel_execute_postops(brg_kernel_ic_tail, 1, addr_batch,
                        (void *)ptr_C, (void *)ptr_D, post_ops_data, scratch);
            } else {
                brgemm_kernel_execute(brg_kernel_ic_tail, 1, addr_batch,
                        (void *)ptr_C, is_amx ? (void *)wsp_tile : nullptr);
            }
            if (is_amx)
                amx_tile_configure(&brg_kernel_palettes_[base_brg_ker_idx][0]);
        }
    };

    const int os_chunks = div_up(jbgp.nb_os, jbgp.nb_os_blocking);
    const int oc_chunks = div_up(jbgp.nb_oc, jbgp.nb_oc_blocking);
    const int work_amount = oc_chunks * os_chunks;

    const auto init_thr_groups
            = [&](const int ithr, const int nthr, int &nthr_ic, int &nthr_oc_mb,
                      int &ithr_ic, int &ithr_oc_mb) {
                  nthr_ic = jbgp.nthr_ic_b <= nthr ? jbgp.nthr_ic_b : 1;
                  nthr_oc_mb = nthr / nthr_ic;
                  ithr_ic = ithr / nthr_oc_mb;
                  ithr_oc_mb = ithr % nthr_oc_mb;
                  if (ithr_oc_mb >= work_amount || ithr_ic >= ic_chunks
                          || ithr >= rnd_dn(nthr, nthr_ic))
                      return false;
                  return true;
              };

    // If work_amount == 1 we limit num_threads to 1 as parallel(1, ...) does
    // not create parallel section at all. We do not limit num_threads
    // for 1 < work_amount < dnnl_get_max_threads() case to avoid potential
    // overhead on spawning different number of OMP threads from layer to layer.
    const int num_threads = (work_amount == 1 ? 1 : jbgp.nthr);
    parallel(num_threads, [&](const int ithr, const int nthr) {
        int nthr_ic {1}, nthr_oc_mb {1}, ithr_ic {0}, ithr_oc_mb {0};
        bool ok = init_thr_groups(
                ithr, nthr, nthr_ic, nthr_oc_mb, ithr_ic, ithr_oc_mb);
        if (!ok) return;

        int start {0}, end {0};
        balance211(work_amount, nthr_oc_mb, ithr_oc_mb, start, end);

        int icc_start {0}, icc_end {ic_chunks};
        if (nthr_ic > 1)
            balance211(ic_chunks, nthr_ic, ithr_ic, icc_start, icc_end);

        const int icc_work = icc_end - icc_start;

        if (is_amx)
            amx_tile_configure(&brg_kernel_palettes_[base_brg_ker_idx][0]);

        int occ {0}, osc {0};
        nd_iterator_init(start, osc, os_chunks, occ, oc_chunks);
        while (start < end) {
            int ocb_s = occ * jbgp.nb_oc_blocking;
            int ocb_e = nstl::min(ocb_s + jbgp.nb_oc_blocking, jbgp.nb_oc);
            int ocb_work = ocb_e - ocb_s;

            int osb_s = osc * jbgp.nb_os_blocking;
            int osb_e = nstl::min(osb_s + jbgp.nb_os_blocking, jbgp.nb_os);
            int osb_work = osb_e - osb_s;

            // Each thread runs the below loops:
            int loop_start = 0, loop_end = icc_work * osb_work * ocb_work;
            int icc = 0, osb = 0, ocb = 0;

            // If buffer is required, then inner-most loop will be over icc_work
            const bool ocb_inner_most = is_f32 && !jbgp.use_buffer;
            if (ocb_inner_most)
                nd_iterator_init(
                        0, icc, icc_work, osb, osb_work, ocb, ocb_work);
            else
                nd_iterator_init(
                        0, osb, osb_work, ocb, ocb_work, icc, icc_work);

            while (loop_start < loop_end) {
                const int n = (osb + osb_s) * jbgp.os_block;
                const int cur_icc = icc + icc_start;
                const bool copy_buffer_a = jbgp.use_buffer_a
                        && IMPLICATION(ocb_inner_most, ocb == 0);
                ker(ithr_oc_mb, nthr_oc_mb, ithr_ic, n, ocb + ocb_s, cur_icc,
                        cur_icc == icc_start, osb, copy_buffer_a);

                ++loop_start;
                if (ocb_inner_most)
                    nd_iterator_step(
                            icc, icc_work, osb, osb_work, ocb, ocb_work);
                else
                    nd_iterator_step(
                            osb, osb_work, ocb, ocb_work, icc, icc_work);
            }

            ++start;
            nd_iterator_step(osc, os_chunks, occ, oc_chunks);
        }
        if (is_amx) amx_tile_release();
    });

    if (jbgp.nthr_ic_b > 1) {
        assert(jbgp.use_buffer && is_f32);

        const auto get_dst_reduced_off = [&](int ithr_ic, int osb, int ocb) {
            assert(jbgp.nthr_ic_b > 1);
            int os = osb * jbgp.os_block;
            int oc = ocb * jbgp.oc_block;
            const size_t dst_off = get_blk_off(dst_d, jbgp.dst_dt, os, oc);
            if (ithr_ic == 0) return dst_off;
            assert(ithr_ic > 0);
            const size_t ic_buf_idx = jbgp.with_sum ? ithr_ic : ithr_ic - 1;
            return dst_off + (ic_buf_idx * jbgp.mb * jbgp.LDC * acc_dt_size);
        };

        parallel(num_threads, [&](const int ithr, const int nthr) {
            int nthr_ic {1}, nthr_oc_mb {1}, ithr_ic {0}, ithr_oc_mb {0};
            bool ok = init_thr_groups(
                    ithr, nthr, nthr_ic, nthr_oc_mb, ithr_ic, ithr_oc_mb);
            if (!ok) return;

            int ocmb_start {0}, ocmb_end {0};
            int start {0}, end {0};
            balance211(
                    work_amount, nthr_oc_mb, ithr_oc_mb, ocmb_start, ocmb_end);
            balance211(ocmb_end - ocmb_start, nthr_ic, ithr_ic, start, end);

            int occ {0}, osc {0};
            nd_iterator_init(
                    ocmb_start + start, osc, os_chunks, occ, oc_chunks);
            while (start < end) {
                int ocb_s = occ * jbgp.nb_oc_blocking;
                int ocb_e = nstl::min(ocb_s + jbgp.nb_oc_blocking, jbgp.nb_oc);

                int osb_s = osc * jbgp.nb_os_blocking;
                int osb_e = nstl::min(osb_s + jbgp.nb_os_blocking, jbgp.nb_os);

                for (int osb = osb_s; osb < osb_e; ++osb) {
                    int cur_os_block = nstl::min(
                            jbgp.os - osb * jbgp.os_block, jbgp.os_block);
                    const bool is_os_tail = cur_os_block < jbgp.os_block;
                    const int cur_oc_chunk_size
                            = nstl::min(jbgp.LDC, ocb_e * jbgp.oc_block)
                            - ocb_s * jbgp.oc_block;
                    char *dst_reduced = (jbgp.with_sum ? c_buffer_global : dst)
                            + get_dst_reduced_off(0, osb, ocb_s);
                    const size_t os_offset = jbgp.LDC * acc_dt_size;
                    for (int ic_buf = 0; ic_buf < nthr_ic - 1; ++ic_buf) {
                        const char *c_buffer = c_buffer_global
                                + get_dst_reduced_off(ic_buf + 1, osb, ocb_s);
                        for (int os = 0; os < cur_os_block; ++os) {
                            acc_ker_->accumulate(
                                    (float *)(dst_reduced + os * os_offset),
                                    (float *)(c_buffer + os * os_offset),
                                    cur_oc_chunk_size);
                        }
                    }
                    if (are_post_ops_applicable) {
                        for (int ocb = ocb_s; ocb < ocb_e; ++ocb) {
                            const bool is_oc_tail
                                    = (jbgp.oc - ocb * jbgp.oc_block
                                            < jbgp.oc_block);
                            const int brg_ker_idx = brgemm_inner_product_utils::
                                    get_brg_kernel_index(jbgp, false, false,
                                            is_os_tail, is_oc_tail, false);
                            const auto brg_kernel
                                    = brg_kernels_[brg_ker_idx].get();
                            const int os = osb * jbgp.os_block;
                            const int oc = ocb * jbgp.oc_block;
                            const auto ptr_bias = jbgp.with_bias
                                    ? bias + bia_dt_size * oc
                                    : nullptr;
                            auto ptr_D = dst
                                    + get_blk_off(dst_d, jbgp.dst_dt, os, oc);
                            auto ptr_C = (jbgp.with_sum ? c_buffer_global : dst)
                                    + get_dst_reduced_off(0, osb, ocb);

                            char *wsp_tile = is_amx ? wsp_tile_base
                                            + ithr
                                                    * jit_brgemm_primitive_conf_t::
                                                            tile_wsp_per_thread
                                                    : nullptr;

                            void *scratch = is_amx
                                    ? static_cast<void *>(wsp_tile)
                                    : (jbgp.signed_input ? static_cast<void *>(
                                               const_cast<int *>(
                                                       &compensation[oc]))
                                                         : nullptr);

                            const brgemm_post_ops_data_t post_ops_data {
                                    static_cast<const void *>(ptr_bias),
                                    &oscales[jbgp.is_oc_scale * oc],
                                    post_ops_binary_rhs_arg_vec.data(),
                                    static_cast<size_t>(oc), 0, dst, 0, nullptr,
                                    nullptr, nullptr, true /* skip_accm */};
                            brgemm_kernel_execute_postops(brg_kernel, 0,
                                    nullptr, (void *)ptr_C, (void *)ptr_D,
                                    post_ops_data, scratch);
                        }
                    }
                }
                ++start;
                nd_iterator_step(osc, os_chunks, occ, oc_chunks);
            }
        });
    }

    return status::success;
}

template struct brgemm_inner_product_fwd_t<avx512_core>;
template struct brgemm_inner_product_fwd_t<avx512_core_bf16>;
template struct brgemm_inner_product_fwd_t<avx512_core_vnni>;
template struct brgemm_inner_product_fwd_t<avx512_core_bf16_amx_bf16>;
template struct brgemm_inner_product_fwd_t<avx512_core_bf16_amx_int8>;

template <cpu_isa_t isa>
void brgemm_inner_product_bwd_data_t<isa>::execute_backward_data(
        const exec_ctx_t &ctx) const {

    auto diff_dst_ = CTX_IN_MEM(const char *, DNNL_ARG_DIFF_DST);
    auto weights_ = CTX_IN_MEM(const char *, DNNL_ARG_WEIGHTS);
    auto diff_src_ = CTX_OUT_MEM(char *, DNNL_ARG_DIFF_SRC);

    auto diff_src = const_cast<char *>(diff_src_);
    auto weights = const_cast<char *>(weights_);
    auto diff_dst = const_cast<char *>(diff_dst_);

    const memory_desc_wrapper diff_src_d(pd()->diff_src_md());
    const memory_desc_wrapper diff_dst_d(pd()->diff_dst_md());
    const memory_desc_wrapper weights_d(pd()->weights_md(0));

    const auto &jbgp = pd()->jbgp_;

    static constexpr bool is_amx = (isa == avx512_core_bf16_amx_bf16);
    const bool is_f32 = everyone_is(f32, jbgp.src_dt, jbgp.wei_dt, jbgp.dst_dt);
    const bool is_bf16_bf16_out = everyone_is(bf16, jbgp.wei_dt, jbgp.dst_dt)
            && everyone_is(bf16, jbgp.src_dt);
    const bool is_bf16_f32_out = everyone_is(bf16, jbgp.wei_dt, jbgp.dst_dt)
            && everyone_is(f32, jbgp.src_dt);

    const dim_t wei_dt_size = types::data_type_size(jbgp.wei_dt);

    memory_tracking::grantor_t scratchpad = ctx.get_scratchpad_grantor();
    brgemm_batch_element_t *addr_batch_global
            = scratchpad.template get<brgemm_batch_element_t>(
                    key_brgemm_primitive_batch);
    char *c_buffer_global = (jbgp.use_buffer)
            ? scratchpad.template get<char>(key_brgemm_primitive_buffer)
            : nullptr;
    char *b_buffer_global = jbgp.use_buffer_b
            ? scratchpad.template get<char>(key_brgemm_primitive_buffer_b)
            : nullptr;
    char *a_buffer_global = jbgp.use_buffer_a
            ? scratchpad.template get<char>(key_brgemm_primitive_buffer_a)
            : nullptr;
    auto wsp_tile_base = is_amx
            ? ctx.get_scratchpad_grantor().template get<char>(
                    key_conv_amx_tile_buffer)
            : nullptr;

    const int oc_chunks = div_up(jbgp.nb_oc, jbgp.nb_oc_blocking);
    bool is_os_tail = (jbgp.mb < jbgp.os_block);
    bool is_ic_tail = (jbgp.ic < jbgp.ic_block);
    bool is_oc_tail = (jbgp.oc < jbgp.oc_block) && !jbgp.use_buffer_a;

    const dim_t acc_dt_sz = types::data_type_size(jbgp.acc_dt);
    const dim_t src_dt_sz = types::data_type_size(jbgp.src_dt);

    const int base_brg_ker_idx = brgemm_inner_product_utils::
            get_brg_kernel_index( // TODO: Can be calculated on initialization stage
                    jbgp, false, false, is_os_tail, is_ic_tail, is_oc_tail);

    const int os_chunks = div_up(jbgp.nb_os, jbgp.nb_os_blocking);
    const int work_amount = jbgp.nb_ic * os_chunks;
    const int num_threads = (work_amount == 1 ? 1 : jbgp.nthr);

    const auto get_weights_ptr = [&](int icb, int ocb) {
        int fwd_ic_block = (is_amx) ? 2 * jbgp.simd_w : jbgp.simd_w;
        int fwd_oc_block = 0;
        switch (jbgp.wei_tag) {
            case OI16i64o:
            case OIw16i64o:
            case OIhw16i64o:
            case OIdhw16i64o:
            case OI8i64o2i:
            case OIw8i64o2i:
            case OIhw8i64o2i:
            case OIdhw8i64o2i:
            case OI16i64o2i:
            case OIw16i64o2i:
            case OIhw16i64o2i:
            case OIdhw16i64o2i: fwd_oc_block = 4 * jbgp.simd_w; break;
            case OI16i32o:
            case OIw16i32o:
            case OIhw16i32o:
            case OIdhw16i32o:
            case OI8i32o2i:
            case OIw8i32o2i:
            case OIhw8i32o2i:
            case OIdhw8i32o2i:
            case OI16i32o2i:
            case OIw16i32o2i:
            case OIhw16i32o2i:
            case OIdhw16i32o2i: fwd_oc_block = 2 * jbgp.simd_w; break;
            default: fwd_oc_block = jbgp.simd_w;
        };
        int fwd_icb = icb * jbgp.ic_block / fwd_ic_block;
        int fwd_ocb = ocb * jbgp.oc_block / fwd_oc_block;
        char *ptr_wei_local = weights
                + get_blk_off(weights_d, jbgp.wei_dt, fwd_ocb, fwd_icb);

        int fwd_ocb_simd = (ocb * jbgp.oc_block) % fwd_oc_block;
        int fwd_icb_simd = (icb * jbgp.ic_block) % fwd_ic_block;
        int blk_sz = jbgp.wei_dt == data_type::bf16 ? 2 : 1;

        return ptr_wei_local
                + wei_dt_size
                * (fwd_icb_simd / blk_sz * blk_sz * fwd_oc_block
                        + blk_sz * fwd_ocb_simd);
    };

    const auto transform_b_chunk
            = [&](char *tr_wei, const char *wei, int trans_batch, int current_N,
                      int current_K) {
                  auto ctx = jit_brgemm_trans_wei_t::ctx_t();
                  ctx.src = (void *)wei;
                  ctx.tr_src = (void *)tr_wei;
                  ctx.current_gemm_batch = trans_batch;
                  ctx.current_N = current_N;
                  ctx.current_K = current_K;
                  (*trans_B_kernel_)(&ctx);
              };

    const auto ker = [&](int ithr_ic_mb, int nthr_ic_mb, int ithr_oc,
                             int nthr_oc, int n, int icb, int occ, bool do_init,
                             bool do_b_transpose) {
        const int ithr = nthr_ic_mb * ithr_oc + ithr_ic_mb;
        brgemm_batch_element_t *addr_batch
                = addr_batch_global + ithr * jbgp.adjusted_batch_size;

        const int ic = icb * jbgp.ic_block;
        const int ocb = occ * jbgp.nb_oc_blocking;
        const int oc = ocb * jbgp.oc_block;
        const size_t dsrc_off = get_blk_off(diff_src_d, jbgp.src_dt, n, ic);
        const int adj_buffers = (jbgp.src_dt == f32) ? 1 : 0;
        const size_t c_buf_shift = jbgp.nthr_oc_b > 1
                ? (ithr_oc - adj_buffers)
                        * static_cast<size_t>(jbgp.mb * jbgp.LDC)
                : ithr * static_cast<size_t>(jbgp.LDC * jbgp.M);
        const size_t c_buf_off
                = types::data_type_size(jbgp.acc_dt) * c_buf_shift
                + (jbgp.nthr_oc_b > 1 ? acc_dt_sz * dsrc_off / src_dt_sz : 0);
        bool use_c_buf = false;
        if ((is_f32 || is_bf16_f32_out) && jbgp.use_buffer) {
            use_c_buf = (jbgp.nthr_oc_b == 1 || ithr_oc > 0);
        } else if (is_bf16_bf16_out && jbgp.use_buffer) {
            if (jbgp.nthr_oc_b > 1)
                use_c_buf = true;
            else
                use_c_buf = (jbgp.nthr_oc_b == 1 || ithr_oc > 0);
        }

        const size_t a_buffer_size_per_thr
                = jbgp.os_block * jbgp.LDA * types::data_type_size(jbgp.dst_dt);
        char *c_buffer = use_c_buf ? c_buffer_global + c_buf_off : nullptr;
        char *a_buffer = jbgp.use_buffer_a
                ? a_buffer_global + ithr * a_buffer_size_per_thr
                : diff_dst;
        char *wsp_tile = is_amx ? wsp_tile_base
                        + ithr
                                * jit_brgemm_primitive_conf_t::
                                        tile_wsp_per_thread
                                : nullptr;

        bool kernel_init = do_init;

        const bool is_os_tail = (jbgp.mb - n < jbgp.os_block);
        const bool is_ic_tail = (jbgp.ic - ic < jbgp.ic_block);
        const bool is_last_oc_chunk = occ == oc_chunks - 1;
        const bool is_oc_tail = is_last_oc_chunk && jbgp.K_tail > 0;

        const int rnd_oc
                = rnd_up(jbgp.oc, jbgp.use_buffer_a ? jbgp.oc_block : 1);
        const int nb_oc_b
                = nstl::min((rnd_oc - oc) / jbgp.oc_block, jbgp.nb_oc_blocking);

        auto is_bs_tail = (nb_oc_b != jbgp.nb_oc_blocking);
        const int brg_ker_idx
                = brgemm_inner_product_utils::get_brg_kernel_index(jbgp,
                        is_bs_tail, kernel_init, is_os_tail, is_ic_tail, false);
        auto brg_kernel = brg_kernels_[brg_ker_idx].get();

        const int size_B = jbgp.LDB * rnd_up(jbgp.K, 2);

        const size_t b_buf_shift = jbgp.ip_bwd_d_global_b_transpose
                ? icb * jbgp.nb_oc + ocb
                : ithr * jbgp.gemm_batch_size;
        const size_t b_buf_off = wei_dt_size * b_buf_shift * size_B;
        char *b_buffer = b_buffer_global + b_buf_off;

        char *ptr_D = diff_src + dsrc_off;
        char *ptr_C = use_c_buf ? c_buffer : ptr_D;

        if (jbgp.use_buffer_a)
            copy_data_chunk(copy_diff_dst_kernel_, a_buffer,
                    diff_dst + get_blk_off(diff_dst_d, jbgp.dst_dt, n, oc),
                    is_os_tail ? jbgp.os - n : jbgp.os_block, is_last_oc_chunk);

        if (nb_oc_b > 0 && brg_kernel != nullptr) {
            if (is_amx && (is_os_tail || is_ic_tail))
                amx_tile_configure(&brg_kernel_palettes_[brg_ker_idx][0]);

            for (int oc_block = 0; oc_block < nb_oc_b; oc_block++) {
                addr_batch[oc_block].ptr.A = jbgp.use_buffer_a ? a_buffer
                                + oc_block * jbgp.oc_block
                                        * types::data_type_size(jbgp.dst_dt)
                                                               : diff_dst
                                + get_blk_off(diff_dst_d, jbgp.dst_dt, n,
                                        oc + oc_block * jbgp.oc_block);
                addr_batch[oc_block].ptr.B
                        = b_buffer + wei_dt_size * (oc_block * size_B);
                if (!jbgp.ip_bwd_d_global_b_transpose && do_b_transpose)
                    transform_b_chunk((char *)addr_batch[oc_block].ptr.B,
                            get_weights_ptr(icb, ocb + oc_block), 1,
                            is_ic_tail ? jbgp.ic % jbgp.ic_block
                                       : jbgp.ic_block,
                            jbgp.oc_block);
            }

            if (jbgp.use_buffer && (jbgp.nthr_oc_b <= 1 || num_threads == 1)
                    && is_last_oc_chunk && !is_oc_tail) {
                void *scratch
                        = is_amx ? static_cast<void *>(wsp_tile) : nullptr;
                const brgemm_post_ops_data_t empty_po_data {};
                brgemm_kernel_execute_postops(brg_kernel, nb_oc_b, addr_batch,
                        (void *)c_buffer, (void *)ptr_D, empty_po_data,
                        scratch);

            } else {
                brgemm_kernel_execute(brg_kernel, nb_oc_b, addr_batch,
                        (void *)ptr_C, is_amx ? (void *)wsp_tile : nullptr);
            }
            if (is_amx && (is_os_tail || is_ic_tail))
                amx_tile_configure(&brg_kernel_palettes_[base_brg_ker_idx][0]);
        }
        if (is_oc_tail) {
            assert(!jbgp.use_buffer_a);

            const int oc_block = nb_oc_b;
            addr_batch[0].ptr.A = diff_dst
                    + get_blk_off(diff_dst_d, jbgp.dst_dt, n,
                            oc + oc_block * jbgp.oc_block);
            addr_batch[0].ptr.B = b_buffer + wei_dt_size * (oc_block * size_B);
            if (!jbgp.ip_bwd_d_global_b_transpose && do_b_transpose) {
                transform_b_chunk((char *)addr_batch[0].ptr.B,
                        get_weights_ptr(icb, ocb + oc_block), 1,
                        is_ic_tail ? jbgp.ic % jbgp.ic_block : jbgp.ic_block,
                        jbgp.K_tail);
            }

            auto use_init_ker = (kernel_init && nb_oc_b == 0);
            const int brg_kernel_oc_tail_idx
                    = brgemm_inner_product_utils::get_brg_kernel_index(jbgp,
                            false, use_init_ker, is_os_tail, is_ic_tail, true);
            auto brg_kernel_oc_tail
                    = brg_kernels_[brg_kernel_oc_tail_idx].get();
            if (is_amx)
                amx_tile_configure(
                        &brg_kernel_palettes_[brg_kernel_oc_tail_idx][0]);
            if (jbgp.use_buffer && jbgp.nthr_oc_b <= 1) {
                void *scratch
                        = is_amx ? static_cast<void *>(wsp_tile) : nullptr;
                const brgemm_post_ops_data_t empty_po_data {};
                brgemm_kernel_execute_postops(brg_kernel_oc_tail, 1, addr_batch,
                        (void *)c_buffer, (void *)ptr_D, empty_po_data,
                        scratch);

            } else {
                brgemm_kernel_execute(brg_kernel_oc_tail, 1, addr_batch,
                        (void *)ptr_C, is_amx ? (void *)wsp_tile : nullptr);
            }
            if (is_amx)
                amx_tile_configure(&brg_kernel_palettes_[base_brg_ker_idx][0]);
        }
    };

    if (jbgp.ip_bwd_d_global_b_transpose && jbgp.use_buffer_b) {
        assert(IMPLICATION(
                jbgp.ip_bwd_d_global_b_transpose, jbgp.nthr_oc_b == 1));
        parallel(num_threads, [&](const int ithr, const int nthr) {
            int start {0}, end {0};
            int max_ch_block = nstl::max(jbgp.ic_block, jbgp.oc_block);
            int ic_chunk_sz = max_ch_block / jbgp.ic_block;
            int oc_chunk_sz = max_ch_block / jbgp.oc_block;
            int nc_ic = utils::div_up(jbgp.nb_ic, ic_chunk_sz);
            int nc_oc = utils::div_up(jbgp.nb_oc, oc_chunk_sz);
            int transp_work_amount = nc_ic * nc_oc;
            balance211(transp_work_amount, nthr, ithr, start, end);
            int icc, occ;
            nd_iterator_init(start, icc, nc_ic, occ, nc_oc);
            while (start < end) {
                int icb_start = icc * ic_chunk_sz;
                int icb_end = nstl::min((icc + 1) * ic_chunk_sz, jbgp.nb_ic);
                int ocb_start = occ * oc_chunk_sz;
                int ocb_end = nstl::min((occ + 1) * oc_chunk_sz, jbgp.nb_oc);
                for_(int icb = icb_start; icb < icb_end; icb++)
                for (int ocb = ocb_start; ocb < ocb_end; ocb++) {
                    int ic = icb * jbgp.ic_block;
                    int oc = ocb * jbgp.oc_block;
                    bool is_ic_tail = (jbgp.ic - ic < jbgp.ic_block);
                    bool is_oc_tail = (jbgp.oc - oc < jbgp.oc_block);
                    const int size_B = jbgp.LDB * rnd_up(jbgp.K, 2);
                    char *b_buffer = b_buffer_global
                            + wei_dt_size
                                    * ((dim_t)icb * jbgp.nb_oc * size_B
                                            + (dim_t)ocb * size_B);

                    transform_b_chunk(b_buffer, get_weights_ptr(icb, ocb), 1,
                            is_ic_tail ? jbgp.ic % jbgp.ic_block
                                       : jbgp.ic_block,
                            is_oc_tail ? jbgp.oc % jbgp.oc_block
                                       : jbgp.oc_block);
                }
                ++start;
                nd_iterator_step(icc, nc_ic, occ, nc_oc);
            }
        });
    }

    parallel(num_threads, [&](const int ithr, const int nthr) {
        const int nthr_oc = jbgp.nthr_oc_b <= nthr ? jbgp.nthr_oc_b : 1;
        const int nthr_ic_mb = nthr / nthr_oc;
        const int ithr_ic_mb = ithr % nthr_ic_mb;
        const int ithr_oc = ithr / nthr_ic_mb;
        if (ithr_ic_mb >= work_amount || ithr_oc >= oc_chunks
                || ithr >= rnd_dn(nthr, nthr_oc))
            return;

        int start {0}, end {0};
        balance211(work_amount, nthr_ic_mb, ithr_ic_mb, start, end);
        int occ_start {0}, occ_end {oc_chunks};
        if (nthr_oc > 1)
            balance211(oc_chunks, nthr_oc, ithr_oc, occ_start, occ_end);

        if (is_amx)
            amx_tile_configure(&brg_kernel_palettes_[base_brg_ker_idx][0]);

        int icb {0}, oss {0};
        nd_iterator_init(start, oss, os_chunks, icb, jbgp.nb_ic);
        while (start < end) {
            const int nb_os_blocking
                    = nstl::min(jbgp.nb_os - oss * jbgp.nb_os_blocking,
                            jbgp.nb_os_blocking);
            const int occ_work = occ_end - occ_start;
            const int loop_iteration = nb_os_blocking * occ_work;

            for (int iter = 0; iter < loop_iteration; ++iter) {
                int osb = 0, occ = occ_start;
                if (jbgp.use_buffer || !is_f32) {
                    osb += iter / occ_work;
                    occ += iter % occ_work;
                } else {
                    occ += iter / nb_os_blocking;
                    osb += iter % nb_os_blocking;
                }
                int n = (oss * jbgp.nb_os_blocking + osb) * jbgp.os_block;
                ker(ithr_ic_mb, nthr_ic_mb, ithr_oc, nthr_oc, n, icb, occ,
                        occ == occ_start, osb == 0 || occ_work > 1);
            }
            ++start;
            nd_iterator_step(oss, os_chunks, icb, jbgp.nb_ic);
        }
        if (is_amx) amx_tile_release();
    });

    if (jbgp.nthr_oc_b > 1) {
        parallel(num_threads, [&](const int ithr, const int nthr) {
            const int nthr_oc = jbgp.nthr_oc_b <= nthr ? jbgp.nthr_oc_b : 1;
            if (nthr_oc <= 1) return;

            const int ddst_elems = jbgp.LDC * jbgp.os;
            const int reduce_chunk_size = 64;
            int start {0}, end {0};
            balance211(div_up(ddst_elems, reduce_chunk_size), nthr, ithr, start,
                    end);
            const dim_t reduce_start = start * reduce_chunk_size;
            const dim_t reduce_finish
                    = nstl::min(end * reduce_chunk_size, ddst_elems);
            if (reduce_finish <= reduce_start) return;
            const dim_t elems_to_reduce = reduce_finish - reduce_start;
            const dim_t acc_dt_sz = types::data_type_size(jbgp.acc_dt);

            char *dsrc_reduced = diff_src + src_dt_sz * reduce_start;
            char *c_buffer_start = c_buffer_global + acc_dt_sz * reduce_start;

            float *out_buffer = (is_f32 || is_bf16_f32_out)
                    ? (float *)dsrc_reduced
                    : (float *)c_buffer_start;
            int oc_buf_idx = is_bf16_bf16_out;
            int oc_buf_end = !is_bf16_bf16_out;
            for (int oc_buf = oc_buf_idx; oc_buf < nthr_oc - oc_buf_end;
                    oc_buf++) {
                const dim_t c_buf_offt = acc_dt_sz
                        * (oc_buf * jbgp.os * jbgp.LDC + reduce_start);
                char *c_buffer = c_buffer_global + c_buf_offt;

                acc_ker_->accumulate((float *)out_buffer, (float *)c_buffer,
                        elems_to_reduce);
                if (is_bf16_bf16_out && oc_buf == (nthr_oc - oc_buf_end) - 1) {
                    cvt_float_to_bfloat16((bfloat16_t *)dsrc_reduced,
                            (const float *)out_buffer, elems_to_reduce);
                }
            }
        });
    }
}

template struct brgemm_inner_product_bwd_data_t<avx512_core>;
template struct brgemm_inner_product_bwd_data_t<avx512_core_bf16_amx_bf16>;
template struct brgemm_inner_product_bwd_data_t<avx512_core_bf16>;

template <cpu_isa_t isa>
struct brgemm_inner_product_bwd_weights_t<isa>::thread_info_t {
    const char *src;
    const char *diff_dst;
    char *diff_weights;
    char *diff_bias;

    const memory_tracking::grantor_t scratchpad;

    char *buffer_a = nullptr;
    char *buffer_b = nullptr;
    char *buffer_c = nullptr;
    char *buffer_bias = nullptr;
    char *wsp_tile_base = nullptr;

    int ithr;
    int ithr_ic_c, ithr_oc_c, ithr_os_c;
    int nthr;
    int nthr_ic_c, nthr_oc_c, nthr_os_c;

    int os_c_start = 0, os_c_end = 0, os_c_work;
    int oc_c_start = 0, oc_c_end = 0, oc_c_work;
    int ic_c_start = 0, ic_c_end = 0, ic_c_work;
    simple_barrier::ctx_t *barrier_ctx;

    thread_info_t(const brgemm_inner_product_bwd_weights_t *self,
            const exec_ctx_t &ctx, int ithr)
        : scratchpad(ctx.get_scratchpad_grantor()), ithr(ithr) {

        constexpr bool is_amx = (isa == avx512_core_bf16_amx_bf16);

        src = CTX_IN_MEM(const char *, DNNL_ARG_SRC);
        diff_dst = CTX_IN_MEM(const char *, DNNL_ARG_DIFF_DST);
        diff_weights = CTX_OUT_MEM(char *, DNNL_ARG_DIFF_WEIGHTS);
        diff_bias = CTX_OUT_MEM(char *, DNNL_ARG_DIFF_BIAS);
        const auto &jbgp = self->pd()->jbgp_;

        buffer_c = (jbgp.use_buffer)
                ? scratchpad.template get<char>(key_brgemm_primitive_buffer)
                : nullptr;

        buffer_bias = (jbgp.with_bias
                              && (jbgp.bia_dt == data_type::bf16
                                      || jbgp.nthr_mb > 1))
                ? scratchpad.template get<char>(key_iprod_bias_bf16_convert_wsp)
                : nullptr;

        buffer_a = scratchpad.template get<char>(key_brgemm_primitive_buffer_a);
        buffer_b = jbgp.use_buffer_b
                ? scratchpad.template get<char>(key_brgemm_primitive_buffer_b)
                : nullptr;

        wsp_tile_base = is_amx
                ? ctx.get_scratchpad_grantor().template get<char>(
                        key_conv_amx_tile_buffer)
                : nullptr;

        nthr = jbgp.nthr;
        nthr_ic_c = jbgp.nthr_ic_b;
        nthr_oc_c = jbgp.nthr_oc_b;
        nthr_os_c = jbgp.nthr_mb;

        ithr_ic_c = ithr % nthr_ic_c;
        ithr_oc_c = ithr / nthr_ic_c % nthr_oc_c;
        ithr_os_c = ithr / nthr_ic_c / nthr_oc_c;

        int oc_chunks = utils::div_up(jbgp.nb_oc, jbgp.nb_oc_blocking);
        int ic_chunks = utils::div_up(jbgp.nb_ic, jbgp.nb_ic_blocking);
        int os_chunks = utils::div_up(jbgp.nb_os, jbgp.nb_os_blocking);

        /* reduction dimension */
        balance211(os_chunks, nthr_os_c, ithr_os_c, os_c_start, os_c_end);
        os_c_work = os_c_end - os_c_start;

        balance211(oc_chunks, nthr_oc_c, ithr_oc_c, oc_c_start, oc_c_end);
        oc_c_work = oc_c_end - oc_c_start;

        balance211(ic_chunks, nthr_ic_c, ithr_ic_c, ic_c_start, ic_c_end);
        ic_c_work = ic_c_end - ic_c_start;

        if (dnnl_thr_syncable())
            barrier_ctx = scratchpad.template get<simple_barrier::ctx_t>(
                    key_conv_wei_bia_reduction_bctx);
    }
};

template <cpu_isa_t isa>
void brgemm_inner_product_bwd_weights_t<isa>::transform_matrix_a_chunk(
        char *tr_src, const char *src, int trans_batch, int current_m,
        int current_k) const {
    auto ctx = jit_brgemm_trans_src_t::ctx_t();
    ctx.src = (void *)src;
    ctx.tr_src = (void *)tr_src;
    ctx.current_gemm_batch = trans_batch;
    ctx.current_M = current_m;
    ctx.current_K = current_k;
    (*trans_A_kernel_)(&ctx);
}

template <cpu_isa_t isa>
void brgemm_inner_product_bwd_weights_t<isa>::transform_matrix_b_chunk(
        char *tr_diff_dst, const char *diff_dst, int trans_batch,
        int current_col_size, int current_row_size) const {
    auto ctx = jit_brgemm_trans_to_vnni_t::ctx_t();
    ctx.src = (void *)diff_dst;
    ctx.tr_src = (void *)tr_diff_dst;
    ctx.current_gemm_batch = trans_batch;
    ctx.current_col_size = current_col_size;
    ctx.current_row_size = current_row_size;
    (*trans_B_kernel_)(&ctx);
}

template <cpu_isa_t isa>
void brgemm_inner_product_bwd_weights_t<isa>::transpose_matrix_c_chunk(
        const thread_info_t *ti, const int ocb, const int icb, int oc_size,
        int ic_size, bool is_reduction) const {
    const auto &jbgp = pd()->jbgp_;

    if (isa == avx512_core_bf16_amx_bf16) {
        auto p = jit_amx_ip_trans_diff_wei::ctx_t();

        const dim_t ext_nb_ic = div_up(jbgp.ic, ext_ic_block_);
        dim_t icb_shift = (icb * (jbgp.ic_block / ext_ic_block_))
                * ext_ic_block_ * ext_oc_block_;

        dim_t ocb_shift = (ocb * (jbgp.oc_block / ext_oc_block_)) * ext_nb_ic
                * ext_ic_block_ * ext_oc_block_;
        dim_t out_offset = ocb_shift + icb_shift;

        p.src = get_wei_acc_ptr(ti, ocb, icb, 0);
        p.dst = (void *)(ti->diff_weights
                + types::data_type_size(jbgp.wei_dt) * out_offset);

        p.last_ic_block = (jbgp.ic <= ext_ic_block_
                                  || (jbgp.nb_ic > 1 && icb == jbgp.nb_ic - 1))
                ? 1
                : 0;
        p.last_oc_block = (jbgp.oc <= ext_oc_block_
                                  || (jbgp.nb_oc > 1 && ocb == jbgp.nb_oc - 1))
                ? 1
                : 0;
        (*diff_wei_trans_kernel_)(&p);
    } else {
        auto ctx = jit_brgemm_trans_to_vnni_t::ctx_t();
        ctx.src = (void *)(get_wei_acc_ptr(ti, ocb, icb, 0));

        ctx.tr_src = (void *)(ti->diff_weights
                + types::data_type_size(jbgp.wei_dt)
                        * get_wei_offset(ocb, icb));

        ctx.current_gemm_batch = 1;
        ctx.current_col_size = oc_size;
        ctx.current_row_size = ic_size;
        (*trans_C_kernel_)(&ctx);
    }
}

template <cpu_isa_t isa>
dim_t brgemm_inner_product_bwd_weights_t<isa>::get_wei_offset(
        int ocb, int icb) const {
    if (isa == avx512_core_bf16_amx_bf16) {
        const auto &jbgp = pd()->jbgp_;
        const dim_t offset
                = jbgp.kd * jbgp.kh * jbgp.kw * jbgp.ic_block * jbgp.oc_block;
        return (ocb * jbgp.nb_ic + icb) * offset;
    } else {
        const memory_desc_wrapper diff_weights_d(pd()->diff_weights_md(0));
        return diff_weights_d.blk_off(ocb, icb);
    }
}

template <cpu_isa_t isa>
char *brgemm_inner_product_bwd_weights_t<isa>::get_wei_acc_ptr(
        const thread_info_t *ti, int ocb, int icb,
        int reduction_buf_idx) const {
    constexpr bool is_amx_bf16 = (isa == avx512_core_bf16_amx_bf16);

    const auto &jbgp = pd()->jbgp_;
    const int reduction_buf_start_idx = jbgp.wei_dt == f32;
    // reduction_buf_idx argument allows manually set up required reduction
    // buffer index, required for reduction and transform diff_weights parts.
    // It has value -1 by default. If reduction_buf_idx < 0 then ti->ithr_os_c
    // is used for calculation of the current reduction index.
    const int buf_idx = reduction_buf_idx >= 0
            ? reduction_buf_idx
            : (ti->ithr_os_c - reduction_buf_start_idx);
    const size_t acc_dt_size = types::data_type_size(jbgp.acc_dt);

    if ((jbgp.nthr_mb > 1 && buf_idx < 0)
            || (jbgp.wei_dt == jbgp.acc_dt && reduction_buf_idx < 0
                    && ti->ithr_os_c == 0)) {
        MAYBE_UNUSED(reduction_buf_idx);
        const int icb_scale = (!is_amx_bf16 || jbgp.wei_dt == jbgp.acc_dt)
                ? jbgp.ic_block / jbgp.simd_w
                : 1;
        const memory_desc_wrapper diff_weights_d(pd()->diff_weights_md(0));
        return (char *)ti->diff_weights
                + get_blk_off(
                        diff_weights_d, jbgp.wei_dt, ocb, icb * icb_scale);
    }

    if (!jbgp.use_buffer) return nullptr;

    const int ocb_l = ocb % jbgp.nb_oc_blocking;
    const int icb_l = icb % jbgp.nb_ic_blocking;

    if (jbgp.nthr_mb > 1 || jbgp.harness == harness_mb_reduction) {
        const size_t icc = icb / jbgp.nb_ic_blocking;
        const size_t occ = ocb / jbgp.nb_oc_blocking;
        const size_t num_ic_chunks = div_up(jbgp.nb_ic, jbgp.nb_ic_blocking);
        const size_t num_oc_chunks = div_up(jbgp.nb_oc, jbgp.nb_oc_blocking);
        const size_t block_size = acc_dt_size * jbgp.ic_block * jbgp.oc_block;
        const size_t chunk_size
                = block_size * jbgp.nb_ic_blocking * jbgp.nb_oc_blocking;
        const size_t reduction_buf_shift
                = num_ic_chunks * num_oc_chunks * chunk_size * buf_idx;
        return ti->buffer_c + reduction_buf_shift
                + (occ * num_ic_chunks + icc) * chunk_size
                + (ocb_l * jbgp.nb_ic_blocking + icb_l) * block_size;
    } else if (jbgp.nthr_mb == 1) {
        MAYBE_UNUSED(reduction_buf_idx);
        const size_t blk_size = acc_dt_size * jbgp.ic_block * jbgp.oc_block;
        const size_t buf_size_per_thread
                = blk_size * jbgp.nb_ic_blocking * jbgp.nb_oc_blocking;
        const size_t offset_within_thread_buf
                = blk_size * (jbgp.nb_ic_blocking * ocb_l + icb_l);
        const size_t offset
                = ti->ithr * buf_size_per_thread + offset_within_thread_buf;
        return ti->buffer_c + offset;
    }

    assert(!"unsupported case");
    return nullptr;
};

template <cpu_isa_t isa>
void brgemm_inner_product_bwd_weights_t<isa>::compute_diff_weights_and_bias(
        const thread_info_t *ti) const {
    auto diff_dst = const_cast<char *>(ti->diff_dst);
    auto diff_bias = ti->diff_bias;

    const memory_desc_wrapper diff_dst_d(pd()->diff_dst_md());

    const auto &jbgp = pd()->jbgp_;

    const size_t bia_dt_size
            = jbgp.with_bias ? types::data_type_size(jbgp.bia_dt) : 0;
    const size_t acc_dt_size = types::data_type_size(jbgp.acc_dt);

    const int oc_chunk_sz = jbgp.oc_block * jbgp.nb_oc_blocking;

    brgemm_batch_element_t *addr_batch_global
            = ti->scratchpad.template get<brgemm_batch_element_t>(
                    key_brgemm_primitive_batch);

    constexpr bool is_amx_bf16 = (isa == avx512_core_bf16_amx_bf16);
    char *a_buffer_global = ti->buffer_a;
    char *b_buffer_global = ti->buffer_b;
    char *wsp_tile_global = (is_amx_bf16) ? ti->wsp_tile_base : nullptr;
    int os_chunks = utils::div_up(jbgp.nb_os, jbgp.nb_os_blocking);

    const auto get_bia_acc_ptr = [&](int oc) {
        const int reduction_buf_start_idx = jbgp.bia_dt == f32;
        if (jbgp.bia_dt == data_type::bf16
                || (jbgp.nthr_mb > 1
                        && ti->ithr_os_c >= reduction_buf_start_idx)) {
            return ti->buffer_bias
                    + acc_dt_size * (ti->ithr_os_c - reduction_buf_start_idx)
                    * jbgp.oc
                    + acc_dt_size * oc;
        } else {
            return ti->diff_bias + bia_dt_size * oc;
        }
    };

    const auto ker = [&](const int osc, const int icb, const int ocb) {
        dim_t os_chunks_per_thr = utils::div_up(os_chunks, jbgp.nthr_mb);
        int ic_chunks = utils::div_up(jbgp.nb_ic, jbgp.nb_ic_blocking);
        dim_t ic_chunks_per_thr = utils::div_up(ic_chunks, jbgp.nthr_ic_b);

        int osc_l_idx = osc - ti->os_c_start;
        int icb_l_idx = icb - ti->ic_c_start * jbgp.nb_ic_blocking;
        int ocb_l_idx = ocb - ti->oc_c_start * jbgp.nb_oc_blocking;
        dim_t a_buf_idx = osc_l_idx * ic_chunks_per_thr * jbgp.nb_ic_blocking
                + icb_l_idx;
        int b_buf_idx = osc_l_idx;

        constexpr int tile_size = 1024;

        brgemm_batch_element_t *addr_batch
                = addr_batch_global + ti->ithr * jbgp.adjusted_batch_size;
        const int size_A = jbgp.LDA * jbgp.M;
        const int size_B = jbgp.LDB * rnd_up(jbgp.K, 2);
        char *a_buffer = a_buffer_global
                + types::data_type_size(jbgp.src_dt)
                        * ((ti->ithr * os_chunks_per_thr * ic_chunks_per_thr
                                           * jbgp.nb_ic_blocking
                                   + a_buf_idx)
                                * jbgp.gemm_batch_size * jbgp.os_block
                                * jbgp.ic_block);
        char *b_buffer = b_buffer_global ? b_buffer_global
                        + types::data_type_size(jbgp.dst_dt)
                                * ((ti->ithr * os_chunks_per_thr + b_buf_idx)
                                                * jbgp.gemm_batch_size
                                                * jbgp.os_block * jbgp.LDB
                                        + (ocb_l_idx % jbgp.nb_oc_blocking)
                                                * jbgp.oc_block)
                                         : nullptr;

        char *wsp_tile = is_amx_bf16
                ? wsp_tile_global + ti->ithr * 4 * tile_size
                : nullptr;
        int ic = icb * jbgp.ic_block;
        int oc = ocb * jbgp.oc_block;
        int n = osc * jbgp.nb_os_blocking * jbgp.os_block;

        bool kernel_init = (osc == ti->os_c_start);

        bool is_os_tail = jbgp.mb - n < jbgp.os_block * jbgp.nb_os_blocking;
        bool is_ic_tail = jbgp.ic - ic < jbgp.ic_block;
        bool is_oc_tail = jbgp.oc - oc < jbgp.oc_block;
        const int oc_chunk_tail = jbgp.oc % oc_chunk_sz;
        const bool is_last_oc_chunk = jbgp.oc - oc < oc_chunk_sz;
        const int curr_oc_chunk_sz = oc_chunk_tail > 0 && is_last_oc_chunk
                ? oc_chunk_tail
                : oc_chunk_sz;

        const bool transform_weights_to_vnni = jbgp.wei_dt == bf16
                && (jbgp.nthr_mb == 1 || os_chunks == 1)
                && osc == (os_chunks - 1);

        auto nb_os_b = is_os_tail ? (jbgp.mb - n) / jbgp.os_block
                                  : jbgp.nb_os_blocking;

        auto is_bs_tail = (nb_os_b != jbgp.nb_os_blocking);
        const int brg_ker_idx
                = brgemm_inner_product_utils::get_brg_kernel_index(jbgp,
                        is_bs_tail, kernel_init, is_ic_tail, is_oc_tail, false);
        auto brg_kernel = brg_kernels_[brg_ker_idx].get();

        if (kernel_init && (is_ic_tail || is_oc_tail))
            utils::array_set(get_wei_acc_ptr(ti, ocb, icb), 0,
                    types::data_type_size(jbgp.acc_dt) * jbgp.ic_block
                            * jbgp.oc_block);
        if (nb_os_b > 0 && brg_kernel != nullptr) {
            if (is_amx_bf16)
                amx_tile_configure(&brg_kernel_palettes_[brg_ker_idx][0]);
            if (jbgp.use_buffer_a && ocb_l_idx == 0) {
                const memory_desc_wrapper src_d(pd()->src_md());
                auto src_ptr = ti->src
                        + types::data_type_size(jbgp.src_dt)
                                * src_d.blk_off(n, ic);
                transform_matrix_a_chunk(a_buffer, src_ptr, nb_os_b,
                        is_ic_tail ? jbgp.ic % jbgp.ic_block : jbgp.ic_block,
                        jbgp.os_block);
            }

            if (jbgp.use_buffer_b && icb_l_idx == 0
                    && ocb_l_idx % jbgp.nb_oc_blocking == 0) {
                auto diff_dst_ptr = diff_dst
                        + types::data_type_size(jbgp.dst_dt)
                                * diff_dst_d.blk_off(n, oc);
                transform_matrix_b_chunk(b_buffer, diff_dst_ptr, nb_os_b,
                        curr_oc_chunk_sz, jbgp.os_block);
            }

            for (int os_block = 0; os_block < nb_os_b; os_block++) {
                auto a_ptr = a_buffer
                        + types::data_type_size(jbgp.src_dt) * os_block
                                * size_A;
                addr_batch[os_block].ptr.A = a_ptr;
                auto diff_dst_ptr = diff_dst
                        + types::data_type_size(jbgp.dst_dt)
                                * diff_dst_d.blk_off(
                                        n + os_block * jbgp.os_block, oc);
                if (jbgp.use_buffer_b) {
                    auto b_ptr = b_buffer
                            + types::data_type_size(jbgp.dst_dt) * os_block
                                    * size_B;
                    addr_batch[os_block].ptr.B = b_ptr;
                } else {
                    addr_batch[os_block].ptr.B = diff_dst_ptr;
                }
                if (jbgp.with_bias && icb == 0) {
                    brgemm_kernel_diff_bias_t p;
                    auto bias_ptr = diff_bias + bia_dt_size * oc;
                    p.ptr_diff_dst = (void *)addr_batch[os_block].ptr.B;
                    p.ptr_diff_bias_acc = (void *)get_bia_acc_ptr(oc);
                    p.ptr_diff_bias = (void *)bias_ptr;
                    bool is_first = kernel_init && os_block == 0;
                    bool is_last = (jbgp.nthr_mb == 1 || os_chunks == 1)
                            && osc == os_chunks - 1 && os_block == nb_os_b - 1
                            && !is_os_tail;
                    p.flags = 0 | (is_first ? FLAG_REDUCE_FIRST : 0)
                            | (is_last ? FLAG_REDUCE_LAST : 0);

                    (*kernels_db_[false][is_oc_tail])(&p);
                }
            }
            brgemm_kernel_execute(brg_kernel, nb_os_b, addr_batch,
                    (void *)get_wei_acc_ptr(ti, ocb, icb), wsp_tile);
        }

        if (is_os_tail) {
            int os_block = nb_os_b;
            auto a_ptr = a_buffer
                    + types::data_type_size(jbgp.src_dt) * os_block
                            * jbgp.ic_block * jbgp.os_block;
            if (jbgp.use_buffer_a && ocb_l_idx == 0) {
                const memory_desc_wrapper src_d(pd()->src_md());
                auto src_ptr = ti->src
                        + types::data_type_size(jbgp.src_dt)
                                * src_d.blk_off(
                                        n + os_block * jbgp.os_block, ic);
                transform_matrix_a_chunk(a_ptr, src_ptr, 1,
                        is_ic_tail ? jbgp.ic % jbgp.ic_block : jbgp.ic_block,
                        jbgp.mb % jbgp.os_block);
            }

            addr_batch[0].ptr.A = a_ptr;
            auto diff_dst_ptr = diff_dst
                    + types::data_type_size(jbgp.dst_dt)
                            * diff_dst_d.blk_off(
                                    n + os_block * jbgp.os_block, oc);
            if (jbgp.use_buffer_b) {
                auto b_ptr = b_buffer
                        + types::data_type_size(jbgp.dst_dt) * os_block
                                * jbgp.os_block * jbgp.LDB;
                if (icb_l_idx == 0 && ocb_l_idx % jbgp.nb_oc_blocking == 0)
                    transform_matrix_b_chunk(b_ptr, diff_dst_ptr, 1,
                            curr_oc_chunk_sz, jbgp.mb % jbgp.os_block);
                addr_batch[0].ptr.B = b_ptr;
            } else {
                addr_batch[0].ptr.B = diff_dst_ptr;
            }

            if (jbgp.with_bias && icb == 0) {
                brgemm_kernel_diff_bias_t p;
                auto bias_ptr = diff_bias + bia_dt_size * oc;
                p.ptr_diff_dst = (void *)addr_batch[0].ptr.B;
                p.ptr_diff_bias_acc = (void *)get_bia_acc_ptr(oc);
                p.ptr_diff_bias = (void *)bias_ptr;
                bool is_first = kernel_init && os_block == 0;
                bool is_last = (jbgp.nthr_mb == 1 || os_chunks == 1)
                        && osc == os_chunks - 1;
                p.flags = 0 | (is_first ? FLAG_REDUCE_FIRST : 0)
                        | (is_last ? FLAG_REDUCE_LAST : 0);

                (*kernels_db_[true][is_oc_tail])(&p);
            }

            auto use_init_ker = (kernel_init && nb_os_b == 0);
            const int brg_ker_idx_os_tail
                    = brgemm_inner_product_utils::get_brg_kernel_index(jbgp,
                            false, use_init_ker, is_ic_tail, is_oc_tail, true);
            auto brg_kernel_os_tail = brg_kernels_[brg_ker_idx_os_tail].get();
            if (brg_kernel_os_tail != nullptr) {
                if (is_amx_bf16)
                    amx_tile_configure(
                            &brg_kernel_palettes_[brg_ker_idx_os_tail][0]);
                brgemm_kernel_execute(brg_kernel_os_tail, 1, addr_batch,
                        (void *)get_wei_acc_ptr(ti, ocb, icb), wsp_tile);
            }
        }

        if (transform_weights_to_vnni) {
            transpose_matrix_c_chunk(ti, ocb, icb,
                    is_oc_tail ? jbgp.oc % jbgp.oc_block : jbgp.oc_block,
                    is_ic_tail ? jbgp.ic % jbgp.ic_block : jbgp.ic_block);
        }
    };

    const auto occ_work = (ti->oc_c_end - ti->oc_c_start);
    const auto icc_work = (ti->ic_c_end - ti->ic_c_start);
    const auto osc_work = (ti->os_c_end - ti->os_c_start);

    auto loop_idx = 0;
    const auto loop_end = occ_work * icc_work * osc_work;

    int occ_idx = 0, icc_idx = 0, osc_idx = 0;
    //TODO: Introduce loop order enum for brgemm-based implementations
    const bool osc_loop_outermost = jbgp.harness == harness_mb_reduction;
    if (osc_loop_outermost)
        nd_iterator_init(loop_idx, osc_idx, osc_work, occ_idx, occ_work,
                icc_idx, icc_work);
    else
        nd_iterator_init(loop_idx, occ_idx, occ_work, icc_idx, icc_work,
                osc_idx, osc_work);

    while (loop_idx < loop_end) {
        const int occ = ti->oc_c_start + occ_idx;
        const int icc = ti->ic_c_start + icc_idx;
        const int osc = ti->os_c_start + osc_idx;

        const int ocb_work = nstl::min(
                jbgp.nb_oc_blocking, jbgp.nb_oc - occ * jbgp.nb_oc_blocking);
        const int icb_work = nstl::min(
                jbgp.nb_ic_blocking, jbgp.nb_ic - icc * jbgp.nb_ic_blocking);

        for_(int ocb = 0; ocb < ocb_work; ocb++)
        for (int icb = 0; icb < icb_work; icb++) {
            ker(osc, icc * jbgp.nb_ic_blocking + icb,
                    occ * jbgp.nb_oc_blocking + ocb);
        }

        ++loop_idx;
        if (osc_loop_outermost)
            nd_iterator_step(
                    osc_idx, osc_work, occ_idx, occ_work, icc_idx, icc_work);
        else
            nd_iterator_step(
                    occ_idx, occ_work, icc_idx, icc_work, osc_idx, osc_work);
    }
    if (is_amx_bf16) amx_tile_release();
}

template <cpu_isa_t isa>
void brgemm_inner_product_bwd_weights_t<
        isa>::reduce_and_convert_diff_weights_and_bias(const thread_info_t *ti)
        const {
    const auto &jbgp = pd()->jbgp_;

    if (dnnl_thr_syncable() && jbgp.nthr > 1)
        simple_barrier::barrier(ti->barrier_ctx, jbgp.nthr);
    if (ti->nthr_os_c == 1) return;

    const bool is_bf16_out = jbgp.wei_dt == data_type::bf16;
    const int icb_scale = (!is_bf16_out) ? jbgp.ic_block / jbgp.simd_w : 1;

    const int icb_work = ti->ic_c_work * jbgp.nb_ic_blocking;
    const int ocb_work = ti->oc_c_work * jbgp.nb_oc_blocking;
    const int work = ocb_work * icb_work;

    int os_chunks = utils::div_up(jbgp.nb_os, jbgp.nb_os_blocking);
    int reduce_buffers = nstl::min(ti->nthr_os_c, os_chunks);
    int reduce_buf_idx_start = is_bf16_out;
    int reduce_buf_idx_end = reduce_buffers - !is_bf16_out;

    int start = 0, end = 0;
    balance211(work, ti->nthr_os_c, ti->ithr_os_c, start, end);
    if (start == end) return;

    int icb_l = 0, ocb_l = 0;
    const int acc_size = jbgp.ic_block * jbgp.oc_block;

    for (int ir = reduce_buf_idx_start; ir < reduce_buf_idx_end; ++ir) {
        int counter = start;
        nd_iterator_init(start, ocb_l, ocb_work, icb_l, icb_work);
        while (counter < end) {
            const int ocb = ti->oc_c_start * jbgp.nb_oc_blocking + ocb_l;
            const int icb = ti->ic_c_start * jbgp.nb_ic_blocking + icb_l;
            char *wei_to_reduce = get_wei_acc_ptr(ti, ocb, icb, ir);
            const memory_desc_wrapper diff_weights_d(pd()->diff_weights_md(0));
            char *wei_reduced = is_bf16_out ? get_wei_acc_ptr(ti, ocb, icb, 0)
                                            : ti->diff_weights
                            + get_blk_off(diff_weights_d, jbgp.wei_dt, ocb,
                                    icb * icb_scale);
            acc_ker_->accumulate(
                    (float *)(wei_reduced), (float *)(wei_to_reduce), acc_size);
            if (is_bf16_out && ir + 1 == reduce_buf_idx_end) {
                transpose_matrix_c_chunk(ti, ocb, icb * icb_scale,
                        jbgp.oc_block, jbgp.ic_block, true);
            }
            ++counter;
            nd_iterator_step(ocb_l, ocb_work, icb_l, icb_work);
        }
    }

    if (jbgp.with_bias && ti->ithr_ic_c == 0 && ti->ic_c_work > 0
            && ti->ithr_os_c == 0 && ti->os_c_work > 0 && ti->oc_c_work > 0) {
        const bool is_bf16_bias = jbgp.bia_dt == data_type::bf16;
        float *bias_reduced = is_bf16_bias ? (float *)ti->buffer_bias
                                           : (float *)ti->diff_bias;
        int reduce_buf_idx_start = is_bf16_bias;
        int reduce_buf_idx_end = reduce_buffers - 1;
        int oc_chunk_size = jbgp.nb_oc_blocking * jbgp.oc_block;
        int oc = ti->oc_c_start * oc_chunk_size;
        int acc_size = nstl::min(ti->oc_c_work * oc_chunk_size, jbgp.oc - oc);

        int ir = reduce_buf_idx_start;
        for (; ir < reduce_buf_idx_end; ++ir) {
            float *bias_to_reduce = (float *)ti->buffer_bias + ir * jbgp.oc;
            acc_ker_->accumulate(
                    &bias_reduced[oc], &bias_to_reduce[oc], acc_size);
        }

        if (is_bf16_bias) {
            float *bias_to_reduce = (float *)ti->buffer_bias + ir * jbgp.oc;
            add_floats_and_cvt_to_bfloat16((bfloat16_t *)(ti->diff_bias) + oc,
                    &bias_reduced[oc], &bias_to_reduce[oc], acc_size);
        }
    }
}

template <cpu_isa_t isa>
void brgemm_inner_product_bwd_weights_t<isa>::execute_backward_weights(
        const exec_ctx_t &ctx) const {
    const auto &jbgp = pd()->jbgp_;

    if (dnnl_thr_syncable() && jbgp.nthr > 1) {
        auto scratchpad = ctx.get_scratchpad_grantor();
        simple_barrier::ctx_init(scratchpad.template get<simple_barrier::ctx_t>(
                key_conv_wei_bia_reduction_bctx));
    }

    parallel(jbgp.nthr, [&](const int ithr, const int nthr) {
        thread_info_t thread_info(this, ctx, ithr);
        compute_diff_weights_and_bias(&thread_info);

        if (dnnl_thr_syncable()) {
            reduce_and_convert_diff_weights_and_bias(&thread_info);
        }
    });

    if (!dnnl_thr_syncable()) {
        parallel(jbgp.nthr, [&](const int ithr, const int nthr) {
            thread_info_t thread_info(this, ctx, ithr);
            reduce_and_convert_diff_weights_and_bias(&thread_info);
        });
    }
}

template struct brgemm_inner_product_bwd_weights_t<avx512_core_bf16_amx_bf16>;
template struct brgemm_inner_product_bwd_weights_t<avx512_core_bf16>;
template struct brgemm_inner_product_bwd_weights_t<avx512_core>;

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

// vim: et ts=4 sw=4 cindent cino+=l0,\:4,N-s
