#include "cl_viv_vx_ext.h"

_viv_uniform int width;
_viv_uniform int height;
_viv_uniform VXC_512Bits UniFP16toFP32Lo4_dp4x4;
_viv_uniform VXC_512Bits uniConvertHalfToFp16_2x8;

_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 float input_scale;
_viv_uniform int inputZP;

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_U8toF16(
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari,
    image2d_array_t output, float eps, int is2D, float rSpaceOrg, int pStride)
{
    int gidy = get_global_id(1);
    int gidz = get_global_id(2);
    int4 coord = (int4)(get_global_id(0), gidy, gidz, 0);
    int4 coord_para = (int4)((convert_int(get_global_id(0) * rSpaceOrg) + gidy * pStride), gidz, 0, 0);
    vxc_uchar16 src0;
    vxc_short8 src1, outval;
    vxc_half8 scale_h, dst;
    float scale_vari, bias_val;
    vxc_float4 bias_f, scale_f;

    vxc_float4 mean_vari = read_imagef(meanVari, coord_para.zy);
    VXC_ReadImage(src1, scale, coord_para.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    bias_f = read_imagef(bias, coord_para.xy);
    VXC_ReadImage2DArray(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));

    _viv_asm(COPY, scale_h, src1, 16);
    VXC_DP4x4(scale_f, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), UniFP16toFP32Lo4_dp4x4);

    scale_vari = scale_f.s0 * mean_vari.s1;
    short zp = inputZP;
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3, norm;
    half4 tmpVal0, tmpVal1;
    float alpha = input_scale * scale_vari;
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0);

    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);
    norm = alpha * tmpData0 + bias_val;
    _viv_asm(CONV, tmpVal0, norm);
    norm = alpha * tmpData1 + bias_val;
    _viv_asm(CONV, tmpVal1, norm);
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalfToFp16_2x8);
    _viv_asm(COPY, outval, dst, 16);
    VXC_WriteImage2DArray(output, coord, outval, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
    coord.x += 8;
    norm = alpha * tmpData2 + bias_val;
    _viv_asm(CONV, tmpVal0, norm);
    norm = alpha * tmpData3 + bias_val;
    _viv_asm(CONV, tmpVal1, norm);
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalfToFp16_2x8);
    _viv_asm(COPY, outval, dst, 16);
    VXC_WriteImage2DArray(output, coord, outval, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_U8toF16_2D(
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari,
    image2d_array_t output, float eps, int is2D, float rSpaceOrg, int pStride)
{
    int gidz = get_global_id(1);
    int2 coord = (int2)(get_global_id(0), gidz);
    int4 coord_para = (int4)(convert_int(get_global_id(0) * rSpaceOrg), gidz, 0, 0);
    vxc_uchar16 src0;
    vxc_short8 src1, outval;
    vxc_half8 scale_h, dst;
    float scale_vari, bias_val;
    vxc_float4 bias_f, scale_f;

    vxc_float4 mean_vari = read_imagef(meanVari, coord_para.zy);
    VXC_ReadImage(src1, scale, coord_para.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0));
    bias_f = read_imagef(bias, coord_para.xy);
    VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));

    _viv_asm(COPY, scale_h, src1, 16);
    VXC_DP4x4(scale_f, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), UniFP16toFP32Lo4_dp4x4);

    scale_vari = scale_f.s0 * mean_vari.s1;
    short zp = inputZP;
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3, norm;
    half4 tmpVal0, tmpVal1;
    float alpha = input_scale * scale_vari;
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0);

    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);
    norm = alpha * tmpData0 + bias_val;
    _viv_asm(CONV, tmpVal0, norm);
    norm = alpha * tmpData1 + bias_val;
    _viv_asm(CONV, tmpVal1, norm);
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalfToFp16_2x8);
    _viv_asm(COPY, outval, dst, 16);
    VXC_WriteImage(output, coord.xy, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    coord.x += 8;
    norm = alpha * tmpData2 + bias_val;
    _viv_asm(CONV, tmpVal0, norm);
    norm = alpha * tmpData3 + bias_val;
    _viv_asm(CONV, tmpVal1, norm);
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalfToFp16_2x8);
    _viv_asm(COPY, outval, dst, 16);
    VXC_WriteImage(output, coord.xy, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
}
