#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniSumX_16x1;
_viv_uniform VXC_512Bits uniSumX2_16x1;
_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 VXC_512Bits uniExtract8Data_2x8;
_viv_uniform int width;
_viv_uniform float inv_multiplier;
_viv_uniform float output_scale;
_viv_uniform float output_zp;

#define CONV2F32(dst, src, section) \
        VXC_DP4x4(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
            uniDataToFP32_##section##_4x4);

#define LAYER_NORM_8BITS_IMPL(name, src_type) \
__kernel void layer_norm_axis0_##name( \
    __read_only  image2d_array_t input, \
    __read_only  image2d_t       bias, \
    __read_only  image2d_t       scale, \
    __write_only image2d_array_t output, \
                 float           eps) \
{ \
    int4 coord = (int4)(0, get_global_id(1), get_global_id(2), get_global_id(2)); \
    int4 coord_out = coord; \
 \
    src_type src0, dst; \
    vxc_short8 src1; \
    vxc_half8 scale_h; \
    float sum = 0, sqr = 0; \
    float4 bias_f0, bias_f1, scale_f0, scale_f1; \
    uint2 _sums = 0, sum_x_x2; \
 \
    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.z, baseAddr); \
 \
    for (coord.x = 0; coord.x < width; ) \
    { \
        VXC_OP4(img_load_3d, src0, input, coord, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        coord.x += 16; \
        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; \
    } \
 \
    float2 sums = convert_float2(_sums) * inv_multiplier; \
 \
    sums.y = sums.y - sums.x * sums.x + eps; \
    sums.y = rsqrt(sums.y); \
    int4 tmpVal0, tmpVal1; \
    float4  tmpData0, tmpData1, tmpData2, tmpData3; \
    int2 coord_bias = (int2)(0, 0); \
 \
    for(coord.x = 0; coord.x < width; coord.x += 16) \
    { \
        VXC_OP4(img_load_3d, src0, input, coord, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        VXC_ReadImage(src1, scale, coord.xw, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        coord_bias.x = coord.x; \
        _viv_asm(COPY, scale_h, src1, 16); \
        CONV2F32(scale_f0, scale_h, 0); \
        CONV2F32(scale_f1, scale_h, 1); \
        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); \
        CONV2F32(tmpData0, src0, 0); \
        CONV2F32(tmpData1, src0, 1); \
        CONV2F32(tmpData2, src0, 2); \
        CONV2F32(tmpData3, src0, 3); \
 \
        float4 norm; \
        tmpData0 = tmpData0 - sums.x; \
        norm = scale_f0 * sums.y * tmpData0 + bias_f0; \
        bias_f0 = read_imagef(bias, coord_bias); \
        CONV2F32(scale_f0, scale_h, 0); \
        coord_bias.x += 4; \
        tmpVal0 = convert_int4_rte(norm * output_scale + output_zp); \
 \
        tmpData1 = tmpData1 - sums.x; \
        norm = scale_f1 * sums.y * tmpData1 + bias_f1; \
        bias_f1 = read_imagef(bias, coord_bias); \
        CONV2F32(scale_f1, scale_h, 1); \
        tmpVal1 = convert_int4_rte(norm * output_scale + output_zp); \
        VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
            uniExtract8Data_2x8); \
 \
        tmpData2 = tmpData2 - sums.x; \
        norm = scale_f0 * sums.y * tmpData2 + bias_f0; \
        tmpVal0 = convert_int4_rte(norm * output_scale + output_zp); \
 \
        tmpData3 = tmpData3 - sums.x; \
        norm = scale_f1 * sums.y * tmpData3 + bias_f1; \
        tmpVal1 = convert_int4_rte(norm * output_scale + output_zp); \
        VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(8, 15, 0, VXC_RM_TowardZero, 1), \
            uniExtract8Data_2x8); \
        coord_out.x = coord.x; \
        VXC_OP4_NoDest(img_store_3d, output, coord_out, dst, VXC_MODIFIER(0, 15, 0,VXC_RM_TowardZero, 0)); \
    } \
}
LAYER_NORM_8BITS_IMPL(U8_F16toU8, vxc_uchar16)
LAYER_NORM_8BITS_IMPL(I8_F16toI8, vxc_char16)

