#include "cl_viv_vx_ext.h"


_viv_uniform VXC_512Bits uniExtract8Data_2x8;

_viv_uniform VXC_512Bits uniDataToFP32_0_4x4;
_viv_uniform VXC_512Bits uniDataToFP32_1_4x4;
_viv_uniform VXC_512Bits uniDataToFP32_2_4x4;
_viv_uniform VXC_512Bits uniDataToFP32_3_4x4;
_viv_uniform float input_scale;
_viv_uniform float input_zp;

#define GROUP_NORM_8BITSTOF16_IMPL(name, src_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_##name( \
    __read_only  image2d_array_t input, \
    __read_only  image2d_t       bias, \
    __read_only  image2d_array_t scale, \
    __read_only  image2d_t       meanVari, \
    __write_only image2d_array_t output, \
    float eps, int is2D, float rSpaceOrg, float pStride) \
{ \
    int gidx = get_global_id(0); \
    int gidy = get_global_id(1); \
    int gidz = get_global_id(2); \
    int4 coord = (int4)(gidx, gidy, gidz, 0); \
    int4 coord_para = (int4)((convert_int(gidx * rSpaceOrg) + convert_int(gidy * pStride)), gidz, 0, 0); \
    src_type 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, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
    bias_f = read_imagef(bias, coord_para.xy); \
    VXC_ReadImage2DArray(src0, input, coord, 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), uniDataToFP32_0_4x4); \
 \
    scale_vari = scale_f.s0 * mean_vari.s1; \
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    half4 tmpVal0, tmpVal1; \
    float alpha = scale_vari * input_scale; \
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0); \
    bias_val = bias_val - input_zp * alpha; \
 \
    VXC_DP4x4(tmpData0, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_0_4x4); \
    VXC_DP4x4(tmpData1, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_1_4x4); \
    VXC_DP4x4(tmpData2, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_2_4x4); \
    VXC_DP4x4(tmpData3, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_3_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), uniExtract8Data_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), uniExtract8Data_2x8); \
    _viv_asm(COPY, outval, dst, 16); \
    VXC_WriteImage2DArray(output, coord, outval, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0)); \
}
GROUP_NORM_8BITSTOF16_IMPL(U8_F16toF16, vxc_uchar16)
GROUP_NORM_8BITSTOF16_IMPL(I8_F16toF16, vxc_char16)


#define GROUP_NORM_8BITSTOF16_IMPL_2D(name, src_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_##name##_2D( \
    __read_only  image2d_array_t input, \
    __read_only  image2d_t       bias, \
    __read_only  image2d_array_t scale, \
    __read_only  image2d_t       meanVari, \
    __write_only image2d_array_t output, \
    float eps, int is2D, float rSpaceOrg, float 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); \
    src_type 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, 0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0)); \
    bias_f = read_imagef(bias, coord_para.xy); \
    VXC_ReadImage(src0, input, coord.xy, 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), uniDataToFP32_0_4x4); \
 \
    scale_vari = scale_f.s0 * mean_vari.s1; \
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    half4 tmpVal0, tmpVal1; \
    float alpha = scale_vari; \
    alpha = scale_vari * input_scale; \
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0); \
    bias_val = bias_val - input_zp * alpha; \
 \
    VXC_DP4x4(tmpData0, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_0_4x4); \
    VXC_DP4x4(tmpData1, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_1_4x4); \
    VXC_DP4x4(tmpData2, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_2_4x4); \
    VXC_DP4x4(tmpData3, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_3_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), uniExtract8Data_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), uniExtract8Data_2x8); \
    _viv_asm(COPY, outval, dst, 16); \
    VXC_WriteImage(output, coord.xy, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}
GROUP_NORM_8BITSTOF16_IMPL_2D(U8_F16toF16, vxc_uchar16)
GROUP_NORM_8BITSTOF16_IMPL_2D(I8_F16toF16, vxc_char16)

