#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniAccumulateSum_2x8;
_viv_uniform int index_num;
_viv_uniform int update_width;
_viv_uniform int output_width;
_viv_uniform int zeropoint;

_viv_uniform int offsetX;
_viv_uniform int offsetY;
_viv_uniform int offsetZ;

inline uchar* get_image2D_array_ptr(image2d_t  input)
{
    int8 desc;
    _viv_asm(COPY, desc, input, sizeof(desc));
    uchar *src_ptr = (uchar*)desc.s0;

    return src_ptr;
}

__kernel void scatter_nd_F16toF16_big(
    __read_only image2d_t   input0,
    __read_only image2d_t   input1,
    image2d_t  output,
    int width,
    int area,
    int coord_dim
    )
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);

    vxc_short8 tmpVal = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    vxc_half8 sum;
    _viv_asm(COPY, sum, tmpVal, 16);
    __global int* index_ptr = (__global int*)get_image2D_array_ptr(input0);
    __global short* update_ptr = (__global short*)get_image2D_array_ptr(input1);
    __global short* output_ptr = (__global short*)get_image2D_array_ptr(output);
    for(int i = 0; i < index_num; i++)
    {
        int4 indice = vload4(0, index_ptr + i * coord_dim);

        int idx = indice.x * offsetX + indice.y * offsetY + indice.z * offsetZ;
        if(gidy == idx)
        {
            vxc_half8 src;
            short tmpData = update_ptr[i * update_width + gidx];
            _viv_asm(COPY, src, tmpData, 4);
            VXC_DP2x8(sum, sum, src, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniAccumulateSum_2x8);
        }
    }
    short dst;
    _viv_asm(COPY, dst, sum, 4);
    output_ptr[gidy * output_width+ gidx] = dst;
}

#define SCATTER_ND_QINT_BIG(src0_type_name, data_type, ptr_type) \
__kernel void scatter_nd_##src0_type_name##to##src0_type_name##_big( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    image2d_t  output, \
    int width, \
    int area, \
    int coord_dim \
    ) \
{ \
    int gidx = get_global_id(0);  \
    int gidy = get_global_id(1); \
    int firstFlg = 1; \
 \
    data_type sum = (data_type)(0, 0, 0, 0, 0, 0, 0, 0); \
    __global int* index_ptr = (__global int*)get_image2D_array_ptr(input0); \
    __global ptr_type* update_ptr = (__global ptr_type*)get_image2D_array_ptr(input1); \
    __global ptr_type* output_ptr = (__global ptr_type*)get_image2D_array_ptr(output); \
    for(int i = 0; i < index_num; i++) \
    { \
        int4 indice = vload4(0, index_ptr + i * coord_dim); \
        int idx = indice.x * offsetX + indice.y * offsetY + indice.z * offsetZ; \
        if(gidy == idx) \
        { \
            data_type src; \
            ptr_type tmpData = update_ptr[i * update_width + gidx]; \
            _viv_asm(COPY, src, tmpData, 4); \
            VXC_DP2x8(sum, sum, src, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniAccumulateSum_2x8); \
            if(firstFlg) \
            { \
                firstFlg = 0; \
            } \
        } \
    } \
    ptr_type dst; \
    if(firstFlg) \
    { \
        _viv_asm(COPY, dst, zeropoint, 4); \
    } \
    else \
    { \
        _viv_asm(COPY, dst, sum, 4); \
    } \
    output_ptr[gidy * output_width+ gidx] = dst; \
}
SCATTER_ND_QINT_BIG(U8,  vxc_uchar8, uchar)
SCATTER_ND_QINT_BIG(I8,  vxc_char8, char)
SCATTER_ND_QINT_BIG(I16, vxc_short8, short)
