#include "cl_viv_vx_ext.h"

/*****************************layernorm uint8 to fp16****************************/
_viv_uniform int width;
_viv_uniform float dimRatio;
_viv_uniform VXC_512Bits UniFP16toFP32Lo4_dp4x4;
_viv_uniform VXC_512Bits uniConvert1stUint8SubZpToFp32_4x4;
_viv_uniform VXC_512Bits uniConvert2ndUint8SubZpToFp32_4x4;
_viv_uniform VXC_512Bits uniConvert3rdUint8SubZpToFp32_4x4;
_viv_uniform VXC_512Bits uniConvert4thUint8SubZpToFp32_4x4;
_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 int tmpZp2;
_viv_uniform float e2InScale;
_viv_uniform VXC_512Bits uniConvertSecFp16Fp32_4x4;
_viv_uniform VXC_512Bits UniPackFP16even_2x8;

__kernel void layer_norm_U8toF16(
    image2d_array_t input,
    image2d_t bias,
    image2d_t scale,
    image2d_array_t output,
              float eps)
{
    int4 coord = (int4)(0, get_global_id(1), get_global_id(2), 0);
    int4 coord_out = coord;
    vxc_uchar16 src0;
    float sum = 0, sqr = 0;
    int tmpSum = 0, tmpSqr = 0;
    vxc_int4 tmpSum1;
    vxc_int4 tmpSqr1;

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

    _viv_asm(COPY, output_desc, output, sizeof(output_desc));
    int baseAddr = (int)get_global_id(2) * output_desc.s4 + output_desc.s0;
    _viv_asm(MOV, coord_out.w, baseAddr);

    for(coord.x = 0; coord.x < width; coord.x += 16)
    {
        VXC_OP4(img_load_3d, src0, input, coord.xyzz, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
        VXC_DP16x1(tmpSum1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumU8_16x1);
        tmpSum += (tmpSum1.x);
        VXC_DP16x1(tmpSqr1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSqrSum_16x1);
        tmpSqr += (tmpSqr1.x + tmpZp1 * tmpSum1.x);
    }
    sum = (tmpSum + sumInZp) * input_scale;
    sqr = (tmpSqr + tmpZp2) * e2InScale;

    float mean, vari;
    mean = sum * dimRatio;
    vari = sqr*dimRatio - mean*mean;
    vari += eps;
    vari = rsqrt(vari);
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3;
    int2 coord_bias = (int2)(0, 0);
    vxc_half8 scale_h;
    vxc_float4 bias_f0, bias_f1, scale_f0, scale_f1;
    vxc_short8 src1, outval;
    short zp = inputZP;
    half4 tmpVal0, tmpVal1;
    vxc_half8 dst;

    for(coord.x = 0; coord.x < width; coord.x += 16)
    {
        VXC_OP4(img_load_3d, src0, input, coord.xyzz, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
        VXC_ReadImage(src1, scale, coord.xw, VXC_5BITOFFSET_XY(0, 0),\
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord_bias.x = coord.x;
        _viv_asm(COPY, scale_h, src1, 16);
        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertSecFp16Fp32_4x4);
        bias_f0 = read_imagef(bias, coord_bias);
        coord_bias.x += 4;
        bias_f1 = read_imagef(bias, coord_bias);
        coord_bias.x += 4;

        VXC_ReadImage(src1, scale, coord.xw, VXC_5BITOFFSET_XY(8, 0),\
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, scale_h, src1, 16);
        VXC_DP4x4(tmpData0, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvert1stUint8SubZpToFp32_4x4);
        VXC_DP4x4(tmpData1, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvert2ndUint8SubZpToFp32_4x4);
        VXC_DP4x4(tmpData2, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvert3rdUint8SubZpToFp32_4x4);
        VXC_DP4x4(tmpData3, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvert4thUint8SubZpToFp32_4x4);
        tmpData0 *= input_scale;
        tmpData1 *= input_scale;
        tmpData2 *= input_scale;
        tmpData3 *= input_scale;

        vxc_float4 norm;
        tmpData0 -= mean;
        norm = scale_f0 * vari * tmpData0 + bias_f0;
        bias_f0 = read_imagef(bias, coord_bias);
        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
        coord_bias.x += 4;
        _viv_asm(CONV, tmpVal0, norm);

        tmpData1 -= mean;
        norm = scale_f1 * vari * tmpData1 + bias_f1;
        bias_f1 = read_imagef(bias, coord_bias);
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertSecFp16Fp32_4x4);
        _viv_asm(CONV, tmpVal1, norm);
        VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),\
            UniPackFP16even_2x8);
        _viv_asm(COPY, outval, dst, 16);
        coord_out.x = coord.x;
        VXC_OP4_NoDest(img_store_3d, output, coord_out.xyww, outval, \
                VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));

        tmpData2 -= mean;
        norm = scale_f0 * vari * tmpData2 + bias_f0;
        _viv_asm(CONV, tmpVal0, norm);

        tmpData3 -= mean;
        norm = scale_f1 * vari * tmpData3 + bias_f1;
        _viv_asm(CONV, tmpVal1, norm);
        VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),\
            UniPackFP16even_2x8);
        _viv_asm(COPY, outval, dst, 16);
        coord_out.x += 8;
        VXC_OP4_NoDest(img_store_3d, output, coord_out.xyww, outval, \
                VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
    }
}

