#include "cl_viv_vx_ext.h"

_viv_uniform int width;

_viv_uniform VXC_512Bits uniConvBF16toF32_Part0_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_Part1_2x8;
_viv_uniform VXC_512Bits uniExtractOddData_2x8;

__kernel __attribute__((reqd_work_group_size(16, 1, 1))) void instance_norm_sums_BF16(
    image2d_array_t input, image2d_array_t output, float eps, int height)
{
    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, gidz);
    vxc_short8 src0, src1, src2;
    float4 srcA, srcB;
    float sum = 0, sqr = 0;
    float4 one = (float4)(1.0, 1.0, 1.0, 1.0);
    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 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)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, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            VXC_DP2x8(src1, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                         uniConvBF16toF32_Part0_2x8);
            VXC_DP2x8(src2, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                         uniConvBF16toF32_Part1_2x8);
            _viv_asm(COPY, srcA, src1, 16);
            _viv_asm(COPY, srcB, src2, 16);
            sum += dot(srcA, one) + dot(srcB, one);
            sqr += dot(srcA * srcA, one) + dot(srcB * srcB, one);
        }
    }

    lcl_sum[lidx] = sum;
    lcl_sqr[lidx] = sqr;
    barrier(CLK_LOCAL_MEM_FENCE);

    int4 coord_out = (int4)(get_group_id(0) << 2, gidz, 0, 0);
    if(lidx == 0)
    {
        __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 instance_norm_sums_BF16_2D(
    image2d_array_t input, image2d_array_t output, float eps, int height)
{
    int gidx = get_global_id(0) << 3;
    int lidx = get_local_id(0);
    int gidz = get_global_id(1);
    int gidy = gidz * height;

    int2 coord = (int2)(gidx, gidy);
    vxc_short8 src0, src1, src2;
    float4 srcA, srcB;
    float sum = 0, sqr = 0;
    float4 one = (float4)(1.0, 1.0, 1.0, 1.0);
    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);

    __local float lcl_sum[16];
    __local float lcl_sqr[16];

    int endH = gidy + height;
    if(gidx < width)
    {
        for(; coord.y < endH;)
        {
            VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
            coord.y++;
            VXC_DP2x8(src1, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                         uniConvBF16toF32_Part0_2x8);
            VXC_DP2x8(src2, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                         uniConvBF16toF32_Part1_2x8);
            _viv_asm(COPY, srcA, src1, 16);
            _viv_asm(COPY, srcB, src2, 16);
            sum += dot(srcA, one) + dot(srcB, one);
            sqr += dot(srcA * srcA, one) + dot(srcB * srcB, one);
        }
    }

    lcl_sum[lidx] = sum;
    lcl_sqr[lidx] = sqr;
    barrier(CLK_LOCAL_MEM_FENCE);

    int4 coord_out = (int4)(get_group_id(0) << 2, gidz, 0, 0);
    if(lidx == 0)
    {
        __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 void instance_norm_BF16_F32toBF16(
    __read_only  image2d_array_t input,
    __read_only  image2d_t       means,
    __write_only image2d_array_t output,
                 int height)
{
    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);
    vxc_short8 src0, src1, src2;
    float4 coef;

    coef = read_imagef(means, coord.yz);

    float4  tmpData0, tmpData1;

    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_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    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_DP2x8(src1, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                    uniConvBF16toF32_Part0_2x8);
    VXC_DP2x8(src2, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                    uniConvBF16toF32_Part1_2x8);
    _viv_asm(COPY, tmpData0, src1, 16);
    _viv_asm(COPY, tmpData1, src2, 16);

    float4 norm;
    norm = tmpData0 * coef.x + coef.y;
    _viv_asm(COPY, src0, norm, 16);
    norm = tmpData1 * coef.x + coef.y;
    _viv_asm(COPY, src1, norm, 16);
    VXC_DP2x8(src2, src0, src1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniExtractOddData_2x8);
    VXC_OP4_NoDest(img_store_3d, output, coord, src2, \
                    VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0));
    }
}

__kernel void instance_norm_BF16_F32toBF16_2D(
    __read_only  image2d_array_t input,
    __read_only  image2d_t       means,
    __write_only image2d_array_t output,
                 int height)
{
    int gidz = get_global_id(1);
    int gidy = gidz * height;
    int2 coord = (int2)(get_global_id(0), gidy);
    int2 coord_para = (int2)(0, gidz);
    int endH = gidy + height;
    vxc_short8 src0, src1, src2;
    float4 coef;

    coef = read_imagef(means, coord_para);

    float4  tmpData0, tmpData1;

    for(; coord.y < endH; coord.y++)
    {
    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
        VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    VXC_DP2x8(src1, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                    uniConvBF16toF32_Part0_2x8);
    VXC_DP2x8(src2, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),
                    uniConvBF16toF32_Part1_2x8);
    _viv_asm(COPY, tmpData0, src1, 16);
    _viv_asm(COPY, tmpData1, src2, 16);

    float4 norm;
    norm = tmpData0 * coef.x + coef.y;
    _viv_asm(COPY, src0, norm, 16);
    norm = tmpData1 * coef.x + coef.y;
    _viv_asm(COPY, src1, norm, 16);
    VXC_DP2x8(src2, src0, src1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniExtractOddData_2x8);
    VXC_WriteImage(output, coord.xy, src2, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
}