#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniSumX_16x1;
_viv_uniform VXC_512Bits uniSumX2_16x1;
_viv_uniform VXC_512Bits uniSum_X_X2_8x2;
_viv_uniform int width;
_viv_uniform int height;



__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void layernorm_axis01_sums_F16toF32(
    image2d_array_t input, image2d_t output)
{
    int gidx = get_global_id(0) << 3;
    int lidx = get_local_id(0);
    int gidz = get_global_id(1);
    int4 coord = (int4)(gidx, 0, gidz, 0);
    vxc_short8 src0;
    vxc_half8 in_h;
    vxc_float4 sumsqr;
    vxc_float4 tmpSumSqr = (vxc_float4)(0);

    __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)coord.z * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord.w, baseAddr_a);

    if(gidx < width)
    {
        for(coord.y = 0; coord.y < height;)
        {
            VXC_OP4(img_load_3d, src0, input, coord.xywz, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            _viv_asm(COPY, in_h, src0, 16);
            VXC_DP8x2(sumsqr, in_h, in_h, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                uniSum_X_X2_8x2);
            tmpSumSqr += sumsqr;
        }
    }

    lcl_sum[lidx] = tmpSumSqr.x;
    lcl_sqr[lidx] = tmpSumSqr.y;
    barrier(CLK_LOCAL_MEM_FENCE);

    int2 coord_out = (int2)(get_group_id(0) << 2, gidz);
    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 = 0;
        float sqr = 0;
        for(int i = 0; i < 4; i++)
        {
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }

        float4 data = (float4)(sum, sqr, 0, 0);
        write_imagef(output, coord_out, data);
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void layernorm_axis01_sums_I16toF32(
    image2d_array_t input, image2d_t output)
{
    int gidx = get_global_id(0) << 3;
    int lidx = get_local_id(0);
    int gidz = get_global_id(1);
    int4 coord = (int4)(gidx, 0, gidz, 0);
    vxc_short8 src0;
    float4 tmpSumSqr = (float4)(0);

    __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)coord.z * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord.w, baseAddr_a);

    if(gidx < width)
    {
        for(coord.y = 0; coord.y < height;)
        {
            VXC_OP4(img_load_3d, src0, input, coord.xywz, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            vxc_float4 sumsqr;
            VXC_DP8x2(sumsqr, src0, src0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0),\
                    uniSum_X_X2_8x2);
            tmpSumSqr += sumsqr;
        }
    }
    lcl_sum[lidx] = tmpSumSqr.x;
    lcl_sqr[lidx] = tmpSumSqr.y;
    barrier(CLK_LOCAL_MEM_FENCE);

    int2 coord_out = (int2)(get_group_id(0) << 2, gidz);
    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;
        float4 data = (float4)(0);
        for(int i = 0; i < 4; i++)
        {
            data.x += dot(tmp_sum[i], one);
            data.y += dot(tmp_sqr[i], one);
        }
        write_imagef(output, coord_out, data);
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void layernorm_axis01_sums_U8toF32(
    image2d_array_t input, image2d_t output)
{
    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);
    vxc_uchar16 src0;
    float sum = 0, sqr = 0;
    int tmpSum = 0, tmpSqr = 0, tmpSum1, tmpSqr1;

    __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)coord.z * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord.w, baseAddr_a);

    if(gidx < width)
    {
        for(coord.y = 0; coord.y < height;)
        {
            VXC_OP4(img_load_3d, src0, input, coord.xywz, 0,
                    VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            VXC_DP16x1(tmpSum1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumX_16x1);
            tmpSum += (tmpSum1);
            VXC_DP16x1(tmpSqr1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumX2_16x1);
            tmpSqr += (tmpSqr1);
        }
        sqr += convert_float(tmpSqr);
        sum = convert_float(tmpSum);
    }
    lcl_sum[lidx] = sum;
    lcl_sqr[lidx] = sqr;
    barrier(CLK_LOCAL_MEM_FENCE);

    int2 coord_out = (int2)(get_group_id(0) << 2, gidz);
    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;
        sum = 0; sqr = 0;
        for(int i = 0; i < 4; i++)
        {
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }
        float4 data = (float4)(sum, sqr, 0, 0);
        write_imagef(output, coord_out, data);
    }
}

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void layernorm_axis01_sums_I8toF32(
    image2d_array_t input, image2d_t output)
{
    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);
    vxc_char16 src0;
    float sum = 0, sqr = 0;
    int tmpSum = 0, tmpSqr = 0, tmpSum1, tmpSqr1;

    __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)coord.z * input_desc.s4 + input_desc.s0;
    _viv_asm(MOV, coord.w, baseAddr_a);

    if(gidx < width)
    {
        for(coord.y = 0; coord.y < height;)
        {
            VXC_OP4(img_load_3d, src0, input, coord.xywz, 0,
                    VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            VXC_DP16x1(tmpSum1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumX_16x1);
            tmpSum += (tmpSum1);
            VXC_DP16x1(tmpSqr1, src0, src0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0), uniSumX2_16x1);
            tmpSqr += (tmpSqr1);
        }
        sqr += convert_float(tmpSqr);
        sum = convert_float(tmpSum);
    }
    lcl_sum[lidx] = sum;
    lcl_sqr[lidx] = sqr;
    barrier(CLK_LOCAL_MEM_FENCE);

    int2 coord_out = (int2)(get_group_id(0) << 2, gidz);
    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;
        sum = 0; sqr = 0;
        for(int i = 0; i < 4; i++)
        {
            sum += dot(tmp_sum[i], one);
            sqr += dot(tmp_sqr[i], one);
        }
        float4 data = (float4)(sum, sqr, 0, 0);
        write_imagef(output, coord_out, data);
    }
}