#define LAYER_NORM_SUMS_2D() \
    uint2 _sums = 0, sum_x_x2; \
 \
    for (coord.x = 0; coord.x < width; ) \
    { \
        VXC_ReadImage(src0, input, coord.xy, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        coord.x += 16; \
        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; \
    } \
 \
    float2 sums = convert_float2(_sums) * inv_multiplier; \
 \
    sums.y = sums.y - sums.x * sums.x + eps; \
    sums.y = rsqrt(sums.y);

#define LAYER_NORM_8BITS_IMPL_2D(name, src_type) \
__kernel void layer_norm_axis0_##name##_2D( \
    __read_only  image2d_array_t input, \
    __read_only  image2d_t       bias, \
    __read_only  image2d_t       scale, \
    __write_only image2d_array_t output, \
                 float           eps) \
{ \
    int4 coord = (int4)(0, get_global_id(1), 0, 0); \
 \
    src_type src0, dst; \
    vxc_short8 src1; \
    vxc_half8 scale_h; \
    float sum = 0, sqr = 0; \
    float4 bias_f0, bias_f1, scale_f0, scale_f1; \
 \
    LAYER_NORM_SUMS_2D(); \
 \
    int4 tmpVal0, tmpVal1; \
    float4  tmpData0, tmpData1, tmpData2, tmpData3; \
    int2 coord_bias = (int2)(0, 0); \
 \
    for (coord.x = 0; coord.x < width; coord.x += 16) \
    { \
        VXC_ReadImage(src0, input, coord.xy, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        VXC_ReadImage(src1, scale, coord.xw, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        coord_bias.x = coord.x; \
        _viv_asm(COPY, scale_h, src1, 16); \
        CONV2F32(scale_f0, scale_h, 0); \
        CONV2F32(scale_f1, scale_h, 1); \
        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); \
        CONV2F32(tmpData0, src0, 0); \
        CONV2F32(tmpData1, src0, 1); \
        CONV2F32(tmpData2, src0, 2); \
        CONV2F32(tmpData3, src0, 3); \
 \
        float4 norm; \
        tmpData0 = tmpData0 - sums.x; \
        norm = scale_f0 * sums.y * tmpData0 + bias_f0; \
        bias_f0 = read_imagef(bias, coord_bias); \
        CONV2F32(scale_f0, scale_h, 0); \
        coord_bias.x += 4; \
        tmpVal0 = convert_int4_rte(norm * output_scale + output_zp); \
 \
        tmpData1 = tmpData1 - sums.x; \
        norm = scale_f1 * sums.y * tmpData1 + bias_f1; \
        bias_f1 = read_imagef(bias, coord_bias); \
        CONV2F32(scale_f1, scale_h, 1); \
        tmpVal1 = convert_int4_rte(norm * output_scale + output_zp); \
        VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
            uniExtract8Data_2x8); \
 \
        tmpData2 = tmpData2 - sums.x; \
        norm = scale_f0 * sums.y * tmpData2 + bias_f0; \
        tmpVal0 = convert_int4_rte(norm * output_scale + output_zp); \
 \
        tmpData3 = tmpData3 - sums.x; \
        norm = scale_f1 * sums.y * tmpData3 + bias_f1; \
        tmpVal1 = convert_int4_rte(norm * output_scale + output_zp); \
        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)); \
    } \
}
LAYER_NORM_8BITS_IMPL_2D(U8_F16toU8, vxc_uchar16)
LAYER_NORM_8BITS_IMPL_2D(I8_F16toI8, vxc_char16)

#define LAYER_NORM_8TOF16_IMPL(name, src_type) \
__kernel void layer_norm_axis0_##name( \
    __read_only  image2d_array_t input, \
    __read_only  image2d_t       bias, \
    __read_only  image2d_t       scale, \
    __write_only image2d_array_t output, \
                 float           eps) \
{ \
    int4 coord = (int4)(0, get_global_id(1), get_global_id(2), get_global_id(2)); \
    int4 coord_out = coord; \
 \
    src_type src0; \
    vxc_short8 src1; \
    vxc_half8 scale_h; \
    float sum = 0, sqr = 0; \
    float4 bias_f0, bias_f1, scale_f0, scale_f1; \
    uint2 _sums = 0, sum_x_x2; \
 \
    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.z, baseAddr); \
 \
    for (coord.x = 0; coord.x < width; ) \
    { \
        VXC_OP4(img_load_3d, src0, input, coord, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        coord.x += 16; \
        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; \
    } \
 \
    float2 sums = convert_float2(_sums) * inv_multiplier; \
 \
    sums.y = sums.y - sums.x * sums.x + eps; \
    sums.y = rsqrt(sums.y); \
    half4 tmpVal0, tmpVal1; \
    float4  tmpData0, tmpData1, tmpData2, tmpData3; \
    int2 coord_bias = (int2)(0, 0); \
 \
    vxc_short8 dst; \
    vxc_half8 result; \
    for(coord.x = 0; coord.x < width; coord.x += 16) \
    { \
        VXC_OP4(img_load_3d, src0, input, coord, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        VXC_ReadImage(src1, scale, coord.xw, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        coord_bias.x = coord.x; \
        _viv_asm(COPY, scale_h, src1, 16); \
        CONV2F32(scale_f0, scale_h, 0); \
        CONV2F32(scale_f1, scale_h, 1); \
        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); \
        CONV2F32(tmpData0, src0, 0); \
        CONV2F32(tmpData1, src0, 1); \
        CONV2F32(tmpData2, src0, 2); \
        CONV2F32(tmpData3, src0, 3); \
 \
        float4 norm; \
        tmpData0 = tmpData0 - sums.x; \
        norm = scale_f0 * sums.y * tmpData0 + bias_f0; \
        bias_f0 = read_imagef(bias, coord_bias); \
        CONV2F32(scale_f0, scale_h, 0); \
        coord_bias.x += 4; \
        _viv_asm(CONV, tmpVal0, norm); \
 \
        tmpData1 = tmpData1 - sums.x; \
        norm = scale_f1 * sums.y * tmpData1 + bias_f1; \
        bias_f1 = read_imagef(bias, coord_bias); \
        CONV2F32(scale_f1, scale_h, 1); \
        _viv_asm(CONV, tmpVal1, norm); \
        VXC_DP2x8(result, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
            uniExtract8Data_2x8); \
        _viv_asm(COPY, dst, result, 16); \
        coord_out.x = coord.x; \
        VXC_OP4_NoDest(img_store_3d, output, coord_out, dst, \
                VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0)); \
 \
        tmpData2 = tmpData2 - sums.x; \
        norm = scale_f0 * sums.y * tmpData2 + bias_f0; \
        _viv_asm(CONV, tmpVal0, norm); \
 \
        tmpData3 = tmpData3 - sums.x; \
        norm = scale_f1 * sums.y * tmpData3 + bias_f1; \
        _viv_asm(CONV, tmpVal1, norm); \
        VXC_DP2x8(result, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
            uniExtract8Data_2x8); \
        _viv_asm(COPY, dst, result, 16); \
        coord_out.x += 8; \
        VXC_OP4_NoDest(img_store_3d, output, coord_out, dst, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0)); \
    } \
}
LAYER_NORM_8TOF16_IMPL(U8_F16toF16, vxc_uchar16)
LAYER_NORM_8TOF16_IMPL(I8_F16toF16, vxc_char16)

