#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniExtact8Bit_2x8;
_viv_uniform VXC_512Bits uniFp16toFp32_left_4x4;
_viv_uniform VXC_512Bits uniRightSubLeft_4x4;
_viv_uniform VXC_512Bits uniExtactHalf8_2x8;
_viv_uniform float2 scale_xy;
_viv_uniform int depth;
_viv_uniform float uint8Scale;
_viv_uniform float output_ZP;
_viv_uniform VXC_512Bits uniFp16toFp32_part1_4x4;
_viv_uniform VXC_512Bits uniConvertI32toI16_2x8;
_viv_uniform VXC_512Bits uniGetMaskShift_2x8;
_viv_uniform float half_pixel_value;

__kernel void resize_bilinear_F16toF16_DOWN
    (
    __read_only     image2d_array_t input,
    __write_only    image2d_array_t output,
    int   align_corners,
    int   half_pixel_centers
    )
{
    int4 coord_out  =  (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0);
    int4   coord_x     = coord_out.xxxx + (int4)(0, 1, 2, 3);
    float4 in_x        = (convert_float4(coord_x) + half_pixel_value) * scale_xy.xxxx - half_pixel_value;
    float4 left_x_f    = floor(in_x);
    float4 x_lerp      = in_x - left_x_f;
    int4   left_x_idx  = convert_int4(left_x_f);
    float  in_y        = (convert_float(coord_out.y) + half_pixel_value) * scale_xy.y - half_pixel_value;
    float  top_y_f     = floor(in_y);
    float  y_lerp      = in_y - top_y_f;
    int    top_y_idx   = convert_int(top_y_f);
    vxc_short8 top_short, bottom_short, dst;
    vxc_half8  top, bottom, result;
    int4 coord_in = (int4)(left_x_idx.x, top_y_idx, coord_out.z, 0);

    int8 input_desc;
    _viv_asm(COPY, input_desc, input, sizeof(input_desc));
    int baseAddr = (int)coord_in.z * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord_in.w, baseAddr);
    VXC_OP4(img_load_3d, top_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0));
    VXC_OP4(img_load_3d, bottom_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1),
            VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0));
    coord_in.x = left_x_idx.y;
    VXC_OP4(img_load_3d, top_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0));
    VXC_OP4(img_load_3d, bottom_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1),
            VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0));
    coord_in.x = left_x_idx.z;
    VXC_OP4(img_load_3d, top_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(4, 5, 0, VXC_RM_TowardZero, 0));
    VXC_OP4(img_load_3d, bottom_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1),
            VXC_MODIFIER(4, 5, 0, VXC_RM_TowardZero, 0));
    coord_in.x = left_x_idx.w;
    VXC_OP4(img_load_3d, top_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(6, 7, 0, VXC_RM_TowardZero, 0));
    VXC_OP4(img_load_3d, bottom_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1),
            VXC_MODIFIER(6, 7, 0, VXC_RM_TowardZero, 0));
    _viv_asm(COPY, top,    top_short, 16);
    _viv_asm(COPY, bottom, bottom_short, 16);

    float4 left4;
    float4 right4;
    float4 top4;
    float4 bottom4;

    VXC_DP4x4(left4, top, top, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniFp16toFp32_left_4x4);
    VXC_DP4x4(right4, top, top, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniRightSubLeft_4x4);
    top4       = right4 * x_lerp + left4;
    VXC_DP4x4(left4, bottom, bottom, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniFp16toFp32_left_4x4);
    VXC_DP4x4(right4, bottom, bottom, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniRightSubLeft_4x4);
    bottom4    = right4 * x_lerp + left4;
    bottom4   -= top4;
    float4 dst4       = bottom4 * y_lerp + top4;

    half4 tmp;
    _viv_asm(CONV, tmp, dst4);
    VXC_DP2x8(result, tmp, tmp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniExtactHalf8_2x8);
    _viv_asm(COPY, dst, result, 16);

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

    VXC_OP4_NoDest(img_store_3d, output, coord_out.xywz, dst,
        VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));
}

