#include "cl_viv_vx_ext.h"

_viv_uniform int channel;
_viv_uniform float dimRatio;

_viv_uniform float input_scale;
_viv_uniform int inputZP;
_viv_uniform float e2InScale;
_viv_uniform VXC_512Bits UniFP16toFP32Lo4_dp4x4;
_viv_uniform VXC_512Bits uniConvert1stUint8SubZpToFp32_4x4;
_viv_uniform VXC_512Bits uniConvertHalftoFp16_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_Part0_2x8;
_viv_uniform VXC_512Bits uniExtractOddData_2x8;

#define MOMENTS_AXIS2_QINT(src0_type_name, read0_type) \
__kernel void moments_axis2_##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); \
    int gidy = get_global_id(1); \
    int4 coord = (int4)(gidx, gidy, 0, 0); \
    read0_type src0; \
    float4 sum = 0, sqr = 0; \
    short zp = inputZP;\
    float4 tmpData0;\
 \
    for(coord.z = 0; coord.z < channel; coord.z++) \
    { \
        VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
                VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_DP4x4(tmpData0, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
                uniConvert1stUint8SubZpToFp32_4x4); \
        sum += (tmpData0); \
        sqr += (tmpData0 * tmpData0); \
    } \
    sum *= input_scale; \
    sqr *= e2InScale; \
 \
    float4 mean = sum * dimRatio; \
    float4 vari = sqr * dimRatio; \
    vari = vari - mean * mean; \
 \
    int2 coord_out = (int2)(gidx, gidy); \
    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, 3, 0, VXC_RM_TowardZero, 0)); \
    VXC_WriteImage(output_vari, coord_out, dst.s4567, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
}

MOMENTS_AXIS2_QINT(U8, vxc_uchar16)
MOMENTS_AXIS2_QINT(I8, vxc_char16)
MOMENTS_AXIS2_QINT(I16, vxc_short8)

__kernel void moments_axis2_F16toF16(
    image2d_array_t input,
    image2d_t output_mean,
    image2d_t output_vari,
              int axis,
              int axis_num)
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int4 coord = (int4)(gidx, gidy, 0, 0);
    vxc_short8 src0;
    vxc_half8 in_h0;
    vxc_float4 tmpSrc0;
    vxc_float4 sum = (vxc_float4)(0);
    vxc_float4 sqr = (vxc_float4)(0);

    for(coord.z = 0; coord.z < channel; coord.z++)
    {
        VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, in_h0, src0, 16);
        VXC_DP4x4(tmpSrc0, in_h0, in_h0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), UniFP16toFP32Lo4_dp4x4);
        sum += tmpSrc0;
        sqr += (tmpSrc0 * tmpSrc0);
    }

    vxc_float4 mean = sum * dimRatio;
    vxc_float4 vari = sqr * dimRatio;
    vari = vari - mean * mean;

    int2 coord_out = (int2)(gidx, gidy);
    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, 3, 0, VXC_RM_TowardZero, 0));
    VXC_WriteImage(output_vari, coord_out, dst.s4567, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
}

__kernel void moments_axis2_BF16toBF16(
    image2d_array_t input,
    image2d_t output_mean,
    image2d_t output_vari,
              int axis,
              int axis_num)
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int4 coord = (int4)(gidx, gidy, 0, 0);
    vxc_ushort8 src0;
    vxc_ushort8 val;
    vxc_float4 sum = (vxc_float4)(0);
    vxc_float4 sqr = (vxc_float4)(0);
    vxc_ushort8 zero = (vxc_ushort8)(0, 0, 0, 0, 0, 0, 0, 0);

    for(coord.z = 0; coord.z < channel; coord.z++)
    {
        VXC_ReadImage2DArray(val, input, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        float4 vec0;
        VXC_DP2x8(src0, val, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                uniConvBF16toF32_Part0_2x8);
        _viv_asm(COPY, vec0, src0, 16);

        sum += vec0;
        sqr += (vec0 * vec0);
    }

    vxc_float4 mean = sum * dimRatio;
    vxc_float4 vari = sqr * dimRatio;
    vari = vari - mean * mean;

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

    vxc_short8 tmpdst0, tmpdst1, dst;
    _viv_asm(COPY, tmpdst0, mean, 16);
    _viv_asm(COPY, tmpdst1, vari, 16);
    VXC_DP2x8(dst, tmpdst0, tmpdst1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniExtractOddData_2x8);

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