#include "cl_viv_vx_ext.h"

_viv_uniform int height;
_viv_uniform float dimRatio;
_viv_uniform int group_num;
_viv_uniform VXC_512Bits uniConvertInt32toUint8_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 int inputZP;
_viv_uniform float scale_inOut;
_viv_uniform float outputScale;
_viv_uniform int output_ZP;
_viv_uniform float inOut_fl_scale;
_viv_uniform float output_fl_scale;
_viv_uniform VXC_512Bits uniConvertInt16Fp32Fst_4x4;
_viv_uniform VXC_512Bits uniConvertInt16Fp32Secd_4x4;
_viv_uniform VXC_512Bits uniConvertInt32toInt16_2x8;

#define INSTANCENORM_8BITS_F32(src1_type_name, read_type) \
__kernel void instance_norm_##src1_type_name##F32to##src1_type_name( \
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari, \
    image2d_array_t output, float eps, int rsFlg) \
{ \
    int gidz = get_global_id(1); \
    int4 coord = (int4)(get_global_id(0), 0, gidz, gidz); \
    int4 coord_in = (int4)(get_global_id(0), 0, gidz, gidz); \
    int2 coord_para = (int2)(gidz, 0); \
    read_type src0, src2; \
    float scale_vari, bias_val; \
    vxc_float4 mean_vari = (vxc_float4)(0); \
 \
    Image img1 = create_image_from_image2d(bias, 4); \
    Image img2 = create_image_from_image2d(scale, 4); \
    Image img3 = create_image_from_image2d(meanVari, 4); \
    __global float* bias_ptr = (__global float*)img1.ptr; \
    __global float* scal_ptr = (__global float*)img2.ptr; \
    __global uchar* sumVari_ptr = (__global uchar*)get_image_ptr_from_coord(img3, coord_para.yx); \
    __global float4* vari_ptr = (__global float4*)sumVari_ptr; \
 \
    float bval = bias_ptr[gidz]; \
    float sval = scal_ptr[gidz]; \
 \
    for(int i = 0; i < group_num; i++) \
    { \
        mean_vari += vari_ptr[i]; \
    } \
    mean_vari *= dimRatio; \
    mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps; \
    mean_vari.s1 = rsqrt(mean_vari.s1); \
 \
    scale_vari = sval * mean_vari.s1; \
    short zp = inputZP; \
    vxc_int4 tmpVal0, tmpVal1; \
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    float alpha = scale_inOut * scale_vari; \
    bias_val = (bval - scale_vari * mean_vari.s0) * outputScale + output_ZP; \
 \
    int8 input_desc, output_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_in.z, baseAddr_a); \
    _viv_asm(COPY, output_desc, output, sizeof(output_desc)); \
    int baseAddr = (int)get_global_id(1) * output_desc.s4 + output_desc.s0; \
    _viv_asm(MOV, coord.z, baseAddr); \
 \
    for(coord.y = 0; coord.y < height; coord.y++) \
    { \
    VXC_OP4(img_load_3d, src0, input, coord_in, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
    coord_in.y ++; \
    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 = tmpData0 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData1 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(src2, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniConvertInt32toUint8_2x8); \
    norm = tmpData2 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData3 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(src2, tmpVal0, tmpVal1, VXC_MODIFIER(8, 15, 0, VXC_RM_TowardZero, 1), uniConvertInt32toUint8_2x8); \
    VXC_OP4_NoDest(img_store_3d, output, coord, src2, VXC_MODIFIER(0, 15, 0,VXC_RM_TowardZero, 0)); \
    } \
}
INSTANCENORM_8BITS_F32(U8, vxc_uchar16)
INSTANCENORM_8BITS_F32(I8, vxc_char16)

