#include "cl_viv_vx_ext.h"

_viv_uniform float2 half_input0_wh;
_viv_uniform float2 add_float_value;
_viv_uniform int depth;


_viv_uniform float input1_scale;
_viv_uniform VXC_512Bits uniDFPtoFp32_part0_4x4;
_viv_uniform VXC_512Bits uniDFPtoFp32_part1_4x4;
_viv_uniform VXC_512Bits uniConvertI8toI8_2x8;

#define GRID_SAMPLE_I8_PROCESS() \
    fxy0 = fxy0 * half_input0_wh.xyxy + add_float_value.xyxy; \
    fxy1 = fxy1 * half_input0_wh.xyxy + add_float_value.xyxy; \
    float4 in_x        = (float4)(fxy0.xz, fxy1.xz); \
    int4   x_idx       = convert_int4(in_x); \
    float4  in_y        = (float4)(fxy0.yw, fxy1.yw); \
    int4    y_idx        = convert_int4(in_y); \
    int4 coord_in = (int4)(x_idx.x, y_idx.x, 0, 0); \
    int8 input_desc; \
    _viv_asm(COPY, input_desc, input0, sizeof(input_desc)); \
    int baseAddr = input_desc.s0; \
    _viv_asm(MOV, coord_in.w, baseAddr); \
    vxc_char16 src, dst; \
    VXC_OP4(img_load_3d, src, input0, coord_in.xywz, \
    VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
    coord_in.x = x_idx.y; \
    coord_in.y = y_idx.y; \
    VXC_OP4(img_load_3d, src, input0, coord_in.xywz, \
    VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0)); \
    coord_in.x = x_idx.z; \
    coord_in.y = y_idx.z; \
    VXC_OP4(img_load_3d, src, input0, coord_in.xywz, \
    VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0)); \
    coord_in.x = x_idx.w; \
    coord_in.y = y_idx.w; \
    VXC_OP4(img_load_3d, src, input0, coord_in.xywz, \
    VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0)); \
    int8 output_desc; \
    _viv_asm(COPY, output_desc, output, sizeof(output_desc)); \
    baseAddr = (int)coord_out.z * output_desc.s4 + output_desc.s0; \
    _viv_asm(MOV, coord_out.w, baseAddr); \
    int loop = depth - 1; \
    while (coord_in.z < loop) \
    { \
        VXC_DP2x8(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_ToNearestEven, 1), uniConvertI8toI8_2x8); \
        VXC_OP4_NoDest(img_store_3d, output, coord_out.xywz, dst, VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
        coord_out.zw = coord_out.zw + (int2)(1, output_desc.s4); \
        coord_in.zw = coord_in.zw + (int2)(1, input_desc.s4); \
        coord_in.x = x_idx.x; \
        coord_in.y = y_idx.x; \
        VXC_OP4(img_load_3d, src, input0, coord_in.xywz, \
        VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
        coord_in.x = x_idx.y; \
        coord_in.y = y_idx.y; \
        VXC_OP4(img_load_3d, src, input0, coord_in.xywz, \
        VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0)); \
        coord_in.x = x_idx.z; \
        coord_in.y = y_idx.z; \
        VXC_OP4(img_load_3d, src, input0, coord_in.xywz, \
        VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0)); \
        coord_in.x = x_idx.w; \
        coord_in.y = y_idx.w; \
        VXC_OP4(img_load_3d, src, input0, coord_in.xywz, \
        VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0)); \
    } \
    VXC_DP2x8(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_ToNearestEven, 1), uniConvertI8toI8_2x8); \
    VXC_OP4_NoDest(img_store_3d, output, coord_out.xywz, dst, VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));


__kernel void nearest_grid_sample_I8_I8toI8(
    __read_only image2d_array_t   input0,
    __read_only image2d_t   input1,
    __write_only image2d_array_t  output,
    int align_corners)
{
    int4 coord_out = (int4)(get_global_id(0), get_global_id(1), 0, 0);
    int4 coord_in1 = coord_out.xyxy;

    coord_in1.xz = coord_in1.xz * 2;
    vxc_char16 read_coord;
    VXC_ReadImage(read_coord, input1, coord_in1.xy, VXC_5BITOFFSET_XY(0, 0), \
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    float4 fxy0;
    float4 fxy1;

    VXC_DP4x4(fxy0, read_coord, read_coord, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDFPtoFp32_part0_4x4);
    VXC_DP4x4(fxy1, read_coord, read_coord, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDFPtoFp32_part1_4x4);

    fxy0 = fxy0 * input1_scale;
    fxy1 = fxy1 * input1_scale;

    GRID_SAMPLE_I8_PROCESS();

}
