#include "cl_viv_vx_ext.h"

_viv_uniform int width;
_viv_uniform int height;
_viv_uniform VXC_512Bits uniExtract8Data_2x8;
_viv_uniform float output_scale;
_viv_uniform float output_zp;

_viv_uniform float input_scale;
_viv_uniform float input_scale2;
_viv_uniform float input_zp;
_viv_uniform float sum_x_tail;
_viv_uniform float sum_x2_tail0;
_viv_uniform float sum_x2_tail1;

_viv_uniform VXC_512Bits uniSumX_16x1;
_viv_uniform VXC_512Bits uniSumX2_16x1;
_viv_uniform VXC_512Bits uniResetFp32_4x4;
_viv_uniform int group_stride;

#define GROUP_NORM_SUMS_8BITS_IMPL(name, src_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_sums_##name( \
    __read_only  image2d_array_t input, \
    __write_only image2d_array_t output, \
                 float eps, int is2D) \
{ \
    int gidx = get_global_id(0) << 4; \
    int lidx = get_local_id(0); \
    int gidz = get_global_id(1); \
    int4 coord = (int4)(gidx, 0, gidz, 0); \
    src_type src0; \
    float2 sums_f32 = 0; \
    int2 sums = 0, sum_x_x2; \
 \
    __local float lcl_sum[16]; \
    __local float lcl_sqr[16]; \
    int8 input_desc; \
    _viv_asm(COPY, input_desc, input, sizeof(input_desc)); \
    int baseAddr_a = (int)get_global_id(1) * input_desc.s4 + input_desc.s0; \
    _viv_asm(MOV, coord.z, baseAddr_a); \
    if (gidx < width) \
    { \
        for(coord.y = 0; coord.y < height;) \
        { \
            VXC_OP4(img_load_3d, src0, input, coord, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
            coord.y++; \
            VXC_DP16x1(sum_x_x2, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumX_16x1); \
            VXC_DP16x1(sum_x_x2, src0, src0, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0), uniSumX2_16x1); \
            sums = sums + sum_x_x2; \
        } \
        sums_f32 = convert_float2(sums); \
        sums_f32.y = sums_f32.y * input_scale2 + sum_x2_tail0 + sum_x2_tail1 * sums_f32.x; \
        sums_f32.x = sums_f32.x * input_scale + sum_x_tail; \
    } \
    lcl_sum[lidx] = sums_f32.x; \
    lcl_sqr[lidx] = sums_f32.y; \
    barrier(CLK_LOCAL_MEM_FENCE); \
 \
    int4 coord_out = (int4)(get_group_id(0) << 2, gidz, 0, 0); \
    if(lidx == 0) \
    { \
        float4 one = (float4)(1, 1, 1, 1); \
        __local float4* tmp_sum = (__local float4*)lcl_sum; \
        __local float4* tmp_sqr = (__local float4*)lcl_sqr; \
        float sum_x = 0,sum_x2 = 0; \
        for(int i = 0; i < 4; i++) \
        { \
            sum_x += dot(tmp_sum[i], one); \
            sum_x2 += dot(tmp_sqr[i], one); \
        } \
        float4 data = (float4)(sum_x, sum_x2, 0, 0); \
        write_imagef(output, coord_out, data); \
    } \
}
GROUP_NORM_SUMS_8BITS_IMPL(U8, vxc_uchar16)
GROUP_NORM_SUMS_8BITS_IMPL(I8, vxc_char16)

#define GROUP_NORM_SUMS_8BITS_IMPL_2D(name, src_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_sums_##name##_2D( \
    __read_only  image2d_array_t input, \
    __write_only image2d_array_t output, \
                 float eps, int is2D) \
{ \
    int gidx = get_global_id(0) << 4; \
    int lidx = get_local_id(0); \
 \
    int2 coord = (int2)(gidx, get_global_id(1)); \
    src_type src0; \
    float2 sums = 0; \
 \
    __local float lcl_sum[16]; \
    __local float lcl_sqr[16]; \
    if(gidx < width) \
    { \
        VXC_ReadImage(src0, input, coord, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        VXC_DP16x1(sums, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumX_16x1); \
        VXC_DP16x1(sums, src0, src0, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0), uniSumX2_16x1); \
        sums.y = sums.y * input_scale2 + sum_x2_tail0 + sum_x2_tail1 * sums.x; \
        sums.x = sums.x * input_scale + sum_x_tail; \
    } \
    lcl_sum[lidx] = sums.x; \
    lcl_sqr[lidx] = sums.y; \
    barrier(CLK_LOCAL_MEM_FENCE); \
 \
    int4 coord_out = (int4)(get_group_id(0) << 2, get_global_id(1), 0, 0); \
    if(lidx == 0) \
    { \
        float4 one = (float4)(1, 1, 1, 1); \
        __local float4* tmp_sum = (__local float4*)lcl_sum; \
        __local float4* tmp_sqr = (__local float4*)lcl_sqr; \
        float sum_x = 0,sum_x2 = 0; \
        for(int i = 0; i < 4; i++) \
        { \
            sum_x += dot(tmp_sum[i], one); \
            sum_x2 += dot(tmp_sqr[i], one); \
        } \
        float4 data = (float4)(sum_x, sum_x2, 0, 0); \
        write_imagef(output, coord_out, data); \
    } \
}
GROUP_NORM_SUMS_8BITS_IMPL_2D(U8, vxc_uchar16)
GROUP_NORM_SUMS_8BITS_IMPL_2D(I8, vxc_char16)

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void group_norm_means(
    image2d_t input, image2d_t output, float eps, float group_ratio)
{
    int gidx = get_global_id(0);
    int lidx = get_local_id(0);

    int2 coord = (int2)(gidx, get_global_id(1));
    vxc_uchar16 src0 = 1;
    float2 sum_sqr = (float2)(0);
    float4 mean_vari;
    VXC_DP4x4(mean_vari, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniResetFp32_4x4);

    __local float2 lcl_data[16];
    __local float2 lcl_sum[4];

    for(; coord.x < group_stride; coord.x += 64)
    {
        mean_vari += read_imagef(input, coord);
    }
    lcl_data[lidx] = mean_vari.xy;
    barrier(CLK_LOCAL_MEM_FENCE);
    if(lidx < 4)
    {
        float2 tmpSum = (float2)(0);
        for(int i = lidx; i < 16; i+=4)
        {
            tmpSum += lcl_data[i];
        }
        lcl_sum[lidx] = tmpSum;
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    if(lidx == 0)
    {
        for(int i = 0; i < 4; i++)
        {
            sum_sqr += lcl_sum[i];
        }
        mean_vari.xy = sum_sqr * group_ratio;
        mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps;
        mean_vari.s1 = rsqrt(mean_vari.s1);

        coord.x = 0;
        write_imagef(output, coord, mean_vari);
    }
}

_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;
#define CONVERT_INPUT_TO_F32() \
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);

