# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""Core kernel of dot product of 4 Int8 operations"""
#pylint: disable=invalid-name
import tvm
from tvm import te
import tvm.target.codegen


def dot_16x1x16_uint8_int8_int32():
    """Dispatch the most optimized intrin depending on the target"""
    mcpu = tvm.target.Target.current().mcpu

    assert mcpu in ("skylake-avx512", "cascadelake"), \
        "An old Intel machine that does not have fast Int8 support."
    if mcpu == "skylake-avx512":
        return dot_16x1x16_uint8_int8_int32_skylake()
    # cascadelake
    return dot_16x1x16_uint8_int8_int32_cascadelake()


def dot_16x1x16_uint8_int8_int32_skylake():
    """
    Int8 dot product by every 4 elements using AVX512 Skylake instructions.
    This function takes two arrays of uint8 and int8 datatype -- data[4] and
    kernel[16][4] -- and computes a dot product of data[4] with every
    4 elements of kernels, resulting in output[16] of int32 datatype.
    The pseudo code is as follows.
    .. code-block:: c
        void dot_16x1x16_uint8_int8_int32(uint8 data[4], int8 kernel[16][4],
                int32 output[16]){
            for (int i = 0; i < 16; i++){
                output[i] = 0;
                for (int k = 0; k < 4; k++){
                    output[i] += data[k] * kernel[i][k]
                }
            }
        }

    Physically, the kernel array sits in an AVX512 vector register and
    the data[4] is broadcasted to another AVX512 vector register. This
    function returns a TensorIntrin that can be used to tensorize
    a schedule.

    Returns
    -------
    intrin : TensorIntrin
        The Skylake int8 TensorIntrin that can be used in tensorizing schedule
    """

    int32_lanes = 16 # 16 int32 lanes in AVX512
    num_int8_elements = 4 # 4 int8 elements in int32
    data = te.placeholder((num_int8_elements,), dtype='uint8', name='data')
    kernel = te.placeholder((int32_lanes, num_int8_elements), dtype='int8', name='kernel')
    k = te.reduce_axis((0, num_int8_elements), name='k')
    C = te.compute((int32_lanes,),
                   lambda i: te.sum(data[k].astype('int32') *
                                    kernel[i, k].astype('int32'),
                                    axis=k),
                   name="C")

    a_buffer = tvm.tir.decl_buffer(data.shape, dtype='uint8', name="a_buffer",
                                   offset_factor=1,
                                   strides=[1])
    b_buffer = tvm.tir.decl_buffer(kernel.shape, dtype='int8', name="b_buffer",
                                   offset_factor=1,
                                   strides=[te.var('ldw'), 1])

    def _intrin_func(ins, outs):
        def _instr(index):
            ib = tvm.tir.ir_builder.create()
            if index == 1:
                ib.emit(outs[0].vstore(0, tvm.tir.const(0, 'int32x16')))
                return ib.get()

            a_int8 = ins[0].vload([0], "uint8x4")
            re_int32 = tvm.tir.call_pure_intrin('int32', 'reinterpret', a_int8)
            vec_ai32 = re_int32.astype('int32x16')
            vec_a = tvm.tir.call_pure_intrin('int8x64', 'reinterpret', vec_ai32)
            vec_b = ins[1].vload([0, 0], "int8x64")
            vec_one = tvm.tir.const(1, "int16x32")
            pair_reduction = tvm.tir.call_llvm_intrin('int16x32',
                                                      'llvm.x86.avx512.pmaddubs.w.512',
                                                      tvm.tir.const(0, 'uint32'),
                                                      vec_a, vec_b)
            quad_reduction = tvm.tir.call_llvm_intrin('int32x16',
                                                      'llvm.x86.avx512.pmaddw.d.512',
                                                      tvm.tir.const(0, 'uint32'),
                                                      pair_reduction, vec_one)
            if index == 0:
                ib.emit(outs[0].vstore(0, quad_reduction))
            else:
                ib.emit(outs[0].vstore(0, quad_reduction + outs[0].vload([0], 'int32x16')))
            return ib.get()

        # body, reset, update
        return _instr(0), _instr(1), _instr(2)

    with tvm.target.build_config(offset_factor=1, partition_const_loop=True):
        return te.decl_tensor_intrin(C.op, _intrin_func, binds={data:a_buffer, kernel:b_buffer})