#define RESIZE_BILINEAR_F16TOQINT_DOWN(out_name, dst_type) \
__kernel void resize_bilinear_F16to##out_name##_DOWN( \
    __read_only     image2d_array_t input, \
    __write_only    image2d_array_t output, \
    int   align_corners, \
    int   half_pixel_centers \
    ) \
{ \
    int4 coord_out  =  (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0); \
    int4   coord_x     = coord_out.xxxx + (int4)(0, 1, 2, 3); \
    float4 in_x        = (convert_float4(coord_x) + half_pixel_value) * scale_xy.xxxx - half_pixel_value; \
    float4 left_x_f    = floor(in_x); \
    float4 x_lerp      = in_x - left_x_f; \
    int4   left_x_idx  = convert_int4(left_x_f); \
    float  in_y        = (convert_float(coord_out.y) + half_pixel_value) * scale_xy.y - half_pixel_value; \
    float  top_y_f     = floor(in_y); \
    float  y_lerp      = in_y - top_y_f; \
    int    top_y_idx   = convert_int(top_y_f); \
 \
    vxc_short8 top_short, bottom_short; \
    vxc_half8  top, bottom; \
    int4 coord_in = (int4)(left_x_idx.x, top_y_idx, coord_out.z, 0); \
 \
    int8 input_desc; \
    _viv_asm(COPY, input_desc, input, sizeof(input_desc)); \
    int baseAddr = (int)coord_in.z * input_desc.s4 + input_desc.s0; \
    _viv_asm(MOV, coord_in.w, baseAddr); \
    VXC_OP4(img_load_3d, top_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0), \
            VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0)); \
    VXC_OP4(img_load_3d, bottom_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1), \
            VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0)); \
    coord_in.x = left_x_idx.y; \
    VXC_OP4(img_load_3d, top_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0), \
            VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0)); \
    VXC_OP4(img_load_3d, bottom_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1), \
            VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0)); \
    coord_in.x = left_x_idx.z; \
    VXC_OP4(img_load_3d, top_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0), \
            VXC_MODIFIER(4, 5, 0, VXC_RM_TowardZero, 0)); \
    VXC_OP4(img_load_3d, bottom_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1), \
            VXC_MODIFIER(4, 5, 0, VXC_RM_TowardZero, 0)); \
    coord_in.x = left_x_idx.w; \
    VXC_OP4(img_load_3d, top_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0), \
            VXC_MODIFIER(6, 7, 0, VXC_RM_TowardZero, 0)); \
    VXC_OP4(img_load_3d, bottom_short, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1), \
            VXC_MODIFIER(6, 7, 0, VXC_RM_TowardZero, 0)); \
    _viv_asm(COPY, top,    top_short, 16); \
    _viv_asm(COPY, bottom, bottom_short, 16); \
 \
    float4 left4, right4, top4, bottom4; \
    VXC_DP4x4(left4, top, top, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniFp16toFp32_left_4x4); \
    VXC_DP4x4(right4, top, top, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniRightSubLeft_4x4); \
    top4        = right4 * x_lerp + left4; \
    VXC_DP4x4(left4, bottom, bottom, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniFp16toFp32_left_4x4); \
    VXC_DP4x4(right4, bottom, bottom, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniRightSubLeft_4x4); \
    bottom4      = right4 * x_lerp + left4; \
    bottom4     -= top4; \
    float4 dst4  = bottom4 * y_lerp + top4; \
    dst4         = dst4 * uint8Scale + output_ZP; \
    int4 dst     = convert_int4_rte(dst4); \
    dst_type dst_uchar; \
    VXC_DP2x8(dst_uchar, dst, dst, VXC_MODIFIER(0, 3, 0, VXC_RM_ToNearestEven, 1), uniExtact8Bit_2x8); \
 \
    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); \
 \
    VXC_OP4_NoDest(img_store_3d, output, coord_out.xywz, dst_uchar, \
        VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
}
RESIZE_BILINEAR_F16TOQINT_DOWN(U8, vxc_uchar8)
RESIZE_BILINEAR_F16TOQINT_DOWN(U16, vxc_ushort8)