#define GROUP_NORM_8TOF16_F32_IMPL(name, src_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_##name( \
    __read_only  image2d_array_t input, \
    __read_only  image2d_t       bias, \
    __read_only  image2d_t       scale, \
    __read_only  image2d_t       meanVari, \
    __write_only image2d_array_t output, \
    float eps, int is2D, float rSpaceOrg, float pStride) \
{ \
    int gidx = get_global_id(0); \
    int gidy = get_global_id(1); \
    int gidz = get_global_id(2); \
    int4 coord = (int4)(gidx, gidy, gidz, 0); \
    int4 coord_para = (int4)((convert_int(gidx * rSpaceOrg) + convert_int(gidy * pStride)), gidz, 0, 0); \
    src_type src0; \
    vxc_short8 outval; \
    vxc_half8 dst; \
    float scale_vari, bias_val; \
    vxc_float4 bias_f, scale_f; \
 \
    vxc_float4 mean_vari = read_imagef(meanVari, coord_para.zy); \
    bias_f = read_imagef(bias, coord_para.xy); \
    scale_f = read_imagef(scale, coord_para.xy); \
    VXC_ReadImage2DArray(src0, input, coord, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
 \
    scale_vari = scale_f.s0 * mean_vari.s1; \
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    half4 tmpVal0, tmpVal1; \
    float alpha = scale_vari * input_scale; \
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0); \
    bias_val = bias_val - input_zp * alpha; \
 \
    VXC_DP4x4(tmpData0, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_0_4x4); \
    VXC_DP4x4(tmpData1, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_1_4x4); \
    VXC_DP4x4(tmpData2, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_2_4x4); \
    VXC_DP4x4(tmpData3, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_3_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), uniExtract8Data_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), uniExtract8Data_2x8); \
    _viv_asm(COPY, outval, dst, 16); \
    VXC_WriteImage2DArray(output, coord, outval, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0)); \
}
GROUP_NORM_8TOF16_F32_IMPL(U8_F32toF16, vxc_uchar16)
GROUP_NORM_8TOF16_F32_IMPL(I8_F32toF16, vxc_char16)

#define GROUP_NORM_8TOF16_F32_IMPL_2D(name, src_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_##name##_2D( \
    __read_only  image2d_array_t input, \
    __read_only  image2d_t       bias, \
    __read_only  image2d_t       scale, \
    __read_only  image2d_t       meanVari, \
    __write_only image2d_array_t output, \
    float eps, int is2D, float rSpaceOrg, float 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); \
    src_type src0; \
    vxc_short8 outval; \
    vxc_half8 dst; \
    float scale_vari, bias_val; \
    vxc_float4 bias_f, scale_f; \
 \
    vxc_float4 mean_vari = read_imagef(meanVari, coord_para.zy); \
    bias_f = read_imagef(bias, coord_para.xy); \
    scale_f = read_imagef(scale, coord_para.xy); \
    VXC_ReadImage(src0, input, coord.xy, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
 \
    scale_vari = scale_f.s0 * mean_vari.s1; \
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    half4 tmpVal0, tmpVal1; \
    float alpha = scale_vari * input_scale; \
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0); \
    bias_val = bias_val - input_zp * alpha; \
 \
    VXC_DP4x4(tmpData0, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_0_4x4); \
    VXC_DP4x4(tmpData1, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_1_4x4); \
    VXC_DP4x4(tmpData2, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_2_4x4); \
    VXC_DP4x4(tmpData3, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_3_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), uniExtract8Data_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), uniExtract8Data_2x8); \
    _viv_asm(COPY, outval, dst, 16); \
    VXC_WriteImage(output, coord.xy, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}
GROUP_NORM_8TOF16_F32_IMPL_2D(U8_F32toF16, vxc_uchar16)
GROUP_NORM_8TOF16_F32_IMPL_2D(I8_F32toF16, vxc_char16)
