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

namespace mkldnn {
namespace impl {
namespace cpu {

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

template <cpu_isa_t isa, bool with_relu, impl::data_type_t src_type, data_type_t dst_type>
void _jit_uni_x8s8s32x_convolution_fwd_t<isa, with_relu, src_type, dst_type>::execute_forward() {
    auto src = reinterpret_cast<const src_data_t *>(this->input_memory(0));
    auto weights = reinterpret_cast<const wei_data_t *>(this->input_memory(1));
    auto bias = reinterpret_cast<const char *>(this->input_memory(2));
    auto dst = reinterpret_cast<dst_data_t *>(this->memory());

    const memory_desc_wrapper src_d(conf_.src_pd());
    const memory_desc_wrapper dst_d(conf_.dst_pd());
    const memory_desc_wrapper weights_d(conf_.weights_pd(0));
    const memory_desc_wrapper bias_d(conf_.weights_pd(1));

    const auto &jcp = kernel_->jcp;

    size_t offset = (size_t)jcp.ngroups * rnd_up(jcp.oc, jcp.oc_block) * rnd_up(jcp.ic, jcp.ic_block) * jcp.kh * jcp.kw;
    auto w = const_cast<wei_data_t *>(weights);
    int32_t* compensation = (jcp.signed_input) ? reinterpret_cast<int32_t *>(&w[offset]) : 0;

    const size_t bia_dt_size = conf_.with_bias() ? types::data_type_size(conf_.cdesc()->bias_desc.data_type) : 0;
    float* scales = conf_.attr()->output_scales_.scales_;

    int ocb_work = div_up(jcp.nb_oc, jcp.nb_oc_blocking);
    const size_t work_amount = jcp.mb * jcp.ngroups * ocb_work * jcp.oh;

    auto ker = [&](const int ithr, const int nthr) {
        size_t start{0}, end{0};
        balance211(work_amount, nthr, ithr, start, end);

        size_t n{0}, g{0}, ocbb{0}, oh{0};
        nd_iterator_init(start, n, jcp.mb, g, jcp.ngroups, ocbb, ocb_work,
                         oh, jcp.oh);
        for (size_t iwork = start; iwork < end; ++iwork) {
            int ocb = ocbb * jcp.nb_oc_blocking;
            int ocb_num = jcp.nb_oc_blocking;

            jit_conv_call_s par_conv = {};

            const int ij = oh * jcp.stride_h;
            const int i_t_overflow = nstl::min(jcp.kh, div_up(nstl::max(0, jcp.t_pad - ij), (jcp.dilate_h+1)));
            const int i_b_overflow = nstl::min(jcp.kh, div_up(nstl::max(jcp.ih, ij + (jcp.kh-1) * (jcp.dilate_h+1) -
                                               jcp.t_pad+1) - jcp.ih, (jcp.dilate_h + 1)));

            const size_t _oc = g * jcp.nb_oc + ocb;
            const size_t _ic = g * jcp.nb_ic;

            const int ih = nstl::max(ij - jcp.t_pad + i_t_overflow * (jcp.dilate_h + 1), 0);
            par_conv.src = &src[src_d.blk_off(n, _ic*jcp.ic_block, ih, 0)];

            size_t dst_off = dst_d.blk_off(n, _oc*jcp.oc_block, oh, 0);
            par_conv.dst = &dst[dst_off];

            const int wh = (!jcp.signed_input) ? i_t_overflow : 0;
            par_conv.filt = &weights[conf_.with_groups()
                                ? weights_d.blk_off(g, ocb, 0, wh, 0)
                                : weights_d.blk_off(ocb, 0, wh, 0)];

            if (bias)
                par_conv.bias = &bias[bias_d.blk_off(_oc * jcp.oc_block*bia_dt_size)];

            par_conv.oc_work =
                    nstl::min((ocb + ocb_num) * jcp.oc_block, jcp.oc) - ocb*jcp.oc_block;

            par_conv.kw_padding = 0;
            const int kh_padding = jcp.kh - i_t_overflow - i_b_overflow;
            par_conv.kh_padding = nstl::max(0, kh_padding);

            par_conv.scales = (jcp.signed_input) ? &local_scales_[jcp.is_oc_scale * _oc * jcp.oc_block]
                                                 : &scales[jcp.is_oc_scale * _oc * jcp.oc_block];

            par_conv.compensation = (jcp.signed_input) ? compensation + _oc * jcp.oc_block : 0;
            par_conv.t_overflow = i_t_overflow;
            par_conv.b_overflow = i_b_overflow;

            kernel_->jit_ker(&par_conv);
            nd_iterator_step(n, jcp.mb, g, jcp.ngroups, ocbb, ocb_work, oh, jcp.oh);
        }
    };

    parallel(0, ker);
}

template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, true, data_type::u8, data_type::u8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, true, data_type::u8, data_type::s8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, true, data_type::u8, data_type::s32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, true, data_type::u8, data_type::f32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, false, data_type::u8, data_type::u8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, false, data_type::u8, data_type::s8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, false, data_type::u8, data_type::s32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, false, data_type::u8, data_type::f32>::execute_forward();

template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, true, data_type::s8, data_type::u8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, true, data_type::s8, data_type::s8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, true, data_type::s8, data_type::s32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, true, data_type::s8, data_type::f32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, false, data_type::s8, data_type::u8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, false, data_type::s8, data_type::s8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, false, data_type::s8, data_type::s32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<avx2, false, data_type::s8, data_type::f32>::execute_forward();

template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, true, data_type::u8, data_type::u8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, true, data_type::u8, data_type::s8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, true, data_type::u8, data_type::s32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, true, data_type::u8, data_type::f32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, false, data_type::u8, data_type::u8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, false, data_type::u8, data_type::s8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, false, data_type::u8, data_type::s32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, false, data_type::u8, data_type::f32>::execute_forward();

template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, true, data_type::s8, data_type::u8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, true, data_type::s8, data_type::s8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, true, data_type::s8, data_type::s32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, true, data_type::s8, data_type::f32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, false, data_type::s8, data_type::u8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, false, data_type::s8, data_type::s8>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, false, data_type::s8, data_type::s32>::execute_forward();
template void _jit_uni_x8s8s32x_convolution_fwd_t<sse42, false, data_type::s8, data_type::f32>::execute_forward();

}
}
}
