#include "cl_viv_vx_ext.h"

_viv_uniform int width;
_viv_uniform int height;
_viv_uniform VXC_512Bits UniFP16toFP32Lo4_dp4x4;
_viv_uniform VXC_512Bits uniConvertHalfToFp16_2x8;
_viv_uniform VXC_512Bits uniConvertInt32toUint8_2x8;

_viv_uniform VXC_512Bits uniFp16SumSqr_dp8x2;
_viv_uniform VXC_512Bits uniConvertEndInt16Fp32_4x4;

_viv_uniform float outputScale;
_viv_uniform int output_ZP;

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_sumsqr_F16(
    image2d_array_t input, image2d_array_t output, float eps, int is2D)
{
    int gidx = get_global_id(0) << 3;
    int lidx = get_local_id(0);
    int gidz = get_global_id(1);
    int4 coord = (int4)(gidx, 0, gidz, 0);
    vxc_short8 src0;
    vxc_half8 in_h;
    vxc_float4 sumsqr;
    vxc_float4 tmpSumSqr = (vxc_float4)(0);

    __local float lcl_sum[16];
    __local float lcl_sqr[16];

    int8 input_desc;
    _viv_asm(COPY, input_desc, input, sizeof(input_desc));
    int baseAddr_a = (int)get_global_id(1) * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord.z, baseAddr_a);

    if(gidx < width)
    {
        for(coord.y = 0; coord.y < height;)
        {
            VXC_OP4(img_load_3d, src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            _viv_asm(COPY, in_h, src0, 16);
            VXC_DP8x2(sumsqr, in_h, in_h, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                uniFp16SumSqr_dp8x2);
            tmpSumSqr += sumsqr;
        }
    }

    lcl_sum[lidx] = tmpSumSqr.x;
    lcl_sqr[lidx] = tmpSumSqr.y;
    barrier(CLK_LOCAL_MEM_FENCE);

    int4 coord_out = (int4)(get_group_id(0) << 2, gidz, 0, 0);
    if(lidx == 0)
    {
        float4 one = (float4)(1, 1, 1, 1);
        __local float4* tmp_sum = (__local float4*)lcl_sum;
        __local float4* tmp_sqr = (__local float4*)lcl_sqr;

        float sum = 0;
        float sqr = 0;
        for(int i = 0; i < 4; i++)
        {
            //sum += lcl_sum[i];
            //sqr += lcl_sqr[i];
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }

        float4 data = (float4)(sum, sqr, 0, 0);
        write_imagef(output, coord_out, data);
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_sumsqr_F16_2D(
    image2d_array_t input, image2d_array_t output, float eps, int is2D)
{
    int gidx = get_global_id(0) << 3;
    int lidx = get_local_id(0);

    int2 coord = (int2)(gidx, get_global_id(1));
    vxc_short8 src0;
    vxc_half8 in_h;
    vxc_float4 sumsqr = (vxc_float4)(0);

    __local float lcl_sum[16];
    __local float lcl_sqr[16];

    if(gidx < width)
    {
        VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, in_h, src0, 16);
        VXC_DP8x2(sumsqr, in_h, in_h, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                uniFp16SumSqr_dp8x2);
    }

    lcl_sum[lidx] = sumsqr.x;
    lcl_sqr[lidx] = sumsqr.y;
    barrier(CLK_LOCAL_MEM_FENCE);

    int4 coord_out = (int4)(get_group_id(0) << 2, get_global_id(1), 0, 0);
    if(lidx == 0)
    {
        float4 one = (float4)(1, 1, 1, 1);
        __local float4* tmp_sum = (__local float4*)lcl_sum;
        __local float4* tmp_sqr = (__local float4*)lcl_sqr;

        float sum = 0;
        float sqr = 0;
        for(int i = 0; i < 4; i++)
        {
            //sum += lcl_sum[i];
            //sqr += lcl_sqr[i];
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }

        float4 data = (float4)(sum, sqr, 0, 0);
        write_imagef(output, coord_out, data);
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_F16toF16(
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari, image2d_array_t output,
    float eps, int is2D, float rSpaceOrg, int pStride)
{
    int gidy = get_global_id(1);
    int gidz = get_global_id(2);
    int4 coord = (int4)(get_global_id(0), gidy, gidz, 0);
    int4 coord_para = (int4)((convert_int(get_global_id(0) * rSpaceOrg) + gidy * pStride), gidz, 0, 0);
    vxc_short8 src0;
    vxc_short8 src1;
    vxc_half8 scale_h, in_h;
    float scale_vari, bias_val;
    vxc_float4 bias_f, scale_f;

    vxc_float4 mean_vari = read_imagef(meanVari, coord_para.zy);
    VXC_ReadImage(src1, scale, coord_para.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    bias_f = read_imagef(bias, coord_para.xy);
    VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    _viv_asm(COPY, scale_h, src1, 16);
    VXC_DP4x4(scale_f, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), UniFP16toFP32Lo4_dp4x4);

    scale_vari = scale_f.s0 * mean_vari.s1;
    vxc_float4  tmpData0, tmpData1;
    vxc_short8 outval;
    half4 tmpVal0, tmpVal1;
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0);
    vxc_half8 dst;

    _viv_asm(COPY, in_h, src0, 16);
    VXC_DP4x4(tmpData0, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
    VXC_DP4x4(tmpData1, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertEndInt16Fp32_4x4);

    vxc_float4 norm;
    norm = scale_vari * tmpData0 + bias_val;
    _viv_asm(CONV, tmpVal0, norm);
    norm = scale_vari * tmpData1 + bias_val;
    _viv_asm(CONV, tmpVal1, norm);
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),\
            uniConvertHalfToFp16_2x8);
    _viv_asm(COPY, outval, dst, 16);
    VXC_WriteImage2DArray(output, coord, outval, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_F16toF16_2D(
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari, image2d_array_t output,
    float eps, int is2D, float rSpaceOrg, int pStride)
{
    int gidz = get_global_id(1);
    int2 coord = (int2)(get_global_id(0), gidz);
    int4 coord_para = (int4)(convert_int(get_global_id(0) * rSpaceOrg), gidz, 0, 0);
    vxc_short8 src0;
    vxc_short8 src1;
    vxc_half8 scale_h, in_h;
    float scale_vari, bias_val;
    vxc_float4 bias_f, scale_f;

    vxc_float4 mean_vari = read_imagef(meanVari, coord_para.zy);
    VXC_ReadImage(src1, scale, coord_para.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    bias_f = read_imagef(bias, coord_para.xy);
    VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    _viv_asm(COPY, scale_h, src1, 16);
    VXC_DP4x4(scale_f, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), UniFP16toFP32Lo4_dp4x4);

    scale_vari = scale_f.s0 * mean_vari.s1;
    vxc_float4  tmpData0, tmpData1;
    vxc_short8 outval;
    half4 tmpVal0, tmpVal1;
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0);
    vxc_half8 dst;

    _viv_asm(COPY, in_h, src0, 16);
    VXC_DP4x4(tmpData0, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
        UniFP16toFP32Lo4_dp4x4);
    VXC_DP4x4(tmpData1, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
        uniConvertEndInt16Fp32_4x4);
    vxc_float4 norm;
    norm = scale_vari * tmpData0 + bias_val;
    _viv_asm(CONV, tmpVal0, norm);
    norm = scale_vari * tmpData1 + bias_val;
    _viv_asm(CONV, tmpVal1, norm);
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),\
        uniConvertHalfToFp16_2x8);
    _viv_asm(COPY, outval, dst, 16);
    VXC_WriteImage(output, coord.xy, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_F16toU8(
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari, image2d_array_t output,
    float eps, int is2D, float rSpaceOrg, int pStride)
{
    int gidy = get_global_id(1);
    int gidz = get_global_id(2);
    int4 coord = (int4)(get_global_id(0), gidy, gidz, 0);
    int4 coord_para = (int4)((convert_int(get_global_id(0) * rSpaceOrg) + gidy * pStride), gidz, 0, 0);
    vxc_short8 src0;
    vxc_short8 src1;
    vxc_half8 scale_h, in_h;
    float scale_vari, bias_val;
    vxc_float4 bias_f, scale_f;

    vxc_float4 mean_vari = read_imagef(meanVari, coord_para.zy);
    VXC_ReadImage(src1, scale, coord_para.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    bias_f = read_imagef(bias, coord_para.xy);
    VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    _viv_asm(COPY, scale_h, src1, 16);
    VXC_DP4x4(scale_f, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), UniFP16toFP32Lo4_dp4x4);

    scale_vari = scale_f.s0 * mean_vari.s1;
    vxc_float4  tmpData0, tmpData1;
    vxc_uchar16 outval;
    vxc_int4 tmpVal0, tmpVal1;
    float alpha = outputScale * scale_vari;
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0) * outputScale + output_ZP;

    _viv_asm(COPY, in_h, src0, 16);
    VXC_DP4x4(tmpData0, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
    VXC_DP4x4(tmpData1, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertEndInt16Fp32_4x4);

    vxc_float4 norm;
    norm = alpha * tmpData0 + bias_val;
    tmpVal0 = convert_int4_rte(norm);
    norm = alpha * tmpData1 + bias_val;
    tmpVal1 = convert_int4_rte(norm);
    VXC_DP2x8(outval, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
            uniConvertInt32toUint8_2x8);
    VXC_WriteImage2DArray(output, coord, outval, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_F16toU8_2D(
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari, image2d_array_t output,
    float eps, int is2D, float rSpaceOrg, int pStride)
{
    int gidz = get_global_id(1);
    int2 coord = (int2)(get_global_id(0), gidz);
    int4 coord_para = (int4)(convert_int(get_global_id(0) * rSpaceOrg), gidz, 0, 0);
    vxc_short8 src0;
    vxc_short8 src1;
    vxc_half8 scale_h, in_h;
    float scale_vari, bias_val;
    vxc_float4 bias_f, scale_f;

    vxc_float4 mean_vari = read_imagef(meanVari, coord_para.zy);
    VXC_ReadImage(src1, scale, coord_para.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    bias_f = read_imagef(bias, coord_para.xy);
    VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    _viv_asm(COPY, scale_h, src1, 16);
    VXC_DP4x4(scale_f, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), UniFP16toFP32Lo4_dp4x4);

    scale_vari = scale_f.s0 * mean_vari.s1;
    vxc_float4  tmpData0, tmpData1;
    vxc_uchar16 outval;
    vxc_int4 tmpVal0, tmpVal1;
    float alpha = outputScale * scale_vari;
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0) * outputScale + output_ZP;

    _viv_asm(COPY, in_h, src0, 16);
    VXC_DP4x4(tmpData0, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
    VXC_DP4x4(tmpData1, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertEndInt16Fp32_4x4);
    vxc_float4 norm;
    norm = alpha * tmpData0 + bias_val;
    tmpVal0 = convert_int4_rte(norm);
    norm = alpha * tmpData1 + bias_val;
    tmpVal1 = convert_int4_rte(norm);
    VXC_DP2x8(outval, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
            uniConvertInt32toUint8_2x8);
    VXC_WriteImage(output, coord.xy, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
}