#include "cl_viv_vx_ext.h"

__kernel void gather_nd_I8toI8_3D(
    __read_only image2d_array_t   input0,
    __read_only image2d_t   input1,
    __write_only image2d_t  output,
    int block_size,
    int coord_dim
    )
{
    int gidx = get_global_id(0);  // block_size
    int gidy = get_global_id(1);  // indices_num

    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];

    indice.x = indice.x * block_size + gidx;
    indice.w = 0;

    vxc_char16 src;
    VXC_ReadImage2DArray(src, input0, indice, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));

    VXC_WriteImage(output, coord.zy, src, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
}

__kernel void gather_nd_U8toU8_3D(
    __read_only image2d_array_t   input0,
    __read_only image2d_t   input1,
    __write_only image2d_t  output,
    int block_size,
    int coord_dim
    )
{
    int gidx = get_global_id(0);  // block_size
    int gidy = get_global_id(1);  // indices_num

    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];

    indice.x = indice.x * block_size + gidx;
    indice.w = 0;

    vxc_uchar16 src;
    VXC_ReadImage2DArray(src, input0, indice, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    VXC_WriteImage(output, coord.zy, src, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
}

__kernel void gather_nd_I16toI16_3D(
    __read_only image2d_array_t   input0,
    __read_only image2d_t   input1,
    __write_only image2d_t  output,
    int block_size,
    int coord_dim
    )
{
    int gidx = get_global_id(0);  // block_size
    int gidy = get_global_id(1);  // indices_num

    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];

    indice.x = indice.x * block_size + gidx;
    indice.w = 0;

    vxc_short8 src;
    VXC_ReadImage2DArray(src, input0, indice, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    VXC_WriteImage(output, coord.zy, src, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
}

__kernel void gather_nd_F16toF16_3D(
    __read_only image2d_array_t   input0,
    __read_only image2d_t   input1,
    __write_only image2d_t  output,
    int block_size,
    int coord_dim
    )
{
    int gidx = get_global_id(0);  // block_size
    int gidy = get_global_id(1);  // indices_num

    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];

    indice.x = indice.x * block_size + gidx;
    indice.w = 0;

    vxc_short8 src;
    VXC_ReadImage2DArray(src, input0, indice, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    VXC_WriteImage(output, coord.zy, src, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
}

