#include "cl_viv_vx_ext.h"

_viv_uniform int indices_num;

_viv_uniform int2 multAndoutZP0;//[0:15] multiplier, [31:63] output zp
_viv_uniform VXC_512Bits uniU8MulAndPostShift_0_Lo_2x8;

_viv_uniform VXC_512Bits uniConvertFp16toU8_2x8;
_viv_uniform int2 multAndoutZP1;//[0:15] multiplier, [31:63] output zp

#define GATHER_8BITS_TO_F16(src0_type_name, read_type) \
__kernel void gather_##src0_type_name##toF16( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    __write_only image2d_t  output, \
    int block_size, \
    int block_num, \
    int axis_num, \
    int is_array \
    ) \
{ \
    int gidx = get_global_id(0); \
    int gidy = get_global_id(1); \
    int gidz = get_global_id(2); \
 \
    int4 coord_in = (int4)(gidy, 0, gidx, 0); \
    int4 indice = read_imagei(input1, coord_in.xy); \
    indice.x = indice.x >= 0 ? indice.x : indice.x + axis_num; \
    coord_in.w = gidz * axis_num + indice.x; \
 \
    read_type src; \
    VXC_ReadImage(src, input0, coord_in.zw, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
    int2 coord = (int2)(gidx, gidz * indices_num + gidy); \
    vxc_half8  src0, src1; \
    vxc_short8 dst0, dst1; \
    vxc_ushort8 ms0; \
    _viv_asm(COPY, ms0, multAndoutZP0, 16); \
    VXC_DP2x8(src0,src,ms0, VXC_MODIFIER(0,7,0, VXC_RM_TowardZero,1), uniU8MulAndPostShift_0_Lo_2x8); \
    VXC_DP2x8(src1,src,ms0, VXC_MODIFIER(0,7,8, VXC_RM_TowardZero,1), uniU8MulAndPostShift_0_Lo_2x8); \
    _viv_asm(COPY, dst0, src0, 16); \
    _viv_asm(COPY, dst1, src1, 16); \
    VXC_WriteImage(output, coord, dst0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    coord.x += 8; \
    VXC_WriteImage(output, coord, dst1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}
GATHER_8BITS_TO_F16(U8, vxc_uchar16)
GATHER_8BITS_TO_F16(I8, vxc_char16)

#define GATHER_F16_TO_QINT(src1_type_name, write_type) \
__kernel void gather_F16to##src1_type_name( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    __write_only image2d_t  output, \
    int block_size, \
    int block_num, \
    int axis_num, \
    int is_array \
    ) \
{ \
    int gidx = get_global_id(0); \
    int gidy = get_global_id(1); \
    int gidz = get_global_id(2); \
    int4 coord_in = (int4)(gidy, 0, gidx, 0); \
 \
    int4 indice = read_imagei(input1, coord_in.xy); \
    indice.x = indice.x >= 0 ? indice.x : indice.x + axis_num; \
    coord_in.w = gidz * axis_num + indice.x; \
 \
    vxc_short8 src; \
    VXC_ReadImage(src, input0, coord_in.zw, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    int2 coord = (int2)(gidx, gidz * indices_num + gidy); \
    vxc_ushort8 mp1; \
    _viv_asm(COPY, mp1, multAndoutZP1, 16); \
    vxc_half8 data; \
    write_type dst; \
    _viv_asm(COPY, data, src, 16); \
    VXC_DP2x8(dst,data,mp1,VXC_MODIFIER(0,7,0,VXC_RM_ToNearestEven, 1),uniConvertFp16toU8_2x8); \
    VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}
GATHER_F16_TO_QINT(U8, vxc_uchar16)
GATHER_F16_TO_QINT(I8, vxc_char16)
GATHER_F16_TO_QINT(I16, vxc_short8)

__kernel void gather_I16toF16(
    __read_only image2d_t   input0,
    __read_only image2d_t   input1,
    __write_only image2d_t  output,
    int block_size,
    int block_num,
    int axis_num,
    int is_array
    )
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int gidz = get_global_id(2);

    int4 coord_in = (int4)(gidy, 0, gidx, 0);
    int4 indice = read_imagei(input1, coord_in.xy);
    indice.x = indice.x >= 0 ? indice.x : indice.x + axis_num;
    coord_in.w = gidz * axis_num + indice.x;

    vxc_short8 src;
    VXC_ReadImage(src, input0, coord_in.zw, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    int2 coord = (int2)(gidx, gidz * indices_num + gidy);

    vxc_half8  src0;
    vxc_short8 dst0;
    vxc_ushort8 ms0;
    _viv_asm(COPY, ms0, multAndoutZP0, 16);
    VXC_DP2x8(src0, src, ms0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
                uniU8MulAndPostShift_0_Lo_2x8);
    _viv_asm(COPY, dst0, src0, 16);

    VXC_WriteImage(output, coord, dst0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
}

#define GATHER_8BITS_TO_F16_AXIS0(src0_type_name, read_type) \
__kernel void gather_##src0_type_name##toF16_axis0( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    __write_only image2d_t  output, \
    int block_size, \
    int block_num, \
    int axis_num, \
    int is_array \
    ) \
{ \
    int2 coord = (int2)(get_global_id(0), get_global_id(1)); \
    int4 indices = read_imagei(input1, coord.xx); \
    indices = indices >= 0 ? indices : indices + axis_num; \
    int2 coord_in = (int2)(indices.x, get_global_id(1)); \
 \
    read_type src; \
    VXC_ReadImage(src, input0, coord_in, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
    indices.x = get_global_id(1); \
    VXC_ReadImage(src, input0, indices.yx, 0, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0)); \
    VXC_ReadImage(src, input0, indices.zx, 0, VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0)); \
    VXC_ReadImage(src, input0, indices.wx, 0, VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0)); \
    vxc_half8  src0; \
    vxc_short8 dst0; \
    vxc_ushort8 ms0; \
    _viv_asm(COPY, ms0, multAndoutZP0, 16); \
    VXC_DP2x8(src0,src,ms0, VXC_MODIFIER(0,7,0, VXC_RM_TowardZero,1), uniU8MulAndPostShift_0_Lo_2x8); \
    _viv_asm(COPY, dst0, src0, 16); \
    VXC_WriteImage(output, coord, dst0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
}
GATHER_8BITS_TO_F16_AXIS0(U8, vxc_uchar16)
GATHER_8BITS_TO_F16_AXIS0(I8, vxc_char16)

#define GATHER_F16_TO_QINT_AXIS0(src1_type_name, write_type) \
__kernel void gather_F16to##src1_type_name##_axis0( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    __write_only image2d_t  output, \
    int block_size, \
    int block_num, \
    int axis_num, \
    int is_array \
    ) \
{ \
    int2 coord = (int2)(get_global_id(0), get_global_id(1)); \
    int4 indices = read_imagei(input1, coord.xx); \
    indices = indices >= 0 ? indices : indices + axis_num; \
    int2 coord_in = (int2)(indices.x, get_global_id(1)); \
 \
    vxc_short8 src; \
    VXC_ReadImage(src, input0, coord_in, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
    indices.x = get_global_id(1); \
    VXC_ReadImage(src, input0, indices.yx, 0, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0)); \
    VXC_ReadImage(src, input0, indices.zx, 0, VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0)); \
    VXC_ReadImage(src, input0, indices.wx, 0, VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0)); \
    vxc_ushort8 mp1; \
    _viv_asm(COPY, mp1, multAndoutZP1, 16); \
    vxc_half8 data; \
    write_type dst; \
    _viv_asm(COPY, data, src, 16); \
    VXC_DP2x8(dst,data,mp1,VXC_MODIFIER(0,7,0,VXC_RM_ToNearestEven, 1),uniConvertFp16toU8_2x8); \
    VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
}
GATHER_F16_TO_QINT_AXIS0(U8, vxc_uchar16)
GATHER_F16_TO_QINT_AXIS0(I8, vxc_char16)
GATHER_F16_TO_QINT_AXIS0(I16, vxc_short8)

__kernel void gather_I16toF16_axis0(
    __read_only image2d_t   input0,
    __read_only image2d_t   input1,
    __write_only image2d_t  output,
    int block_size,
    int block_num,
    int axis_num,
    int is_array
    )
{
    int2 coord = (int2)(get_global_id(0), get_global_id(1));
    int4 indices = read_imagei(input1, coord.xx);
    indices = indices >= 0 ? indices : indices + axis_num;
    int2 coord_in = (int2)(indices.x, get_global_id(1));

    vxc_short8 src;
    VXC_ReadImage(src, input0, coord_in, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    indices.x = get_global_id(1);
    VXC_ReadImage(src, input0, indices.yx, 0, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0));
    VXC_ReadImage(src, input0, indices.zx, 0, VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0));
    VXC_ReadImage(src, input0, indices.wx, 0, VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0));

    vxc_half8  src0;
    vxc_short8 dst0;
    vxc_ushort8 ms0;
    _viv_asm(COPY, ms0, multAndoutZP0, 16);
    VXC_DP2x8(src0, src, ms0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
                uniU8MulAndPostShift_0_Lo_2x8);
    _viv_asm(COPY, dst0, src0, 16);

    VXC_WriteImage(output, coord, dst0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
}
