#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniConvertInt8toFp16_2x8;

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

_viv_uniform VXC_512Bits uniConvertInt16toFp16_2x8;
_viv_uniform VXC_512Bits uinConvertFp16ToInt8_2x8;

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

_viv_uniform VXC_512Bits uinConvertFp16ToInt16_2x8;

#define GATHER_ND_QINT_TO_F16_1D(src0_type_name, read_type) \
__kernel void gather_nd_##src0_type_name##toF16_1D( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    __write_only image2d_t  output, \
    int block_size, \
    int coord_dim \
    ) \
{ \
    int gidx = get_global_id(0); \
    int gidy = get_global_id(1); \
 \
    int4 coord = (int4)(0, gidy, gidx, 0); \
    Image img = create_image_from_image2d(input1, 4); \
    uchar* indice_ptr = get_image_ptr_from_coord(img, coord.xy); \
    int4 indice = ((int4 *)indice_ptr)[0]; \
 \
    coord.w = indice.x; \
 \
    read_type src; \
    VXC_ReadImage(src, input0, coord.zw, 0, VXC_MODIFIER(0, 0, 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.zy, dst0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
}
GATHER_ND_QINT_TO_F16_1D(U8, vxc_uchar16)
GATHER_ND_QINT_TO_F16_1D(I8, vxc_char16)
GATHER_ND_QINT_TO_F16_1D(I16, vxc_short8)

#define GATHER_ND_F16_TO_QINT_1D(src1_type_name, write_type) \
__kernel void gather_nd_F16to##src1_type_name##_1D( \
    __read_only image2d_t   input0, \
    __read_only image2d_t   input1, \
    __write_only image2d_t  output, \
    int block_size, \
    int coord_dim \
    ) \
{ \
    int gidx = get_global_id(0); \
    int gidy = get_global_id(1); \
 \
    int4 coord = (int4)(0, gidy, gidx, 0); \
    Image img = create_image_from_image2d(input1, 4); \
    uchar* indice_ptr = get_image_ptr_from_coord(img, coord.xy); \
    int4 indice = ((int4 *)indice_ptr)[0]; \
 \
    coord.w = indice.x; \
 \
    vxc_short8 src; \
    VXC_ReadImage(src, input0, coord.zw, 0, VXC_MODIFIER(0, 0, 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.zy, dst, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
}
GATHER_ND_F16_TO_QINT_1D(U8, vxc_uchar16)
GATHER_ND_F16_TO_QINT_1D(I8, vxc_char16)
GATHER_ND_F16_TO_QINT_1D(I16, vxc_short8)