def dot_16x1x16_uint8_int8_int16():
    """
    Int8 dot product by every 2 elements using AVX512 Skylake instructions.
    This function takes two arrays of uint8 and int8 datatype -- data[2] and
    kernel[4][32][2] -- and computes a dot product of data[2] with every
    2 elements of kernels, resulting in output[4][32] of int16 datatype.
    The pseudo code is as follows.
    .. code-block:: c
        void dot_16x1x16_uint8_int8_int16(uint8 data[2], int8 kernel[32*4][2],
                int16 output[32*4]){
            for (int i = 0; i< 4; i++){
                for (int j = 0; j < 32; j++){
                    output[i][i] = 0;
                    for (int k = 0; k < 2; k++){
                        output[i][j][k] += data[k] * kernel[i][j][k]
                    }
                }
            }
        }

    Physically, the kernel array sits in four AVX512 vector registers and
    the data[2] is broadcasted to another AVX512 vector register. This
    function returns a TensorIntrin that can be used to tensorize
    a schedule.

    Returns
    -------
    intrin : TensorIntrin
        The Skylake int8 TensorIntrin that can be used in tensorizing schedule
    """

    int16_lanes = 4*32 # 4*32 int32 lanes in 4 AVX512 vector registers
    num_int8_elements = 2 # 2 int8 elements in int16
    data = te.placeholder((num_int8_elements,), dtype='uint8', name='data')
    kernel = te.placeholder((int16_lanes, num_int8_elements), dtype='int8', name='kernel')
    k = te.reduce_axis((0, num_int8_elements), name='k')
    C = te.compute((int16_lanes, ),
                   lambda i: te.sum(data[k].astype('int16') *
                                    kernel[i, k].astype('int16'),
                                    axis=k),
                   name="C")

    a_buffer = tvm.tir.decl_buffer(data.shape, dtype='uint8', name="a_buffer",
                                   offset_factor=1,
                                   strides=[1])
    b_buffer = tvm.tir.decl_buffer(kernel.shape, dtype='int8', name="b_buffer",
                                   offset_factor=1)
    # strides=[te.var('ldw'), 1, 1])

    def _intrin_func(ins, outs):
        def _instr(index):
            ib = tvm.tir.ir_builder.create()
            if index == 1:
                for i in range(4):
                    ib.emit(outs[0].vstore([i*32], tvm.tir.const(0, 'int16x32')))
                return ib.get()

            a_int8 = ins[0].vload([0], "uint8x2")
            re_int16 = tvm.tir.call_pure_intrin('int16', 'reinterpret', a_int8)
            vec_ai16 = re_int16.astype('int16x32')
            vec_a = tvm.tir.call_pure_intrin('int8x64', 'reinterpret', vec_ai16)

            for i in range(4):
                vec_b = ins[1].vload([i*32, 0], "int8x64")
                pair_reduction = tvm.tir.call_llvm_intrin('int16x32',
                                                          'llvm.x86.avx512.pmaddubs.w.512',
                                                          tvm.tir.const(0, 'uint32'),
                                                          vec_a, vec_b)
                if index == 0:
                    ib.emit(outs[0].vstore([i*32], pair_reduction))
                else:
                    ib.emit(outs[0].vstore([i*32], pair_reduction + outs[0].vload([i*32],
                                                                                  'int16x32')))
            return ib.get()

        # body, reset, update
        return _instr(0), _instr(1), _instr(2)

    with tvm.target.build_config(offset_factor=1, partition_const_loop=True):
        return te.decl_tensor_intrin(C.op, _intrin_func, binds={data:a_buffer, kernel:b_buffer})


