#include "cl_viv_vx_ext.h"

_viv_uniform int width;
_viv_uniform int height;
_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 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_AXIS0_QINT(src0_type_name, read0_type) \
__kernel void moments_axis0_##src0_type_name##toF16( \
    image2d_array_t input, \
    image2d_t output_mean, \
    image2d_t output_vari, \
    int axis, int axis_num) \
{ \
    int gidy = get_global_id(0); \
    int gidz = get_global_id(1); \
    int4 coord = (int4)(0, gidy, gidz, 0); \
    read0_type src0; \
    float sum = 0, sqr = 0; \
    int tmpSum = 0, tmpSqr = 0; \
    int4 tmpSum0, tmpSqr0; \
 \
    for(coord.x = 0; coord.x < width; coord.x += 16) \
    { \
        VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
                VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        VXC_DP16x1(tmpSum0, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumU8_16x1); \
        VXC_DP16x1(tmpSqr0, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSqrSum_16x1); \
        tmpSum += (tmpSum0.x); \
        tmpSqr += (tmpSqr0.x + tmpZp1 * tmpSum0.x); \
    } \
    sqr = (convert_float(tmpSqr) * e2InScale + rowSumScale); \
    sum = convert_float(tmpSum + sumInZp) * input_scale; \
 \
    vxc_float4 mean_vari0 = (vxc_float4)(sum, sqr, 0, 0); \
    mean_vari0 *= dimRatio; \
    mean_vari0.s1 = mean_vari0.s1 - mean_vari0.s0 * mean_vari0.s0; \
 \
    int2 coord_out = (int2)(gidy, gidz); \
    half4 tmpData; \
    vxc_half8 tmpVal; \
    vxc_short8 dst; \
    _viv_asm(CONV, tmpData, mean_vari0); \
    VXC_DP2x8(tmpVal, tmpData, tmpData, 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)); \
}

MOMENTS_AXIS0_QINT(U8, vxc_uchar16)
MOMENTS_AXIS0_QINT(I8, vxc_char16)

#define MOMENTS_AXIS0_QINT_2D(src0_type_name, read0_type) \
__kernel void moments_axis0_##src0_type_name##toF16_2D( \
    image2d_t input, \
    image2d_t output_mean, \
    image2d_t output_vari, \
    int axis, int axis_num) \
{ \
    int gidy = get_global_id(0); \
    int2 coord = (int2)(0, gidy); \
    read0_type src0; \
    float sum = 0, sqr = 0; \
    int tmpSum = 0, tmpSqr = 0; \
    int4 tmpSum0, tmpSqr0; \
 \
    for(coord.x = 0; coord.x < width; coord.x += 16) \
    { \
        VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
                VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        VXC_DP16x1(tmpSum0, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumU8_16x1); \
        VXC_DP16x1(tmpSqr0, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSqrSum_16x1); \
        tmpSum += (tmpSum0.x); \
        tmpSqr += (tmpSqr0.x + tmpZp1 * tmpSum0.x); \
    } \
    sqr = (convert_float(tmpSqr) * e2InScale + rowSumScale); \
    sum = convert_float(tmpSum + sumInZp) * input_scale; \
 \
    vxc_float4 mean_vari0 = (vxc_float4)(sum, sqr, 0, 0); \
    mean_vari0 *= dimRatio; \
    mean_vari0.s1 = mean_vari0.s1 - mean_vari0.s0 * mean_vari0.s0; \
 \
    int2 coord_out = (int2)(gidy, 0); \
    half4 tmpData; \
    vxc_half8 tmpVal; \
    vxc_short8 dst; \
    _viv_asm(CONV, tmpData, mean_vari0); \
    VXC_DP2x8(tmpVal, tmpData, tmpData, 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)); \
}
MOMENTS_AXIS0_QINT_2D(U8, vxc_uchar16)
MOMENTS_AXIS0_QINT_2D(I8, vxc_char16)

__kernel void moments_axis0_F16toF16(
    image2d_array_t input,
    image2d_t output_mean,
    image2d_t output_vari,
              int axis, int axis_num)
{
    int gidy = get_global_id(0);
    int gidz = get_global_id(1);
    int4 coord = (int4)(0, gidy, gidz, 0);
    vxc_short8 src0;
    vxc_half8 in_h0;
    vxc_float4 sumsqr0;
    vxc_float4 mean_vari0 = (vxc_float4)(0);
    for(coord.x = 0; coord.x < width; coord.x += 8)
    {
        VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, in_h0, src0, 16);
        VXC_DP8x2(sumsqr0, in_h0, in_h0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                uniFp16SumSqr_dp8x2);
        mean_vari0 += sumsqr0;
    }

    mean_vari0 *= dimRatio;
    mean_vari0.s1 = mean_vari0.s1 - mean_vari0.s0 * mean_vari0.s0;

    int2 coord_out = (int2)(gidy, gidz);

    half4 tmpData;
    vxc_half8 tmpVal;
    vxc_short8 dst;
    _viv_asm(CONV, tmpData, mean_vari0);
    VXC_DP2x8(tmpVal, tmpData, tmpData, 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));
}

