#include "cl_viv_vx_ext.h"

_viv_uniform int height;
_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;

#define MOMENTS_AXIS1_QINT(src0_type_name, read0_type) \
__kernel void moments_axis1_##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 gidz = get_global_id(1); \
    int4 coord = (int4)(gidx, 0, gidz, 0); \
    read0_type src0; \
    float4 sum = 0, sqr = 0; \
    short zp = inputZP;\
    float4 tmpData0;\
 \
    VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
            VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
    for(coord.y = 1; coord.y < height; ) \
    { \
        VXC_DP4x4(tmpData0, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
                uniConvert1stUint8SubZpToFp32_4x4); \
        VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
                VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        coord.y++; \
        sum += (tmpData0); \
        sqr += (tmpData0 * tmpData0); \
    } \
    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, gidz); \
    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_AXIS1_QINT(U8, vxc_uchar16)
MOMENTS_AXIS1_QINT(I8, vxc_char16)
MOMENTS_AXIS1_QINT(I16, vxc_short8)

#define MOMENTS_AXIS1_QINT_2D(src0_type_name, read0_type) \
__kernel void moments_axis1_##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); \
    int2 coord = (int2)(gidx, 0); \
    read0_type src0; \
    float4 sum = 0, sqr = 0; \
    short zp = inputZP;\
    float4 tmpData0;\
    VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
            VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
 \
    for (coord.y = 1; coord.y < height; ) \
    { \
        VXC_DP4x4(tmpData0, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
                uniConvert1stUint8SubZpToFp32_4x4); \
        VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), \
                VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        coord.y++; \
        sum += (tmpData0); \
        sqr += (tmpData0 * tmpData0); \
    } \
    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, 0); \
    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_AXIS1_QINT_2D(U8, vxc_uchar16)
MOMENTS_AXIS1_QINT_2D(I8, vxc_char16)
MOMENTS_AXIS1_QINT_2D(I16, vxc_short8)

__kernel void moments_axis1_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 gidz = get_global_id(1);
    int4 coord = (int4)(gidx, 0, gidz, 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.y = 0; coord.y < height; coord.y++)
    {
        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, gidz);
    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_axis1_F16toF16_2D(
    image2d_t input,
    image2d_t output_mean,
    image2d_t output_vari,
              int axis, int axis_num)
{
    int gidx = get_global_id(0);
    int2 coord = (int2)(gidx, 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.y = 0; coord.y < height; coord.y++)
    {
        VXC_ReadImage(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, 0);
    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));
}