__kernel void resize_bilinear_F16toF16_UP
    (
    image2d_array_t input,
    image2d_array_t output,
    int   align_corners,
    int   half_pixel_centers
    )
{
    int4 coord_out  =  (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0);
    int4   coord_x     = coord_out.xxxx + (int4)(0, 1, 2, 3);
    float4 in_x        = (convert_float4(coord_x) + half_pixel_value) * scale_xy.xxxx - half_pixel_value;
    float4 left_x_f    = floor(in_x);
    float4 x_lerp      = in_x - left_x_f;
    int4   left_x_idx  = convert_int4(left_x_f);
    float4 right_x_f   = ceil(in_x);
    int4   right_x_idx = convert_int4(right_x_f);
    float  in_y        = (convert_float(coord_out.y) + half_pixel_value) * scale_xy.y - half_pixel_value;
    float  top_y_f     = floor(in_y);
    float  y_lerp      = in_y - top_y_f;
    int    top_y_idx   = convert_int(top_y_f);


    vxc_ushort8 src0, src1, src2, src3, dst0, dst1;
    vxc_half8 top;
    vxc_half8 bottom;
    int4 coord_in = (int4)(left_x_idx.x, top_y_idx, coord_out.z, 0);

    int8 input_desc;
    _viv_asm(COPY, input_desc, input, sizeof(input_desc));
    int baseAddr = (int)coord_in.z * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord_in.w, baseAddr);

    VXC_OP4(img_load_3d, src0, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    VXC_OP4(img_load_3d, src1, input, coord_in.xywz, VXC_5BITOFFSET_XY(8, 0),
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    VXC_OP4(img_load_3d, src2, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1),
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    VXC_OP4(img_load_3d, src3, input, coord_in.xywz, VXC_5BITOFFSET_XY(8, 1),
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    vxc_ushort8 bitextract_p0;
    vxc_uchar16 maskShift = {16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16};
    VXC_DP2x8(bitextract_p0, left_x_idx, right_x_idx, \
    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertI32toI16_2x8);
    vxc_ushort8 constData = 16;
    VXC_DP2x8(maskShift, bitextract_p0, constData, \
    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniGetMaskShift_2x8);

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

    float4 left4;
    float4 right4;
    float4 top4;
    float4 bottom4;

    int loop = depth - 1;
    while (coord_in.z < loop)
    {
        VXC_BitExtract(dst0, src0, src1, maskShift, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_BitExtract(dst1, src2, src3, maskShift, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, top, dst0, 16);
        _viv_asm(COPY, bottom, dst1, 16);


        coord_in.zw += (int2)(1, input_desc.s4);
        VXC_OP4(img_load_3d, src0, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, src1, input, coord_in.xywz, VXC_5BITOFFSET_XY(8, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, src2, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 1),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, src3, input, coord_in.xywz, VXC_5BITOFFSET_XY(8, 1),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

        VXC_DP4x4(left4, top, top, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniFp16toFp32_left_4x4);
        VXC_DP4x4(right4, top, top, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniRightSubLeft_4x4);
        top4        = right4 * x_lerp + left4;
        VXC_DP4x4(left4, bottom, bottom, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniFp16toFp32_left_4x4);
        VXC_DP4x4(right4, bottom, bottom, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniRightSubLeft_4x4);
        bottom4      = right4 * x_lerp + left4;
        bottom4     -= top4;
        float4 dst4  = bottom4 * y_lerp + top4;
        half4 tmp;
        _viv_asm(CONV, tmp, dst4);
        VXC_DP2x8(top, tmp, tmp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniExtactHalf8_2x8);
        _viv_asm(COPY, dst0, top, 16);

        VXC_OP4_NoDest(img_store_3d, output, coord_out.xywz, dst0, VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));
        coord_out.zw += (int2)(1, output_desc.s4);
    }

    VXC_BitExtract(dst0, src0, src1, maskShift, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    VXC_BitExtract(dst1, src2, src3, maskShift, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    _viv_asm(COPY, top, dst0, 16);
    _viv_asm(COPY, bottom, dst1, 16);

    VXC_DP4x4(left4, top, top, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniFp16toFp32_left_4x4);
    VXC_DP4x4(right4, top, top, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniRightSubLeft_4x4);
    top4        = right4 * x_lerp + left4;
    VXC_DP4x4(left4, bottom, bottom, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniFp16toFp32_left_4x4);
    VXC_DP4x4(right4, bottom, bottom, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniRightSubLeft_4x4);
    bottom4      = right4 * x_lerp + left4;
    bottom4     -= top4;
    float4 dst4  = bottom4 * y_lerp + top4;
    half4 tmp;
    _viv_asm(CONV, tmp, dst4);
    VXC_DP2x8(top, tmp, tmp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniExtactHalf8_2x8);
    _viv_asm(COPY, dst0, top, 16);

    VXC_OP4_NoDest(img_store_3d, output, coord_out.xywz, dst0, VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));
}
