#include "cl_viv_vx_ext.h"

_viv_uniform int width;
_viv_uniform int height;
_viv_uniform int channel;
_viv_uniform float dimRatio;
_viv_uniform float zpScaleSqr_i16;
_viv_uniform float zpScale2_i16;
_viv_uniform float sumScale_i16;

_viv_uniform VXC_512Bits uniSumU8_16x1;
_viv_uniform VXC_512Bits uniSqrSum_16x1;
_viv_uniform float input_scale;
_viv_uniform int inputZP;
_viv_uniform int sumInZp;
_viv_uniform int tmpZp1;
_viv_uniform float e2InScale;
_viv_uniform float rowSumScale;
_viv_uniform VXC_512Bits uniFp16SumSqr_dp8x2;
_viv_uniform VXC_512Bits uniConvertHalftoFp16_2x8;

#define MOMENTS_AXIS01_QINT(src0_type_name, read0_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void moments_axis01_##src0_type_name##toF16( \
    image2d_array_t input, image2d_t output_mean, image2d_t output_vari, \
    int axis, int axis_num) \
{ \
    int gidx = get_global_id(0) << 4; \
    int lidx = get_local_id(0); \
    int gidz = get_global_id(1); \
    int4 coord = (int4)(gidx, 0, gidz, 0); \
    read0_type src0; \
    float sum = 0, sqr = 0; \
 \
    __local float lcl_sum[16]; \
    __local float lcl_sqr[16]; \
 \
    for(coord.x = gidx; coord.x < width; coord.x += 256) \
    { \
        int tmpSum = 0, tmpSqr = 0, tmpSum1, tmpSqr1; \
        for(coord.y = 0; coord.y < height;) \
        { \
            VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
            coord.y++; \
            VXC_DP16x1(tmpSum1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumU8_16x1); \
            tmpSum += (tmpSum1); \
            VXC_DP16x1(tmpSqr1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSqrSum_16x1); \
            tmpSqr += (tmpSqr1 + tmpZp1 * tmpSum1); \
        } \
        sqr += (tmpSqr * e2InScale + rowSumScale); \
        sum += (tmpSum + sumInZp) * input_scale; \
    } \
    lcl_sum[lidx] = sum; \
    lcl_sqr[lidx] = sqr; \
    barrier(CLK_LOCAL_MEM_FENCE); \
    int2 coord_out = (int2)(gidz, 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; \
 \
        sum = (0); \
        sqr = (0); \
        for(int i = 0; i < 4; i++) \
        { \
            sum += dot(tmp_sum[i], one); \
            sqr += dot(tmp_sqr[i], one); \
        } \
        float4 mean, vari; \
        mean.x = sum * dimRatio; \
        vari.x = sqr * dimRatio; \
        vari.x = vari.x - mean.x * mean.x; \
 \
        half4 tmpMean, tmpVari; \
        vxc_half8 tmpVal; \
        vxc_short8 dst; \
        _viv_asm(CONV, tmpMean, mean); \
        _viv_asm(CONV, tmpVari, vari); \
        VXC_DP2x8(tmpVal, tmpMean, tmpVari, \
                 VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8); \
        _viv_asm(COPY, dst, tmpVal, 16); \
        VXC_WriteImage(output_mean, coord_out, dst.s0123, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
        VXC_WriteImage(output_vari, coord_out, dst.s4567, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
    } \
}
MOMENTS_AXIS01_QINT(U8, vxc_uchar16)
MOMENTS_AXIS01_QINT(I8, vxc_char16)

#define MOMENTS_AXIS01_QINT_2D(src0_type_name, read0_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void moments_axis01_##src0_type_name##toF16_2D( \
    image2d_t input, image2d_t output_mean, image2d_t output_vari, \
    int axis, int axis_num) \
{ \
    int gidx = get_global_id(0) << 4; \
    int lidx = get_local_id(0); \
    int2 coord = (int2)(gidx, 0); \
    read0_type src0; \
    float sum = 0, sqr = 0; \
 \
    __local float lcl_sum[16]; \
    __local float lcl_sqr[16]; \
 \
    for(coord.x = gidx; coord.x < width; coord.x += 256) \
    { \
        int tmpSum = 0, tmpSqr = 0, tmpSum1, tmpSqr1; \
        for(coord.y = 0; coord.y < height;) \
        { \
            VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
            coord.y++; \
            VXC_DP16x1(tmpSum1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumU8_16x1); \
            tmpSum += (tmpSum1); \
            VXC_DP16x1(tmpSqr1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSqrSum_16x1); \
            tmpSqr += (tmpSqr1 + tmpZp1 * tmpSum1); \
        } \
        sqr += (tmpSqr * e2InScale + rowSumScale); \
        sum += (tmpSum + sumInZp) * input_scale; \
    } \
    lcl_sum[lidx] = sum; \
    lcl_sqr[lidx] = sqr; \
    barrier(CLK_LOCAL_MEM_FENCE); \
    int2 coord_out = (int2)(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; \
 \
        sum = (0); \
        sqr = (0); \
        for(int i = 0; i < 4; i++) \
        { \
            sum += dot(tmp_sum[i], one); \
            sqr += dot(tmp_sqr[i], one); \
        } \
        float4 mean, vari; \
        mean.x = sum * dimRatio; \
        vari.x = sqr * dimRatio; \
        vari.x = vari.x - mean.x * mean.x; \
 \
        half4 tmpMean, tmpVari; \
        vxc_half8 tmpVal; \
        vxc_short8 dst; \
        _viv_asm(CONV, tmpMean, mean); \
        _viv_asm(CONV, tmpVari, vari); \
        VXC_DP2x8(tmpVal, tmpMean, tmpVari, \
                 VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8); \
        _viv_asm(COPY, dst, tmpVal, 16); \
        VXC_WriteImage(output_mean, coord_out, dst.s0123, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
        VXC_WriteImage(output_vari, coord_out, dst.s4567, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
    } \
}
MOMENTS_AXIS01_QINT_2D(U8, vxc_uchar16)
MOMENTS_AXIS01_QINT_2D(I8, vxc_char16)

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void moments_axis01_F16toF16(
    image2d_array_t input, image2d_t output_mean, image2d_t output_vari,
    int axis, int axis_num)
{
    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];

    for(coord.x = gidx; coord.x < width; coord.x += 128)
    {
        for(coord.y = 0; coord.y < height;)
        {
            VXC_ReadImage2DArray(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);

    int2 coord_out = (int2)(gidz, 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.0f;
        float sqr = 0.0f;
        for(int i = 0; i < 4; i++)
        {
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }
        float4 mean, vari;
        mean.x = sum * dimRatio;
        vari.x = sqr * dimRatio;
        vari.x = vari.x - mean.x * mean.x;

        half4 tmpMean, tmpVari;
        vxc_half8 tmpVal;
        vxc_short8 dst;
        _viv_asm(CONV, tmpMean, mean);
        _viv_asm(CONV, tmpVari, vari);
        VXC_DP2x8(tmpVal, tmpMean, tmpVari,
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8);
        _viv_asm(COPY, dst, tmpVal, 16);

        VXC_WriteImage(output_mean, coord_out, dst.s0123, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
        VXC_WriteImage(output_vari, coord_out, dst.s4567, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void moments_axis01_F16toF16_2D(
    image2d_array_t input, image2d_t output_mean, image2d_t output_vari,
    int axis, int axis_num)
{
    int gidx = get_global_id(0) << 3;
    int lidx = get_local_id(0);
    int2 coord = (int2)(gidx, 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];

    for(coord.x = gidx; coord.x < width; coord.x += 128)
    {
        for(coord.y = 0; coord.y < height;)
        {
            VXC_ReadImage(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);

    int2 coord_out = (int2)(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.0f;
        float sqr = 0.0f;
        for(int i = 0; i < 4; i++)
        {
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }
        float4 mean, vari;
        mean.x = sum * dimRatio;
        vari.x = sqr * dimRatio;
        vari.x = vari.x - mean.x * mean.x;

        half4 tmpMean, tmpVari;
        vxc_half8 tmpVal;
        vxc_short8 dst;
        _viv_asm(CONV, tmpMean, mean);
        _viv_asm(CONV, tmpVari, vari);
        VXC_DP2x8(tmpVal, tmpMean, tmpVari,
                 VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8);
        _viv_asm(COPY, dst, tmpVal, 16);

        VXC_WriteImage(output_mean, coord_out, dst.s0123, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
        VXC_WriteImage(output_vari, coord_out, dst.s4567, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void moments_axis01_I16toF16(
    image2d_array_t input, image2d_t output_mean, image2d_t output_vari,
    int axis, int axis_num)
{
    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;
    float4 sumsqr;
    float sum = 0, sqr = 0;

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

    for(coord.x = gidx; coord.x < width; coord.x += 128)
    {
        float tmpSum = 0;
        for(coord.y = 0; coord.y < height;)
        {
            VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            VXC_DP8x2(sumsqr, src0, src0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                    uniFp16SumSqr_dp8x2);
            tmpSum += sumsqr.x;
            sqr += (sumsqr.y * e2InScale + zpScaleSqr_i16 + zpScale2_i16 * sumsqr.x);
        }
        sum += tmpSum * input_scale + sumScale_i16;
    }

    lcl_sum[lidx] = sum;
    lcl_sqr[lidx] = sqr;
    barrier(CLK_LOCAL_MEM_FENCE);

    int2 coord_out = (int2)(gidz, 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;

        sum = 0.0f;
        sqr = 0.0f;
        for(int i = 0; i < 4; i++)
        {
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }
        float4 mean, vari;
        mean.x = sum * dimRatio;
        vari.x = sqr * dimRatio;
        vari.x = vari.x - mean.x * mean.x;

        half4 tmpMean, tmpVari;
        vxc_half8 tmpVal;
        vxc_short8 dst;
        _viv_asm(CONV, tmpMean, mean);
        _viv_asm(CONV, tmpVari, vari);
        VXC_DP2x8(tmpVal, tmpMean, tmpVari,
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8);
        _viv_asm(COPY, dst, tmpVal, 16);

        VXC_WriteImage(output_mean, coord_out, dst.s0123, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
        VXC_WriteImage(output_vari, coord_out, dst.s4567, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void moments_axis01_I16toF16_2D(
    image2d_array_t input, image2d_t output_mean, image2d_t output_vari,
    int axis, int axis_num)
{
    int gidx = get_global_id(0) << 3;
    int lidx = get_local_id(0);
    int2 coord = (int2)(gidx, 0);
    vxc_short8 src0;
    float4 sumsqr;
    float sum = 0, sqr = 0;

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

    for(coord.x = gidx; coord.x < width; coord.x += 128)
    {
        float tmpSum = 0;
        for(coord.y = 0; coord.y < height;)
        {
            VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            VXC_DP8x2(sumsqr, src0, src0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                    uniFp16SumSqr_dp8x2);
            tmpSum += sumsqr.x;
            sqr += (sumsqr.y * e2InScale + zpScaleSqr_i16 + zpScale2_i16 * sumsqr.x);
        }
        sum += tmpSum * input_scale + sumScale_i16;
    }

    lcl_sum[lidx] = sum;
    lcl_sqr[lidx] = sqr;
    barrier(CLK_LOCAL_MEM_FENCE);

    int2 coord_out = (int2)(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;

        sum = 0.0f;
        sqr = 0.0f;
        for(int i = 0; i < 4; i++)
        {
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }
        float4 mean_vari;
        mean_vari.x = sum * dimRatio;
        mean_vari.y = sqr * dimRatio;
        mean_vari.y = mean_vari.y - mean_vari.x * mean_vari.x;

        half4 tmpMean, tmpVari;
        vxc_half8 tmpVal;
        vxc_short8 dst;
        _viv_asm(CONV, tmpMean, mean_vari);
        //_viv_asm(CONV, tmpVari, vari);
        VXC_DP2x8(tmpVal, tmpMean, tmpMean,
                 VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8);
        _viv_asm(COPY, dst, tmpVal, 16);

        VXC_WriteImage(output_mean, coord_out, dst.s0123, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
        VXC_WriteImage(output_vari, coord_out, dst.s1023, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    }
}