def dot_16x1x16_uint8_int8_int32_cascadelake():
    """
    Int8 dot product by every 4 elements using AVX512VNNI Cascade Lake instructions.
    This function takes two arrays of uint8 and int8 datatype -- data[4] and
    kernel[16][4] -- and computes a dot product of data[4] with every
    4 elements of kernels, resulting in output[16] of int32 datatype.
    The pseudo code is as follows.
    .. code-block:: c
        void dot_16x1x16_uint8_int8_int32_cascadelake(uint8 data[4], int8 kernel[16][4],
                int32 output[16]){
            for (int i = 0; i < 16; i++){
                output[i] = 0;
                for (int k = 0; k < 4; k++){
                    output[i] += data[k] * kernel[i][k]
                }
            }
        }

    Physically, the kernel array sits in an AVX512 vector register and
    the data[4] is broadcasted to another AVX512 vector register. This
    function returns a TensorIntrin that can be used to tensorize
    a schedule.

    Returns
    -------
    intrin : TensorIntrin
        The Cascade Lake int8 TensorIntrin that can be used in tensorizing schedule
    """

    int32_lanes = 16 # 16 int32 lanes in AVX512
    num_int8_elements = 4 # 4 int8 elements in int32
    data = te.placeholder((num_int8_elements,), dtype='uint8', name='data')
    kernel = te.placeholder((int32_lanes, num_int8_elements), dtype='int8', name='kernel')
    k = te.reduce_axis((0, num_int8_elements), name='k')
    C = te.compute((int32_lanes,),
                   lambda i: te.sum(data[k].astype('int32') *
                                    kernel[i, k].astype('int32'),
                                    axis=k),
                   name="C")

    a_buffer = tvm.tir.decl_buffer(data.shape, dtype='uint8', name="a_buffer",
                                   offset_factor=1,
                                   strides=[1])
    b_buffer = tvm.tir.decl_buffer(kernel.shape, dtype='int8', name="b_buffer",
                                   offset_factor=1,
                                   strides=[te.var('ldw'), 1])

    def _intrin_func(ins, outs):
        def _instr(index):
            ib = tvm.tir.ir_builder.create()
            if index == 1:
                ib.emit(outs[0].vstore(0, tvm.tir.const(0, 'int32x16')))
                return ib.get()

            a_int8 = ins[0].vload([0], "uint8x4")
            re_int32 = tvm.tir.call_pure_intrin('int32', 'reinterpret', a_int8)
            vec_ai32 = re_int32.astype('int32x16')
            vec_b = ins[1].vload([0, 0], "int8x64")

            vnni_inst_name = 'llvm.x86.avx512.vpdpbusd.512'
            llvm_id = tvm.target.codegen.llvm_lookup_intrinsic_id(vnni_inst_name)

            if llvm_id != 0: # VNNI is available for current LLVM version
                vec_bi32 = tvm.tir.call_pure_intrin('int32x16', 'reinterpret', vec_b)
                vec_zero = tvm.tir.const(0, "int32x16")
                quad_reduction = tvm.tir.call_llvm_intrin('int32x16',
                                                          'llvm.x86.avx512.vpdpbusd.512',
                                                          tvm.tir.const(0, 'uint32'),
                                                          vec_zero,
                                                          vec_ai32, vec_bi32)
            else: # Fall back to the normal AVX512
                vec_a = tvm.tir.call_pure_intrin('int8x64', 'reinterpret', vec_ai32)
                vec_one = tvm.tir.const(1, "int16x32")
                pair_reduction = tvm.tir.call_llvm_intrin('int16x32',
                                                          'llvm.x86.avx512.pmaddubs.w.512',
                                                          tvm.tir.const(0, 'uint32'),
                                                          vec_a, vec_b)
                quad_reduction = tvm.tir.call_llvm_intrin('int32x16',
                                                          'llvm.x86.avx512.pmaddw.d.512',
                                                          tvm.tir.const(0, 'uint32'),
                                                          pair_reduction, vec_one)

            if index == 0:
                ib.emit(outs[0].vstore(0, quad_reduction))
            else:
                ib.emit(outs[0].vstore(0, quad_reduction + outs[0].vload([0], 'int32x16')))
            return ib.get()

        # body, reset, update
        return _instr(0), _instr(1), _instr(2)

    with tvm.target.build_config(offset_factor=1, partition_const_loop=True):
        return te.decl_tensor_intrin(C.op, _intrin_func, binds={data:a_buffer, kernel:b_buffer})