__kernel void moments_axis0_F16toF16_2D(
    image2d_t input,
    image2d_t output_mean,
    image2d_t output_vari,
              int axis, int axis_num)
{
    int gidy = get_global_id(0);
    int2 coord = (int2)(0, gidy);
    vxc_short8 src0;
    vxc_half8 in_h0;
    vxc_float4 sumsqr0;
    vxc_float4 mean_vari0 = (vxc_float4)(0);

    for(coord.x = 0; coord.x < width; coord.x += 8)
    {
        VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, in_h0, src0, 16);
        VXC_DP8x2(sumsqr0, in_h0, in_h0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                uniFp16SumSqr_dp8x2);
        mean_vari0 += sumsqr0;
    }
    mean_vari0 *= dimRatio;
    mean_vari0.s1 = mean_vari0.s1 - mean_vari0.s0 * mean_vari0.s0;

    int2 coord_out = (int2)(gidy, 0);

    half4 tmpData;
    vxc_half8 tmpVal;
    vxc_short8 dst;
    _viv_asm(CONV, tmpData, mean_vari0);
    VXC_DP2x8(tmpVal, tmpData, tmpData, 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));
}

__kernel void moments_axis0_I16toF16(
    image2d_array_t input,
    image2d_t output_mean,
    image2d_t output_vari,
              int axis, int axis_num)
{
    int gidy = get_global_id(0);
    int gidz = get_global_id(1);
    int4 coord = (int4)(0, gidy, gidz, 0);
    vxc_short8 src0;
    float4 sumsqr0;
    float sum = 0, sqr = 0;
    float tmpSum = 0;
    for(coord.x = 0; coord.x < width; coord.x += 8)
    {
        VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_DP8x2(sumsqr0, src0, src0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                uniFp16SumSqr_dp8x2);
        tmpSum += sumsqr0.x;
        sqr += (sumsqr0.y * e2InScale + zpScaleSqr_i16 + zpScale2_i16 * sumsqr0.x);
    }
    sum = tmpSum * input_scale + sumScale_i16;

    vxc_float4 mean_vari0 = (vxc_float4)(sum, sqr, 0, 0);
    mean_vari0 *= dimRatio;
    mean_vari0.s1 = mean_vari0.s1 - mean_vari0.s0 * mean_vari0.s0;

    int2 coord_out = (int2)(gidy, gidz);

    half4 tmpData;
    vxc_half8 tmpVal;
    vxc_short8 dst;
    _viv_asm(CONV, tmpData, mean_vari0);
    VXC_DP2x8(tmpVal, tmpData, tmpData, 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));
}

__kernel void moments_axis0_I16toF16_2D(
    image2d_t input,
    image2d_t output_mean,
    image2d_t output_vari,
              int axis, int axis_num)
{
    int gidy = get_global_id(0);
    int2 coord = (int2)(0, gidy);
    vxc_short8 src0;
    float4 sumsqr0;
    float sum = 0, sqr = 0;
    float tmpSum = 0;
    for(coord.x = 0; coord.x < width; coord.x += 8)
    {
        VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_DP8x2(sumsqr0, src0, src0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                uniFp16SumSqr_dp8x2);
        tmpSum += sumsqr0.x;
        sqr += (sumsqr0.y * e2InScale + zpScaleSqr_i16 + zpScale2_i16 * sumsqr0.x);
    }
    sum = tmpSum * input_scale + sumScale_i16;

    vxc_float4 mean_vari0 = (vxc_float4)(sum, sqr, 0, 0);
    mean_vari0 *= dimRatio;
    mean_vari0.s1 = mean_vari0.s1 - mean_vari0.s0 * mean_vari0.s0;

    int2 coord_out = (int2)(gidy, 0);

    half4 tmpData;
    vxc_half8 tmpVal;
    vxc_short8 dst;
    _viv_asm(CONV, tmpData, mean_vari0);
    VXC_DP2x8(tmpVal, tmpData, tmpData, 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));
}