#define INSTANCENORM_8BITS_F32_2D(src1_type_name, read_type) \
__kernel void instance_norm_##src1_type_name##F32to##src1_type_name##_2D( \
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari, \
    image2d_array_t output, float eps, int rsFlg) \
{ \
    int gidz = get_global_id(1); \
    int gidy = gidz * height; \
    int2 coord = (int2)(get_global_id(0), gidy); \
    int2 coord_para = (int2)(gidz, 0); \
    int endH = gidy + height; \
    read_type src0, src2; \
    float scale_vari, bias_val; \
    vxc_float4 mean_vari = (vxc_float4)(0); \
 \
    Image img1 = create_image_from_image2d(bias, 4); \
    Image img2 = create_image_from_image2d(scale, 4); \
    Image img3 = create_image_from_image2d(meanVari, 4); \
    __global float* bias_ptr = (__global float*)img1.ptr; \
    __global float* scal_ptr = (__global float*)img2.ptr; \
    __global uchar* sumVari_ptr = (__global uchar*)get_image_ptr_from_coord(img3, coord_para.yx); \
    __global float4* vari_ptr = (__global float4*)sumVari_ptr; \
 \
    float bval = bias_ptr[gidz]; \
    float sval = scal_ptr[gidz]; \
 \
    for(int i = 0; i < group_num; i++) \
    { \
        mean_vari += vari_ptr[i]; \
    } \
 \
    mean_vari *= dimRatio; \
    mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps; \
    mean_vari.s1 = rsqrt(mean_vari.s1); \
 \
    scale_vari = sval * mean_vari.s1; \
    short zp = inputZP; \
    vxc_int4 tmpVal0, tmpVal1; \
    vxc_float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    float alpha = scale_inOut * scale_vari; \
    bias_val = (bval - scale_vari * mean_vari.s0) * outputScale + output_ZP; \
 \
    for(; coord.y < endH; coord.y++) \
    { \
    VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
    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 = tmpData0 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData1 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(src2, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniConvertInt32toUint8_2x8); \
    norm = tmpData2 * alpha + bias_val; \
    tmpVal0 = convert_int4_rte(norm); \
    norm = tmpData3 * alpha + bias_val; \
    tmpVal1 = convert_int4_rte(norm); \
    VXC_DP2x8(src2, tmpVal0, tmpVal1, VXC_MODIFIER(8, 15, 0, VXC_RM_TowardZero, 1), uniConvertInt32toUint8_2x8); \
    VXC_WriteImage(output, coord, src2, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
    } \
}
INSTANCENORM_8BITS_F32_2D(U8, vxc_uchar16)
INSTANCENORM_8BITS_F32_2D(I8, vxc_char16)

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void instance_norm_I16F32toI16(
    image2d_array_t input, image2d_t bias, image2d_t scale, image2d_t meanVari,
    image2d_array_t output, float eps, int rsFlg)
{
    int gidz = get_global_id(1);
    int4 coord = (int4)(get_global_id(0), 0, gidz, gidz);
    int4 coord_in = (int4)(get_global_id(0), 0, gidz, gidz);
    int2 coord_para = (int2)(gidz, 0);
    vxc_short8 src0, src2;
    float scale_vari, bias_val;
    vxc_float4  mean_vari = (vxc_float4)(0);

    Image img1 = create_image_from_image2d(bias, 4);
    Image img2 = create_image_from_image2d(scale, 4);
    Image img3 = create_image_from_image2d(meanVari, 4);
    __global float* bias_ptr = (__global float*)img1.ptr;
    __global float* scal_ptr = (__global float*)img2.ptr;
    __global uchar* sumVari_ptr = (__global uchar*)get_image_ptr_from_coord(img3, coord_para.yx);
    __global float4* vari_ptr = (__global float4*)sumVari_ptr;

    float bval = bias_ptr[gidz];
    float sval = scal_ptr[gidz];

    for(int i = 0; i < group_num; i++)
    {
        mean_vari += vari_ptr[i];
    }

    mean_vari *= dimRatio;
    mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps;
    mean_vari.s1 = rsqrt(mean_vari.s1);

    scale_vari = sval * mean_vari.s1;
    vxc_int4 tmpVal0, tmpVal1;
    vxc_float4  tmpData0, tmpData1;
    float alpha = inOut_fl_scale * scale_vari;
    bias_val = (bval - scale_vari * mean_vari.s0) * output_fl_scale;

    int8 input_desc, output_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_in.z, baseAddr_a);

    _viv_asm(COPY, output_desc, output, sizeof(output_desc));
    int baseAddr = (int)get_global_id(1) * output_desc.s4 + output_desc.s0;
    _viv_asm(MOV, coord.z, baseAddr);
    for(coord.y = 0; coord.y < height; coord.y++)
    {
    VXC_OP4(img_load_3d, src0, input, coord_in, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    coord_in.y ++;
    VXC_DP4x4(tmpData0, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
        uniConvertInt16Fp32Fst_4x4);
    VXC_DP4x4(tmpData1, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
        uniConvertInt16Fp32Secd_4x4);
    vxc_float4 norm;
    norm = tmpData0 * alpha + bias_val;
    tmpVal0 = convert_int4_rte(norm);
    norm = tmpData1 * alpha + bias_val;
    tmpVal1 = convert_int4_rte(norm);
    VXC_DP2x8(src2, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
        uniConvertInt32toInt16_2x8);
    VXC_OP4_NoDest(img_store_3d, output, coord, src2, \
                    VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void instance_norm_I16F32toI16_2D(
    image2d_t input, image2d_t bias, image2d_t scale,
    image2d_t meanVari, image2d_array_t output, float eps, int rsFlg)
{
    int gidz = get_global_id(1);
    int gidy = gidz * height;
    int2 coord = (int2)(get_global_id(0), gidy);
    int2 coord_para = (int2)(gidz, 0);
    int endH = gidy + height;
    vxc_short8 src0, src2;
    float scale_vari, bias_val;
    vxc_float4 mean_vari = (vxc_float4)(0);

    Image img1 = create_image_from_image2d(bias, 4);
    Image img2 = create_image_from_image2d(scale, 4);
    Image img3 = create_image_from_image2d(meanVari, 4);
    __global float* bias_ptr = (__global float*)img1.ptr;
    __global float* scal_ptr = (__global float*)img2.ptr;
    __global uchar* sumVari_ptr = (__global uchar*)get_image_ptr_from_coord(img3, coord_para.yx);
    __global float4* vari_ptr = (__global float4*)sumVari_ptr;

    float bval = bias_ptr[gidz];
    float sval = scal_ptr[gidz];

    for(int i = 0; i < group_num; i++)
    {
        mean_vari += vari_ptr[i];
    }

    mean_vari *= dimRatio;
    mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps;
    mean_vari.s1 = rsqrt(mean_vari.s1);

    scale_vari = sval * mean_vari.s1;
    vxc_int4 tmpVal0, tmpVal1;
    vxc_float4  tmpData0, tmpData1;
    float alpha = inOut_fl_scale * scale_vari;
    bias_val = (bval - scale_vari * mean_vari.s0) * output_fl_scale;

    for(; coord.y < endH; coord.y++)
    {
    VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),\
        VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    VXC_DP4x4(tmpData0, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
        uniConvertInt16Fp32Fst_4x4);
    VXC_DP4x4(tmpData1, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
        uniConvertInt16Fp32Secd_4x4);
    vxc_float4 norm;
    norm = tmpData0 * alpha + bias_val;
    tmpVal0 = convert_int4_rte(norm);
    norm = tmpData1 * alpha + bias_val;
    tmpVal1 = convert_int4_rte(norm);
    VXC_DP2x8(src2, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
        uniConvertInt32toInt16_2x8);
    VXC_WriteImage(output, coord, src2, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
}