#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniAccumulateSum_2x8;
_viv_uniform VXC_512Bits uniU8MulAndPostShift_0_Lo_2x8;
_viv_uniform VXC_512Bits uniU8MulAndPostShift_1_Lo_2x8;
_viv_uniform int index_num;
_viv_uniform int offset_idx;
_viv_uniform int offsetX;
_viv_uniform int offsetY;
_viv_uniform int offsetZ;
_viv_uniform int offsetW;
_viv_uniform int2 multAndoutZP0;
_viv_uniform int2 multAndoutZP1;//[0:15] multiplier, [31:63] output zp

_viv_uniform VXC_512Bits uniConvBF16toF32_Part0_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_Part1_2x8;
_viv_uniform VXC_512Bits uniExtractOddData_2x8;

__kernel void scatter_nd_update_F16F16toF16(
    __read_only image2d_t   input0,
    __read_only image2d_t   input1,
    __read_only image2d_t   input2,
    image2d_array_t  output,
    int width,
    int area,
    int vol,
    int coord_dim
    )
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int cnt = 0;

    vxc_short8 tmpVal = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    vxc_half8 sum;
    _viv_asm(COPY, sum, tmpVal, 16);
    Image img1 = create_image_from_image2d(input1, 4);
    __global int* index_ptr = (__global int*)img1.ptr;
    for(int i = 0; i < index_num; i++)
    {
        //int4 indice = read_imagei(input1, (int2)(0, i));
        int4 indice = vload4(0, index_ptr + offset_idx);
        index_ptr += coord_dim;
        int idx = indice.x * offsetX + indice.y * offsetY + indice.z * offsetZ + indice.w * offsetW;
        if(gidy == idx)
        {
            vxc_half8 src;
            VXC_ReadImage(tmpVal, input2, (int2)(gidx, i), 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            cnt++;
            _viv_asm(COPY, src, tmpVal, 16);
            VXC_DP2x8(sum, sum, src, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniAccumulateSum_2x8);
        }
    }
    _viv_asm(COPY, tmpVal, sum, 16);
    int2 coord = (int2)(gidx, gidy);
    if(cnt == 0)
    {
        VXC_ReadImage(tmpVal, input0, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
    VXC_WriteImage(output, coord, tmpVal, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
}

__kernel void scatter_nd_update_F16F16toU8(
    __read_only image2d_t   input0,
    __read_only image2d_t   input1,
    __read_only image2d_t   input2,
    image2d_array_t  output,
    int width,
    int area,
    int vol,
    int coord_dim
    )
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int cnt = 0;

    vxc_short8 tmpVal = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    vxc_half8 sum;
    _viv_asm(COPY, sum, tmpVal, 16);
    Image img1 = create_image_from_image2d(input1, 4);
    __global int* index_ptr = (__global int*)img1.ptr;
    for(int i = 0; i < index_num; i++)
    {
        //int4 indice = read_imagei(input1, (int2)(0, i));
        int4 indice = vload4(0, index_ptr + offset_idx);
        index_ptr += coord_dim;
        int idx = indice.x * offsetX + indice.y * offsetY + indice.z * offsetZ + indice.w * offsetW;
        if(gidy == idx)
        {
            vxc_half8 src;
            VXC_ReadImage(tmpVal, input2, (int2)(gidx, i), 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            cnt++;
            _viv_asm(COPY, src, tmpVal, 16);
            VXC_DP2x8(sum, sum, src, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniAccumulateSum_2x8);
        }
    }
    int2 coord = (int2)(gidx, gidy);
    vxc_ushort8 ms0;
    vxc_uchar8 dst;
    if(cnt == 0)
    {
        vxc_half8 src;
        VXC_ReadImage(tmpVal, input0, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, ms0, multAndoutZP0, 16);
        _viv_asm(COPY, src, tmpVal, 16);
        VXC_DP2x8(dst, src, ms0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),
                    uniU8MulAndPostShift_0_Lo_2x8);
        VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
    else
    {
        _viv_asm(COPY, ms0, multAndoutZP1, 16);
        VXC_DP2x8(dst, sum, ms0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),
                    uniU8MulAndPostShift_1_Lo_2x8);
        VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
}

#define SCATTER_ND_UPDATE_QINT(src0_type_name, src2_type_name, out_type_name, data_type) \
__kernel void scatter_nd_update_##src0_type_name##src2_type_name##to##out_type_name##( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    __read_only image2d_t   input2, \
    image2d_array_t  output, \
    int width, \
    int area, \
    int vol, \
    int coord_dim \
    ) \
{ \
    int gidx = get_global_id(0);  \
    int gidy = get_global_id(1); \
    int cnt = 0; \
 \
    data_type sum = (data_type)(0, 0, 0, 0, 0, 0, 0, 0); \
    Image img1 = create_image_from_image2d(input1, 4); \
    __global int* index_ptr = (__global int*)img1.ptr; \
    for(int i = 0; i < index_num; i++) \
    { \
        int4 indice = vload4(0, index_ptr + offset_idx); \
        index_ptr += coord_dim; \
        int idx = indice.x * offsetX + indice.y * offsetY + indice.z * offsetZ + indice.w * offsetW; \
        if(gidy == idx) \
        { \
            data_type src; \
            VXC_ReadImage(src, input2, (int2)(gidx, i), 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
            cnt++; \
            VXC_DP2x8(sum, sum, src, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniAccumulateSum_2x8); \
        } \
    } \
    int2 coord = (int2)(gidx, gidy); \
    vxc_ushort8 ms0; \
    data_type dst; \
    if(cnt == 0) \
    { \
        VXC_ReadImage(sum, input0, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        _viv_asm(COPY, ms0, multAndoutZP0, 16); \
        VXC_DP2x8(dst, sum, ms0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
                    uniU8MulAndPostShift_0_Lo_2x8); \
        VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    } \
    else \
    { \
        _viv_asm(COPY, ms0, multAndoutZP1, 16); \
        VXC_DP2x8(dst, sum, ms0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
                    uniU8MulAndPostShift_1_Lo_2x8); \
        VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    } \
}
SCATTER_ND_UPDATE_QINT(U8,  U8,  U8,  vxc_uchar8)
SCATTER_ND_UPDATE_QINT(I8,  I8,  I8,  vxc_char8)
SCATTER_ND_UPDATE_QINT(I16, I16, I16, vxc_short8)

#define SCATTER_ND_UPDATE_QINT_TO_F16(src0_type, data_type) \
__kernel void scatter_nd_update_##src0_type##src0_type##toF16( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    __read_only image2d_t   input2, \
    image2d_array_t  output, \
    int width, \
    int area, \
    int vol, \
    int coord_dim \
    ) \
{ \
    int gidx = get_global_id(0); \
    int gidy = get_global_id(1); \
    int cnt = 0; \
    vxc_short8 sum = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0); \
    data_type src; \
    Image img1 = create_image_from_image2d(input1, 4); \
    __global int* index_ptr = (__global int*)img1.ptr; \
    for(int i = 0; i < index_num; i++) \
    { \
        int4 indice = vload4(0, index_ptr + offset_idx); \
        index_ptr += coord_dim; \
        int idx = indice.x * offsetX + indice.y * offsetY + indice.z * offsetZ + indice.w * offsetW; \
        if(gidy == idx) \
        { \
            VXC_ReadImage(src, input2, (int2)(gidx, i), 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
            cnt++; \
            VXC_DP2x8(sum, sum, src, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniAccumulateSum_2x8); \
        } \
    } \
    int2 coord = (int2)(gidx, gidy); \
    vxc_ushort8 ms0; \
    vxc_half8 tmpDst; \
    vxc_short8 dst; \
    if(cnt == 0) \
    { \
        VXC_ReadImage(src, input0, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        _viv_asm(COPY, ms0, multAndoutZP0, 16); \
        VXC_DP2x8(tmpDst, src, ms0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
                    uniU8MulAndPostShift_0_Lo_2x8); \
        _viv_asm(COPY, dst, tmpDst, 16); \
        VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    } \
    else \
    { \
        _viv_asm(COPY, ms0, multAndoutZP1, 16); \
        VXC_DP2x8(tmpDst, sum, ms0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
                    uniU8MulAndPostShift_1_Lo_2x8); \
        _viv_asm(COPY, dst, tmpDst, 16); \
        VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    } \
}
SCATTER_ND_UPDATE_QINT_TO_F16(U8,  vxc_uchar8)
SCATTER_ND_UPDATE_QINT_TO_F16(I8,  vxc_char8)
SCATTER_ND_UPDATE_QINT_TO_F16(I16, vxc_short8)

__kernel void scatter_nd_update_BF16BF16toBF16(
    __read_only image2d_t   input0,
    __read_only image2d_t   input1,
    __read_only image2d_t   input2,
    image2d_array_t  output,
    int width,
    int area,
    int vol,
    int coord_dim
    )
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int cnt = 0;

    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    vxc_ushort8 src0, src1, src2;
    float4 srcA, srcB;
    float4 sum0 = (float4)(0);
    float4 sum1 = sum0;

    Image img1 = create_image_from_image2d(input1, 4);
    __global int* index_ptr = (__global int*)img1.ptr;
    for(int i = 0; i < index_num; i++)
    {
        //int4 indice = read_imagei(input1, (int2)(0, i));
        int4 indice = vload4(0, index_ptr + offset_idx);
        index_ptr += coord_dim;
        int idx = indice.x * offsetX + indice.y * offsetY + indice.z * offsetZ + indice.w * offsetW;
        if(gidy == idx)
        {
            VXC_ReadImage(src0, input2, (int2)(gidx, i), 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            cnt++;
            VXC_DP2x8(src1, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                        uniConvBF16toF32_Part0_2x8);
            VXC_DP2x8(src2, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                         uniConvBF16toF32_Part1_2x8);
            _viv_asm(COPY, srcA, src1, 16);
            _viv_asm(COPY, srcB, src2, 16);
            sum0 += srcA;
            sum1 += srcB;
        }
    }
    int2 coord = (int2)(gidx, gidy);
    if(cnt == 0)
    {
        VXC_ReadImage(src2, input0, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_WriteImage(output, coord, src2, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
    else
    {
        _viv_asm(COPY, src0, sum0, 16);
        _viv_asm(COPY, src1, sum1, 16);
        VXC_DP2x8(src2, src0, src1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniExtractOddData_2x8);
        VXC_WriteImage(output, coord, src2, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
}