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

#include "gpu/ocl/gen9_global_pooling.hpp"

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

int calculate_spatial_chunk(const pool_conf_t &conf, engine_t *engine) {
    auto *compute_engine = utils::downcast<compute::compute_engine_t *>(engine);
    const int hw_threads = compute_engine->device_info()->hw_threads();
    const bool is_xe_hp_plus = compute_engine->is_xe_hp()
            || compute_engine->is_xe_hpg() || compute_engine->is_xe_hpc();

    const int spatial_dim = conf.id * conf.ih * conf.iw;
    int chunk_size = spatial_dim;

    // Experimentally selected values for XeHP family
    const int desired_wi_per_thread = is_xe_hp_plus && conf.is_plain ? 1024 : 4;

    const auto get_work_items_num = [&]() {
        return conf.c * conf.mb * utils::div_up(spatial_dim, chunk_size);
    };
    while (get_work_items_num() < hw_threads * desired_wi_per_thread
            && chunk_size > 1) {
        chunk_size = utils::div_up(chunk_size, 2);
    }
    return chunk_size;
}

static status_t init_conf_common(pool_conf_t &conf, offsets_t &off,
        const pooling_pd_t *pd, engine_t *engine) {
    using namespace dnnl::impl::format_tag;

    set_default_pool_conf(conf, *pd->desc(), *pd->invariant_src_md(),
            *pd->invariant_dst_md(), *pd->attr());

    if (conf.id != conf.kd || conf.iw != conf.kw || conf.ih != conf.kh
            || conf.od * conf.ow * conf.oh != 1)
        return status::unimplemented;
    if (!conf.is_backward) return status::unimplemented;

    const memory_desc_wrapper src_mdw(pd->invariant_src_md());
    const memory_desc_wrapper dst_mdw(pd->invariant_dst_md());
    const auto &padded_src_dims = src_mdw.padded_dims();
    const auto &padded_dst_dims = dst_mdw.padded_dims();
    if (utils::array_product(padded_src_dims + 2, conf.ndims - 2)
                    != conf.id * conf.ih * conf.iw
            || utils::array_product(padded_dst_dims + 2, conf.ndims - 2)
                    != conf.od * conf.oh * conf.ow)
        return status::unimplemented;

    set_offsets(src_mdw, off.src_off);
    set_offsets(dst_mdw, off.dst_off);

    auto *compute_engine = utils::downcast<compute::compute_engine_t *>(engine);

    conf.is_plain = src_mdw.is_plain();
    conf.global_pool_spatial_chunk = calculate_spatial_chunk(conf, engine);

    const int spatial_dim_padded = utils::rnd_up(
            conf.id * conf.ih * conf.iw, conf.global_pool_spatial_chunk);
    conf.dispatch = compute_engine->create_dispatch(src_mdw.md_);
    conf.dispatch.define_dim("MB", 0, conf.mb_padded);
    conf.dispatch.define_dim("C", 1, conf.c_padded);
    conf.dispatch.define_dim(
            "SPATIAL", 2, spatial_dim_padded, conf.global_pool_spatial_chunk);
    conf.dispatch.generate();

    conf.attr_info = attr_info_t::create(pd->attr());

    return status::success;
};

static status_t init_kernel_ctx_common(compute::kernel_ctx_t &kernel_ctx,
        const pool_conf_t &conf, const offsets_t &off,
        const post_ops_t &post_ops) {
    using namespace dnnl::impl::alg_kind;
    kernel_ctx.set_data_type(conf.src_dt);

    kernel_ctx.define_int("NDIMS", conf.ndims);
    kernel_ctx.define_int("MB", conf.mb);
    kernel_ctx.define_int("C", conf.c);
    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("SPATIAL_DIM", conf.id * conf.ih * conf.iw);
    kernel_ctx.define_int("SPATIAL_CHUNK", conf.global_pool_spatial_chunk);

    kernel_ctx.define_int("ALG_MAX", (conf.alg == pooling_max));
    kernel_ctx.define_int(
            "ALG_AVG_NP", (conf.alg == pooling_avg_exclude_padding));
    kernel_ctx.define_int(
            "ALG_AVG_P", (conf.alg == pooling_avg_include_padding));
    kernel_ctx.define_int("NEED_ZERO_PADDING",
            (conf.mb != conf.mb_padded || conf.c != conf.c_padded));

    def_attr_info(kernel_ctx, conf.attr_info, post_ops);

    def_offsets(off.src_off, kernel_ctx, "SRC", conf.ndims);
    def_offsets(off.dst_off, kernel_ctx, "DST", conf.ndims);

    def_memory_desc_info(kernel_ctx, conf.src_md_info, "SRC");
    def_memory_desc_info(kernel_ctx, conf.dst_md_info, "DST");

    def_dispatch(kernel_ctx, conf.dispatch);

    return status::success;
}

status_t gen9_global_pooling_bwd_t::pd_t::init_conf(engine_t *engine) {
    return init_conf_common(conf, off, this, engine);
}

status_t gen9_global_pooling_bwd_t::pd_t::init_kernel_ctx(
        compute::kernel_ctx_t &kernel_ctx) const {
    return init_kernel_ctx_common(kernel_ctx, conf, off, attr()->post_ops_);
}

status_t gen9_global_pooling_bwd_t::execute_backward(
        const exec_ctx_t &ctx) const {
    auto &diff_src = CTX_OUT_STORAGE(DNNL_ARG_DIFF_SRC);
    auto &diff_dst = CTX_IN_STORAGE(DNNL_ARG_DIFF_DST);
    auto &ws = CTX_IN_STORAGE(DNNL_ARG_WORKSPACE);

    compute::kernel_arg_list_t arg_list;
    arg_list.set(0, diff_src);
    arg_list.set(1, ws);
    arg_list.set(2, diff_dst);

    auto nd_range = pd()->conf.dispatch.nd_range();

    return parallel_for(ctx, nd_range, kernel_, arg_list);
}

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