#define LAYER_NORM_8TOF16_IMPL_2D(name, src_type) \
__kernel void layer_norm_axis0_##name##_2D( \
    __read_only  image2d_array_t input, \
    __read_only  image2d_t       bias, \
    __read_only  image2d_t       scale, \
    __write_only image2d_array_t output, \
                 float           eps) \
{ \
    int4 coord = (int4)(0, get_global_id(1), 0, 0); \
 \
    src_type src0; \
    vxc_short8 src1; \
    vxc_half8 scale_h; \
    float sum = 0, sqr = 0; \
    float4 bias_f0, bias_f1, scale_f0, scale_f1; \
 \
    LAYER_NORM_SUMS_2D(); \
 \
    half4 tmpVal0, tmpVal1; \
    float4  tmpData0, tmpData1, tmpData2, tmpData3; \
    int2 coord_bias = (int2)(0, 0); \
 \
    vxc_short8 dst; \
    vxc_half8 result; \
    for (coord.x = 0; coord.x < width; coord.x += 16) \
    { \
        VXC_ReadImage(src0, input, coord.xy, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
        VXC_ReadImage(src1, scale, coord.xw, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        coord_bias.x = coord.x; \
        _viv_asm(COPY, scale_h, src1, 16); \
        CONV2F32(scale_f0, scale_h, 0); \
        CONV2F32(scale_f1, scale_h, 1); \
        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); \
        CONV2F32(tmpData0, src0, 0); \
        CONV2F32(tmpData1, src0, 1); \
        CONV2F32(tmpData2, src0, 2); \
        CONV2F32(tmpData3, src0, 3); \
 \
        float4 norm; \
        tmpData0 = tmpData0 - sums.x; \
        norm = scale_f0 * sums.y * tmpData0 + bias_f0; \
        bias_f0 = read_imagef(bias, coord_bias); \
        CONV2F32(scale_f0, scale_h, 0); \
        coord_bias.x += 4; \
        _viv_asm(CONV, tmpVal0, norm); \
 \
        tmpData1 = tmpData1 - sums.x; \
        norm = scale_f1 * sums.y * tmpData1 + bias_f1; \
        bias_f1 = read_imagef(bias, coord_bias); \
        CONV2F32(scale_f1, scale_h, 1); \
        _viv_asm(CONV, tmpVal1, norm); \
        VXC_DP2x8(result, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
            uniExtract8Data_2x8); \
        _viv_asm(COPY, dst, result, 16); \
        VXC_WriteImage(output, coord.xy, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        coord.x += 8; \
 \
        tmpData2 = tmpData2 - sums.x; \
        norm = scale_f0 * sums.y * tmpData2 + bias_f0; \
        _viv_asm(CONV, tmpVal0, norm); \
 \
        tmpData3 = tmpData3 - sums.x; \
        norm = scale_f1 * sums.y * tmpData3 + bias_f1; \
        _viv_asm(CONV, tmpVal1, norm); \
        VXC_DP2x8(result, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), \
            uniExtract8Data_2x8); \
        _viv_asm(COPY, dst, result, 16); \
        VXC_WriteImage(output, coord.xy, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        coord.x -= 8; \
    } \
}
LAYER_NORM_8TOF16_IMPL_2D(U8_F16toF16, vxc_uchar16)
LAYER_NORM_8TOF16_IMPL_2D(I8_F16toF16, vxc_char16)