__kernel void layer_norm_U8toF16_2D(
    image2d_t input,
    image2d_t bias,
    image2d_t scale,
    image2d_t output,
        float eps)
{
    int4 coord = (int4)(0, get_global_id(1), 0, 0);
    vxc_uchar16 src0;
    float sum = 0, sqr = 0;
    int tmpSum = 0, tmpSqr = 0;
    vxc_int4 tmpSum1;
    vxc_int4 tmpSqr1;

    for(coord.x = 0; coord.x < width; coord.x += 16)
    {
        VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
        VXC_DP16x1(tmpSum1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumU8_16x1);
        tmpSum += (tmpSum1.x);
        VXC_DP16x1(tmpSqr1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSqrSum_16x1);
        tmpSqr += (tmpSqr1.x + tmpZp1 * tmpSum1.x);
    }
    sum = (tmpSum + sumInZp) * input_scale;
    sqr = (tmpSqr + tmpZp2) * e2InScale;

    float mean, vari;
    mean = sum * dimRatio;
    vari = sqr*dimRatio - mean*mean;
    vari += eps;
    vari = rsqrt(vari);
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3;
    int2 coord_bias = (int2)(0, 0);
    vxc_half8 scale_h;
    vxc_float4 bias_f0, bias_f1, scale_f0, scale_f1;
    vxc_short8 src1, outval;
    short zp = inputZP;
    half4 tmpVal0, tmpVal1;
    vxc_half8 dst;

    int2 coord_out = (int2)(get_global_id(0), get_global_id(1));

    for(coord.x = 0; coord.x < width; coord.x += 16)
    {
        coord_bias.x = coord.x;
        VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
        VXC_ReadImage(src1, scale, coord.xw, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, scale_h, src1, 16);
        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertSecFp16Fp32_4x4);
        bias_f0 = read_imagef(bias, coord_bias);
        coord_bias.x += 4;
        bias_f1 = read_imagef(bias, coord_bias);
        coord_bias.x += 4;

        VXC_ReadImage(src1, scale, coord.xw, VXC_5BITOFFSET_XY(8, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        _viv_asm(COPY, scale_h, src1, 16);
        VXC_DP4x4(tmpData0, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvert1stUint8SubZpToFp32_4x4);
        VXC_DP4x4(tmpData1, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvert2ndUint8SubZpToFp32_4x4);
        VXC_DP4x4(tmpData2, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvert3rdUint8SubZpToFp32_4x4);
        VXC_DP4x4(tmpData3, src0, zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvert4thUint8SubZpToFp32_4x4);
        tmpData0 *= input_scale;
        tmpData1 *= input_scale;
        tmpData2 *= input_scale;
        tmpData3 *= input_scale;

        vxc_float4 norm;
        tmpData0 -= mean;
        norm = scale_f0 * vari * tmpData0 + bias_f0;
        bias_f0 = read_imagef(bias, coord_bias);
        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
        coord_bias.x += 4;
        _viv_asm(CONV, tmpVal0, norm);

        tmpData1 -= mean;
        norm = scale_f1 * vari * tmpData1 + bias_f1;
        bias_f1 = read_imagef(bias, coord_bias);
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertSecFp16Fp32_4x4);
        _viv_asm(CONV, tmpVal1, norm);
        VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),\
            UniPackFP16even_2x8);
        _viv_asm(COPY, outval, dst, 16);
        coord_out.x = coord.x;
        VXC_WriteImage(output, coord_out, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

        tmpData2 -= mean;
        norm = scale_f0 * vari * tmpData2 + bias_f0;
        _viv_asm(CONV, tmpVal0, norm);

        tmpData3 -= mean;
        norm = scale_f1 * vari * tmpData3 + bias_f1;
        _viv_asm(CONV, tmpVal1, norm);
        VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),\
            UniPackFP16even_2x8);
        _viv_asm(COPY, outval, dst, 16);
        coord_out.x += 8;
        VXC_WriteImage(output, coord_out, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
}