#define GROUP_NORM_8BITS_IMPL(name, src_type, dst_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; \
    dst_type dst; \
    vxc_short8 src1; \
    vxc_half8 scale_h; \
    float scale_vari, bias_val; \
    float4 bias_f, scale_f; \
 \
    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_int4 tmpVal0, tmpVal1; \
    float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    float alpha = input_scale * output_scale * scale_vari; \
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0) * output_scale + output_zp; \
    bias_val = bias_val - input_zp * alpha; \
 \
    CONVERT_INPUT_TO_F32() \
    norm = tmpData0 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData1 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    norm = tmpData2 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData3 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(8, 15, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    VXC_WriteImage2DArray(output, coord, dst, VXC_MODIFIER(0, 15, 0,VXC_RM_TowardZero, 0)); \
}
GROUP_NORM_8BITS_IMPL(U8_F16toU8, vxc_uchar16, vxc_uchar16)
GROUP_NORM_8BITS_IMPL(I8_F16toI8, vxc_char16,  vxc_char16)

#define GROUP_NORM_8BITS_IMPL_2D(name, src_type, dst_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; \
    dst_type dst; \
    vxc_short8 src1; \
    vxc_half8 scale_h; \
    float scale_vari, bias_val; \
    float4 bias_f, scale_f; \
 \
    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_int4 tmpVal0, tmpVal1; \
    float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    float alpha = input_scale * output_scale * scale_vari; \
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0) * output_scale + output_zp; \
    bias_val = bias_val - input_zp * alpha; \
 \
    CONVERT_INPUT_TO_F32() \
    norm = tmpData0 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData1 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    norm = tmpData2 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData3 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(8, 15, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    VXC_WriteImage(output, coord.xy, dst, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
}
GROUP_NORM_8BITS_IMPL_2D(U8_F16toU8, vxc_uchar16, vxc_uchar16)
GROUP_NORM_8BITS_IMPL_2D(I8_F16toI8, vxc_char16,  vxc_char16)

#define GROUP_NORM_8BITS_F32_IMPL(name, src_type, dst_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; \
    dst_type dst; \
    float scale_vari, bias_val; \
    float4 bias_f, scale_f; \
 \
    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_int4 tmpVal0, tmpVal1; \
    float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    float alpha = input_scale * output_scale * scale_vari; \
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0) * output_scale + output_zp; \
    bias_val = bias_val - input_zp * alpha; \
 \
    CONVERT_INPUT_TO_F32() \
    norm = tmpData0 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData1 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    norm = tmpData2 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData3 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(8, 15, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    VXC_WriteImage2DArray(output, coord, dst, VXC_MODIFIER(0, 15, 0,VXC_RM_TowardZero, 0)); \
}
GROUP_NORM_8BITS_F32_IMPL(U8_F32toU8, vxc_uchar16, vxc_uchar16)
GROUP_NORM_8BITS_F32_IMPL(I8_F32toI8, vxc_char16,  vxc_char16)

#define GROUP_NORM_8BITS_F32_IMPL_2D(name, src_type, dst_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; \
    dst_type dst; \
    float scale_vari, bias_val; \
    float4 bias_f, scale_f; \
 \
    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_int4 tmpVal0, tmpVal1; \
    float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    float alpha = input_scale * output_scale * scale_vari; \
    bias_val = (bias_f.s0 - scale_vari * mean_vari.s0) * output_scale + output_zp; \
    bias_val = bias_val - input_zp * alpha; \
 \
    CONVERT_INPUT_TO_F32() \
    norm = tmpData0 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData1 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    norm = tmpData2 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData3 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(8, 15, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    VXC_WriteImage(output, coord.xy, dst, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
}
GROUP_NORM_8BITS_F32_IMPL_2D(U8_F32toU8, vxc_uchar16, vxc_uchar16)
GROUP_NORM_8BITS_F32_IMPL_2D(I8_F32toI8, vxc_char16,  vxc_char16)
