/*******************************************************************************
* 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 <algorithm>
#include "gpu/ocl/ocl_stream.hpp"
#include "gpu/ocl/xe_lp_x8s8x_1x1_convolution.hpp"

namespace dnnl {
namespace impl {
namespace gpu {
namespace ocl {

status_t xe_lp_x8s8x_1x1_convolution_fwd_t::pd_t::init_conf(engine_t *engine) {
    using namespace format_tag;

    const convolution_desc_t &cd = *desc();
    const memory_desc_wrapper src_mdw(src_md());
    const memory_desc_wrapper weights_mdw(weights_md());
    const memory_desc_wrapper dst_mdw(dst_md());
    const memory_desc_wrapper bias_mdw(weights_md(1));
    auto dev_info = utils::downcast<compute::compute_engine_t *>(engine)
                            ->device_info();

    set_default_conf(conf, cd, *src_md(), *weights_md(), *dst_md(),
            *weights_md(1), *attr());

    conf.is_nhwc
            = src_mdw.matches_one_of_tag(nwc, nhwc, ndhwc) != format_tag::undef
            || dst_mdw.matches_one_of_tag(nwc, nhwc, ndhwc)
                    != format_tag::undef;

    if (conf.is_depthwise || conf.kw != 1 || conf.kh != 1 || conf.kd != 1
            || (conf.with_groups && conf.ngroups > 1
                    && (conf.oc % 32 != 0 || conf.ic % 32 != 0)))
        return status::unimplemented;

    conf.src_data_type = src_mdw.data_type();
    conf.dst_data_type = dst_mdw.data_type();

    conf.mb_block = 32;
    conf.oc_block = 32;
    conf.ic_block = 32;
    conf.nchunk = utils::div_up(conf.oc * conf.ngroups, conf.oc_block);
    int ow = conf.ow;
    int oh = conf.oh;
    int od = conf.od;
    const bool is_stride1
            = conf.stride_d == 1 && conf.stride_h == 1 && conf.stride_w == 1;
    const bool is_padded = conf.l_pad > 0 || conf.t_pad > 0 || conf.f_pad > 0;

    // TODO: fix r-padded shapes issue in 1x1 kernel
    if (conf.r_pad > 0) return status::unimplemented;

    if (is_stride1 || is_padded) {
        // reshape to nCx32c
        ow = ow * oh * od;
        oh = od = 1;
    }
    conf.sp = ow;

    if ((conf.mb % 16 == 0) && !conf.is_nhwc && !is_padded) {
        conf.mb_block = 32;
        conf.sp_block = 1;
    } else {
        conf.mb_block = 1;
        conf.sp_block = 4;
        auto approx_clocks = [&](const int block) {
            int ic_chunks = utils::div_up(conf.ic, conf.ic_block);
            bool use_slm = utils::div_up(conf.ow, block) % 8 == 0;
            int mem_clocks = ic_chunks * (16 - use_slm * 6)
                    + block / 2 * (ic_chunks + 1);
            int compute_clocks = 32 * block * ic_chunks;
            int num_threads = conf.nchunk * conf.mb * od * oh
                    * utils::div_up(ow, block);
            return utils::div_up(num_threads, dev_info->hw_threads())
                    * (compute_clocks + mem_clocks);
        };
        auto clock_compare = [&](const int &block1, const int &block2) {
            return approx_clocks(block1) < approx_clocks(block2);
        };
        std::vector<int> sorted_blocks = {4, 8, 12, 16};
        std::sort(sorted_blocks.begin(), sorted_blocks.end(), clock_compare);
        conf.sp_block = sorted_blocks[0];
    }
    conf.src_data_type = src_mdw.data_type();
    conf.dst_data_type = dst_mdw.data_type();

    const int ow_group = (utils::div_up(ow, conf.sp_block) % 8) ? 1 : 8;

    conf.sub_group_size = 8;
    conf.lws_d[0] = conf.sub_group_size;
    conf.lws_d[1] = ow_group;
    conf.lws_d[2] = 1;

    const int num_sp_threads = utils::div_up(ow, conf.sp_block) * oh * od;

    conf.gws_d[0] = utils::rnd_up(conf.nchunk * 8, conf.lws_d[0]);
    conf.gws_d[1] = utils::rnd_up(num_sp_threads, conf.lws_d[1]);
    conf.gws_d[2] = (conf.mb_block == 32 ? 2 : 1)
            * utils::div_up(conf.mb, conf.mb_block);

    conf.with_bias = cd.bias_desc.format_kind != format_kind::undef;

    format_tag_t src_tag, dst_tag, wei_tag;

    if (conf.is_nhwc) {
        src_tag = utils::pick(conf.ndims - 3, nwc, nhwc);
        dst_tag = utils::pick(conf.ndims - 3, nwc, nhwc);
    } else {
        if (conf.mb_block == 32) {
            src_tag = utils::pick(conf.ndims - 3, NCw32n32c, NChw32n32c);
            dst_tag = utils::pick(conf.ndims - 3, NCw32n32c, NChw32n32c);
        } else {
            src_tag = utils::pick(conf.ndims - 3, nCw32c, nChw32c);
            dst_tag = utils::pick(conf.ndims - 3, nCw32c, nChw32c);
        }
    }

    wei_tag = conf.with_groups
            ? utils::pick(conf.ndims - 3, gOIw4o8i8o4i, gOIhw4o8i8o4i)
            : utils::pick(conf.ndims - 3, OIw4o8i8o4i, OIhw4o8i8o4i);

    conf.src_tag = src_mdw.format_kind() == format_kind::any
            ? src_tag
            : src_mdw.matches_one_of_tag(src_tag);
    conf.wei_tag = weights_mdw.format_kind() == format_kind::any
            ? wei_tag
            : weights_mdw.matches_one_of_tag(wei_tag);
    conf.dst_tag = dst_mdw.format_kind() == format_kind::any
            ? dst_tag
            : dst_mdw.matches_one_of_tag(dst_tag);

    if (conf.src_tag != src_tag || conf.wei_tag != wei_tag
            || conf.dst_tag != dst_tag)
        return status::unimplemented;

    return status::success;
}

status_t xe_lp_x8s8x_1x1_convolution_fwd_t::pd_t::init_kernel_ctx(
        compute::kernel_ctx_t &kernel_ctx) const {
    kernel_ctx.define_int("G", conf.ngroups);
    kernel_ctx.define_int("MB", conf.mb);
    kernel_ctx.define_int("IC", conf.ic_without_padding);
    kernel_ctx.define_int("ID", conf.id);
    kernel_ctx.define_int("IH", conf.ih);
    kernel_ctx.define_int("IW", conf.iw);
    kernel_ctx.define_int("OC", conf.oc_without_padding);
    kernel_ctx.define_int("OD", conf.od);
    kernel_ctx.define_int("OH", conf.oh);
    kernel_ctx.define_int("OW", conf.ow);
    kernel_ctx.define_int("KD", conf.kd);
    kernel_ctx.define_int("KH", conf.kh);
    kernel_ctx.define_int("KW", conf.kw);
    kernel_ctx.define_int("SD", conf.stride_d);
    kernel_ctx.define_int("SH", conf.stride_h);
    kernel_ctx.define_int("SW", conf.stride_w);
    kernel_ctx.define_int("PD", conf.f_pad);
    kernel_ctx.define_int("PH", conf.t_pad);
    kernel_ctx.define_int("PW", conf.l_pad);

    kernel_ctx.define_int("SP_BLOCK", conf.sp_block);
    kernel_ctx.define_int("MB_BLOCK", conf.mb_block);
    kernel_ctx.define_int("OC_BLOCK", conf.oc_block);
    kernel_ctx.define_int("IC_BLOCK", conf.ic_block);

    kernel_ctx.define_int("WITH_BIAS", conf.with_bias);

    def_attr_info(
            kernel_ctx, conf.attr_info, attr()->post_ops_, &(dst_md()->dims));

    kernel_ctx.define_int("SUB_GROUP_SIZE", conf.sub_group_size);

    kernel_ctx.define_int("LWS_0", conf.lws_d[0]);
    kernel_ctx.define_int("LWS_1", conf.lws_d[1]);
    kernel_ctx.define_int("LWS_2", conf.lws_d[2]);

    kernel_ctx.define_int("OC_NCHUNK", utils::div_up(conf.oc, conf.oc_block));
    kernel_ctx.define_int("IC_NCHUNK", utils::div_up(conf.ic, conf.ic_block));

    kernel_ctx.define_int(
            "INT8_WEI_SLM", utils::div_up(conf.sp, conf.sp_block) % 8 == 0);
    kernel_ctx.define_int("SP_TAIL",
            utils::div_up(conf.sp, conf.sp_block) % conf.lws_d[1] == 0);
    kernel_ctx.define_int("OUT_SP_TAIL", conf.sp % conf.sp_block);

    kernel_ctx.define_int("WEI_4O8I8O4I", 1);

    kernel_ctx.set_data_type(conf.dst_data_type);
    def_data_type(kernel_ctx, conf.src_data_type, "SRC");
    def_data_type(kernel_ctx, conf.dst_data_type, "DST");
    def_data_type(kernel_ctx,
            conf.attr_info.sum_data_type == dnnl_data_type_undef
                    ? conf.dst_data_type
                    : conf.attr_info.sum_data_type,
            "SUM");
    kernel_ctx.add_option("-Dcl_intel_subgroups_char");

    return status::success;
}

void xe_lp_x8s8x_1x1_convolution_fwd_t::pd_t::init_scratchpad() {
    if (conf.attr_info.with_src_zpoints) {
        size_t size = conf.ngroups * utils::rnd_up(conf.oc, 32);

        auto scratchpad = scratchpad_registry().registrar();
        scratchpad.book(memory_tracking::names::key_conv_wei_reduction, size,
                types::data_type_size(data_type::s32), OCL_BUFFER_ALIGNMENT);
    }
}

status_t xe_lp_x8s8x_1x1_convolution_fwd_t::execute_forward(
        const exec_ctx_t &ctx) const {
    auto &src = CTX_IN_STORAGE(DNNL_ARG_SRC);
    auto &weights = CTX_IN_STORAGE(DNNL_ARG_WEIGHTS);
    auto &bias = CTX_IN_STORAGE(DNNL_ARG_BIAS);
    auto &oscales = CTX_IN_STORAGE(DNNL_ARG_ATTR_OUTPUT_SCALES);
    auto &dst = CTX_OUT_STORAGE(DNNL_ARG_DST);
    auto &src_zpoints
            = CTX_IN_STORAGE(DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC);
    auto &dst_zpoints
            = CTX_IN_STORAGE(DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_DST);

    const auto &conf = pd()->conf;

    std::unique_ptr<memory_storage_t> temp_src_compensation;
    if (conf.attr_info.with_src_zpoints) {
        temp_src_compensation = ctx.get_scratchpad_grantor().get_memory_storage(
                memory_tracking::names::key_conv_wei_reduction);

        compute::kernel_arg_list_t arg_list;
        arg_list.set(0, src_zpoints);
        arg_list.set(1, weights);
        arg_list.set(2, *temp_src_compensation);

        auto nd_range = compute::nd_range_t(
                {8, utils::div_up(conf.oc, 32), conf.ngroups}, {8, 1, 1});
        status_t status = parallel_for(
                ctx, nd_range, src_compensation_kernel_, arg_list);
        if (status != status::success) return status::runtime_error;
    }

    compute::kernel_arg_list_t arg_list;
    arg_list.set(0, src);
    arg_list.set(1, weights);
    arg_list.set(2, bias);
    arg_list.set(3, dst);

    unsigned arg_idx = append_post_ops_to_arg_list(
            ctx, arg_list, 4, pd()->attr()->post_ops_);

    if (conf.attr_info.with_common_oscales
            || conf.attr_info.with_per_oc_oscales) {
        arg_list.set(arg_idx++, oscales);
    } else {
        arg_list.set(arg_idx++, memory_storage_t::empty_storage());
    }

    if (conf.attr_info.with_src_zpoints)
        arg_list.set(arg_idx++, *temp_src_compensation);
    else
        arg_list.set(arg_idx++, memory_storage_t::empty_storage());

    if (conf.attr_info.with_dst_zpoints)
        arg_list.set(arg_idx++, dst_zpoints);
    else
        arg_list.set(arg_idx++, memory_storage_t::empty_storage());

    auto nd_range = compute::nd_range_t(conf.gws_d, conf.lws_d);
    status_t status = parallel_for(ctx, nd_range, kernel_, arg_list);

    if (!post_ops_preserves_zeroes(ctx, pd()->attr()->post_ops_)) {
        ctx.zero_pad_output(DNNL_ARG_DST);
    }

    return status;
}

} // namespace ocl
} // namespace gpu
} // namespace impl
} // namespace dnnl
