#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniDataToFP32_0_4x4;
_viv_uniform VXC_512Bits uniDataToFP32_1_4x4;
_viv_uniform VXC_512Bits uniExtract8Data_2x8;
_viv_uniform int height;
_viv_uniform uint group_num;
_viv_uniform float output_zp;
_viv_uniform float output_scale;
_viv_uniform float inv_multiplier;


__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void layernorm_axis01_I8_F16toF16(
    image2d_array_t input, image2d_array_t bias, image2d_array_t scale, image2d_t meanVari,
    image2d_array_t output, float eps)
{
    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_sum = (int2)(0, gidz);
    int4 coord_para = coord;
    vxc_char16 src0;
    vxc_short8 src1, outval;
    vxc_half8 scale_h, dst;
    vxc_float4 bias_f0, bias_f1, scale_f0, scale_f1;
    vxc_float4 mean_vari = (vxc_float4)(0);

    for(int i = 0; i < group_num; i++)
    {
        mean_vari += read_imagef(meanVari, coord_sum);
        coord_sum.x += 4;
    }
    mean_vari *= inv_multiplier;
    mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps;
    mean_vari.s1 = rsqrt(mean_vari.s1);
    coord_para.z = 0;
    coord_para.w = 0;
    int4 coord_bias = coord_para;

    int8 input_desc, scale_desc, output_desc;
    _viv_asm(COPY, input_desc, input, sizeof(input_desc));
    int baseAddr_a = (int)coord_in.z * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord_in.z, baseAddr_a);

    _viv_asm(COPY, scale_desc, scale, sizeof(scale_desc));
    int baseAddr_c = (int)coord_para.z * scale_desc.s4 + scale_desc.s0;
    _viv_asm(MOV, coord_para.w, baseAddr_c);

    _viv_asm(COPY, output_desc, output, sizeof(output_desc));
    int baseAddr = (int)coord.z * output_desc.s4 + output_desc.s0;
    _viv_asm(MOV, coord.z, baseAddr);

    vxc_float4  tmpData0, tmpData1, norm;
    half4 tmpVal0, tmpVal1;

    for(coord.y = 0; coord.y < height; coord.y++)
    {
        VXC_OP4(img_load_3d, src0, input, coord_in, 0,
                        VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord_in.y ++;
        coord_para.y = coord.y; coord_bias.y = coord.y;
        VXC_OP4(img_load_3d, src1, scale, coord_para.xywz, 0,
                        VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        bias_f0 = read_imagef(bias, coord_bias);
        coord_bias.x += 4;
        bias_f1 = read_imagef(bias, coord_bias);
        coord_bias.x = coord.x;

        _viv_asm(COPY, scale_h, src1, 16);
        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
                uniDataToFP32_0_4x4);
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
                uniDataToFP32_1_4x4);
        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);
        tmpData0 = tmpData0 - mean_vari.s0;
        tmpData1 = tmpData1 - mean_vari.s0;

        norm = scale_f0 * mean_vari.s1 * tmpData0 + bias_f0;
        _viv_asm(CONV, tmpVal0, norm);
        norm = scale_f1 * mean_vari.s1 * tmpData1 + bias_f1;
        _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_OP4_NoDest(img_store_3d, output, coord, outval, \
                VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void layernorm_axis01_I8_F32toF16(
    image2d_array_t input, image2d_array_t bias, image2d_array_t scale, image2d_t meanVari,
    image2d_array_t output, float eps)
{
    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_sum = (int2)(0, gidz);
    int4 coord_para = coord;
    vxc_char16 src0;
    vxc_short8 outval;
    vxc_half8 scale_h, dst;
    vxc_float4 bias_f0, bias_f1, scale_f0, scale_f1;
    vxc_float4 mean_vari = (vxc_float4)(0);

    for(int i = 0; i < group_num; i++)
    {
        mean_vari += read_imagef(meanVari, coord_sum);
        coord_sum.x += 4;
    }
    mean_vari *= inv_multiplier;
    mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps;
    mean_vari.s1 = rsqrt(mean_vari.s1);
    coord_para.z = 0;
    coord_para.w = 0;
    int4 coord_bias = coord_para;

    int8 input_desc, scale_desc, output_desc;
    _viv_asm(COPY, input_desc, input, sizeof(input_desc));
    int baseAddr_a = (int)coord_in.z * 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)coord.z * output_desc.s4 + output_desc.s0;
    _viv_asm(MOV, coord.z, baseAddr);

    vxc_float4  tmpData0, tmpData1, norm;
    half4 tmpVal0, tmpVal1;

    for(coord.y = 0; coord.y < height; coord.y++)
    {
        VXC_OP4(img_load_3d, src0, input, coord_in, 0,
                        VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord_in.y ++;
        coord_bias.y = coord.y;

        bias_f0 = read_imagef(bias, coord_bias);
        scale_f0 = read_imagef(scale, coord_bias);
        coord_bias.x += 4;
        bias_f1 = read_imagef(bias, coord_bias);
        scale_f1 = read_imagef(scale, coord_bias);
        coord_bias.x = coord.x;

        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);
        tmpData0 = tmpData0 - mean_vari.s0;
        tmpData1 = tmpData1 - mean_vari.s0;

        norm = scale_f0 * mean_vari.s1 * tmpData0 + bias_f0;
        _viv_asm(CONV, tmpVal0, norm);
        norm = scale_f1 * mean_vari.s1 * tmpData1 + bias_f1;
        _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_OP4_NoDest(img_store_3d, output, coord, outval, \
                VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void layernorm_axis01_I8_F16toU8(
    image2d_array_t input, image2d_array_t bias, image2d_array_t scale, image2d_t meanVari,
    image2d_array_t output, float eps)
{
    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_sum = (int2)(0, gidz);
    int4 coord_para = coord;
    vxc_char16 src0;
    vxc_uchar16 outval;
    vxc_short8 src1;
    vxc_half8 scale_h;
    vxc_float4 bias_f0, bias_f1, scale_f0, scale_f1;
    vxc_float4 mean_vari = (vxc_float4)(0);

    for(int i = 0; i < group_num; i++)
    {
        mean_vari += read_imagef(meanVari, coord_sum);
        coord_sum.x += 4;
    }
    mean_vari *= inv_multiplier;
    mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps;
    mean_vari.s1 = rsqrt(mean_vari.s1);
    coord_para.z = 0;
    coord_para.w = 0;
    int4 coord_bias = coord_para;

    int8 input_desc, scale_desc, output_desc;
    _viv_asm(COPY, input_desc, input, sizeof(input_desc));
    int baseAddr_a = (int)coord_in.z * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord_in.z, baseAddr_a);

    _viv_asm(COPY, scale_desc, scale, sizeof(scale_desc));
    int baseAddr_c = (int)coord_para.z * scale_desc.s4 + scale_desc.s0;
    _viv_asm(MOV, coord_para.w, baseAddr_c);

    _viv_asm(COPY, output_desc, output, sizeof(output_desc));
    int baseAddr = (int)coord.z * output_desc.s4 + output_desc.s0;
    _viv_asm(MOV, coord.z, baseAddr);

    vxc_float4  tmpData0, tmpData1, norm;
    vxc_int4 tmpVal0, tmpVal1;

    for(coord.y = 0; coord.y < height; coord.y++)
    {
        VXC_OP4(img_load_3d, src0, input, coord_in, 0,
                        VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord_in.y ++;
        coord_para.y = coord.y; coord_bias.y = coord.y;
        VXC_OP4(img_load_3d, src1, scale, coord_para.xywz, 0,
                        VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        bias_f0 = read_imagef(bias, coord_bias);
        coord_bias.x += 4;
        bias_f1 = read_imagef(bias, coord_bias);
        coord_bias.x = coord.x;

        _viv_asm(COPY, scale_h, src1, 16);
        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
                uniDataToFP32_0_4x4);
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
                uniDataToFP32_1_4x4);

        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);
        tmpData0 = tmpData0 - mean_vari.s0;
        tmpData1 = tmpData1 - mean_vari.s0;

        norm = scale_f0 * mean_vari.s1 * tmpData0 + bias_f0;
        tmpVal0 = convert_int4_rte(norm * output_scale + output_zp);
        norm = scale_f1 * mean_vari.s1 * tmpData1 + bias_f1;
        tmpVal1 = convert_int4_rte(norm * output_scale + output_zp);

        VXC_DP2x8(outval, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
                uniExtract8Data_2x8);
        VXC_OP4_NoDest(img_store_3d, output, coord, outval, \
                VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
    }
}
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void layernorm_axis01_I8_F32toU8(
    image2d_array_t input, image2d_array_t bias, image2d_array_t scale, image2d_t meanVari,
    image2d_array_t output, float eps)
{
    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_sum = (int2)(0, gidz);
    int4 coord_para = coord;
    vxc_char16 src0;
    vxc_uchar16 outval;
    vxc_float4 bias_f0, bias_f1, scale_f0, scale_f1;
    vxc_float4 mean_vari = (vxc_float4)(0);

    for(int i = 0; i < group_num; i++)
    {
        mean_vari += read_imagef(meanVari, coord_sum);
        coord_sum.x += 4;
    }
    mean_vari *= inv_multiplier;
    mean_vari.s1 = mean_vari.s1 - mean_vari.s0 * mean_vari.s0 + eps;
    mean_vari.s1 = rsqrt(mean_vari.s1);
    coord_para.z = 0;
    coord_para.w = 0;
    int4 coord_bias = coord_para;

    int8 input_desc, scale_desc, output_desc;
    _viv_asm(COPY, input_desc, input, sizeof(input_desc));
    int baseAddr_a = (int)coord_in.z * 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)coord.z * output_desc.s4 + output_desc.s0;
    _viv_asm(MOV, coord.z, baseAddr);

    vxc_float4  tmpData0, tmpData1, norm;
    vxc_int4 tmpVal0, tmpVal1;

    for(coord.y = 0; coord.y < height; coord.y++)
    {
        VXC_OP4(img_load_3d, src0, input, coord_in, 0,
                        VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord_in.y ++;
        coord_bias.y = coord.y;

        bias_f0 = read_imagef(bias, coord_bias);
        scale_f0 = read_imagef(scale, coord_bias);
        coord_bias.x += 4;
        bias_f1 = read_imagef(bias, coord_bias);
        scale_f1 = read_imagef(scale, coord_bias);
        coord_bias.x = coord.x;

        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);
        tmpData0 = tmpData0 - mean_vari.s0;
        tmpData1 = tmpData1 - mean_vari.s0;

        norm = scale_f0 * mean_vari.s1 * tmpData0 + bias_f0;
        tmpVal0 = convert_int4_rte(norm * output_scale + output_zp);
        norm = scale_f1 * mean_vari.s1 * tmpData1 + bias_f1;
        tmpVal1 = convert_int4_rte(norm * output_scale + output_zp);

        VXC_DP2x8(outval, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
                uniExtract8Data_2x8);
        VXC_OP4_NoDest(img_store_3d, output, coord, outval, \
                VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
    }
}