#include "cl_viv_vx_ext.h"

_viv_uniform float scale_x;
_viv_uniform int out_height;
_viv_uniform VXC_512Bits uniConvertI32toI16_2x8;
_viv_uniform VXC_512Bits uniGetMaskShift_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_odd_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_even_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_Part0_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_Part1_2x8;
_viv_uniform float half_pixel_value;

__kernel void resize_1d_bilinear_BF16toBF16_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_x - 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);

    vxc_short8 top;
    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    int4 coord_in = (int4)(left_x_idx.x, coord_out.y, 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);

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

    do
    {
        VXC_OP4(img_load_3d, top, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0));
        coord_in.x = left_x_idx.y;
        VXC_OP4(img_load_3d, top, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0));
        coord_in.x = left_x_idx.z;
        VXC_OP4(img_load_3d, top, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(4, 5, 0, VXC_RM_TowardZero, 0));
        coord_in.x = left_x_idx.w;
        VXC_OP4(img_load_3d, top, input, coord_in.xywz, VXC_5BITOFFSET_XY(0, 0),
            VXC_MODIFIER(6, 7, 0, VXC_RM_TowardZero, 0));
        coord_in.x = left_x_idx.x;

        vxc_ushort8 src;
        float4 left4;
        float4 right4;
        float4 dst4;

        VXC_DP2x8(src, top, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_even_2x8);
        _viv_asm(COPY, right4, src, 16);
        VXC_DP2x8(src, top, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_odd_2x8);
        _viv_asm(COPY, left4,  src, 16);
        right4    -= left4;
        dst4       = right4 * x_lerp + left4;
        vxc_ushort8 tmp, dst;
        _viv_asm(COPY, tmp, dst4, 16);
        dst.s0123 = tmp.s1357;
        VXC_OP4_NoDest(img_store_3d, output, coord_out.xywz, dst,
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));
        coord_in.y++;
        coord_out.y ++;
    } while (coord_out.y < out_height);
}

__kernel void resize_1d_bilinear_BF16toBF16_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_x - 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);

    vxc_ushort8 src0, src1, dst0;
    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    int4 coord_in = (int4)(left_x_idx.x, coord_out.y, 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_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);

    do
    {
        VXC_BitExtract(dst0, src0, src1, maskShift, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

        coord_in.y ++;
        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_ushort8 dst_tmp;
        float4 left4;
        float4 right4;

        VXC_DP2x8(dst_tmp, dst0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part0_2x8);
        _viv_asm(COPY, left4, dst_tmp, 16);
        VXC_DP2x8(dst_tmp, dst0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part1_2x8);
        _viv_asm(COPY, right4, dst_tmp, 16);
        right4     -= left4;
        float4 dst4 = right4 * x_lerp + left4;

        vxc_ushort8 tmp, dst;
        _viv_asm(COPY, tmp, dst4, 16);
        dst.s0123 = tmp.s1357;

        VXC_OP4_NoDest(img_store_3d, output, coord_out.xywz, dst,
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));
        coord_out.y++;
    } while (coord_out.y < out_height);
}
