#include "cl_viv_vx_ext.h"

_viv_uniform int inputWidth;
_viv_uniform float output_ZP;
_viv_uniform float zP2x;
_viv_uniform int inputZP;

_viv_uniform float inOutScale;
_viv_uniform float e2InScale;
_viv_uniform float zpSqr8x;
_viv_uniform VXC_512Bits uniInt16SumSqr_dp8x2;
_viv_uniform VXC_512Bits UniFP16toFP32Lo4_dp4x4;
_viv_uniform VXC_512Bits uniConvertSecFp16Fp32_4x4;
_viv_uniform VXC_512Bits uniConvert1stUint8SubZpToFp32_4x4;
_viv_uniform VXC_512Bits uniConvert2ndUint8SubZpToFp32_4x4;
_viv_uniform VXC_512Bits uniConvertInt32toUint8_2x8;
_viv_uniform VXC_512Bits uniConvertHalfToFp16_2x8;

_viv_uniform VXC_512Bits uniFp16SumSqr_dp8x2;

__kernel __attribute__((reqd_work_group_size(16, 1, 1)))
     void l2normalizescale_axis0_F16_F16toF16_2D(
    __read_only image2d_t input, __read_only image2d_t scale, __write_only image2d_t output, int axis)
{
    int4 coord = (int4)(get_global_id(0), get_global_id(1), 0, 0);
    int lidx = get_local_id(0);
    vxc_short8 src0, src1, dst;
    vxc_half8 in_h, scale_h, tmpDst;
    float sum = 0;
    vxc_float4 scale_f0, scale_f1, sumsqr, tmpData0, tmpData1;
    __local float lcl_sum[16];
    float4 one = (float4)(1, 1, 1, 1);
    for(; coord.x < inputWidth;)
    {
        VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord.x += 128;
        _viv_asm(COPY, in_h, src0, 16);
        VXC_DP8x2(sumsqr, in_h, in_h, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), \
                    uniFp16SumSqr_dp8x2);
        sum += sumsqr.y;
    }
    lcl_sum[lidx] = sum;
    barrier(CLK_LOCAL_MEM_FENCE);
    float4 *pLocalPtr = (float4 *)&lcl_sum[0];
    float4 data0;
    data0 = pLocalPtr[0] + pLocalPtr[1] + pLocalPtr[2] + pLocalPtr[3];
    sum = dot(data0, one);
    float alpha = rsqrt(sum);

    for(coord.x = get_global_id(0); coord.x < inputWidth; coord.x += 128)
    {
        VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_ReadImage(src1, scale, coord.xz, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

        _viv_asm(COPY, in_h, src0, 16);
        _viv_asm(COPY, scale_h, src1, 16);
        VXC_DP4x4(tmpData0, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
        VXC_DP4x4(tmpData1, in_h, in_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertSecFp16Fp32_4x4);

        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4);
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertSecFp16Fp32_4x4);

        half4 tmpVal0, tmpVal1;
        tmpData0 *= scale_f0 * alpha;
        tmpData1 *= scale_f1 * alpha;
        _viv_asm(CONV, tmpVal0, tmpData0);
        _viv_asm(CONV, tmpVal1, tmpData1);
        VXC_DP2x8(tmpDst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),\
            uniConvertHalfToFp16_2x8);
        _viv_asm(COPY, dst, tmpDst, 16);
        VXC_WriteImage(output, coord.xy, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
}


#define L2NORMSCALE_QINTF16TOQINT_AXIS0_2D(in0_name, in1_name, out_name, read_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) \
     void l2normalizescale_axis0_##in0_name##_##in1_name##to##out_name##_2D( \
    __read_only image2d_t input, __read_only image2d_t scale, __write_only image2d_t output, int axis)\
{ \
    int4 coord = (int4)(get_global_id(0), get_global_id(1), 0, 0); \
    int lidx = get_local_id(0); \
    read_type src0, dst; \
    vxc_short8 src1; \
    vxc_half8 scale_h; \
    float sum = 0; \
    vxc_float4 scale_f0, scale_f1, sumsqr; \
    __local float lcl_sum[16]; \
    float4 one = (float4)(1, 1, 1, 1); \
    for(; coord.x < inputWidth;) \
    { \
        VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        coord.x += 128; \
        VXC_DP8x2(sumsqr, src0, src0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), \
                    uniInt16SumSqr_dp8x2); \
        sum += sumsqr.y - zP2x * sumsqr.x + zpSqr8x; \
    } \
    sum *= e2InScale; \
    lcl_sum[lidx] = sum; \
    barrier(CLK_LOCAL_MEM_FENCE); \
    float4 *pLocalPtr = (float4 *)&lcl_sum[0]; \
    float4 data0 = pLocalPtr[0] + pLocalPtr[1] + pLocalPtr[2] + pLocalPtr[3]; \
    sum = dot(data0, one); \
    float alpha = rsqrt(sum) * inOutScale; \
    short zp = inputZP; \
    vxc_float4  tmpData0, tmpData1; \
    for(coord.x = get_global_id(0); coord.x < inputWidth; coord.x += 128) \
    { \
        VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        VXC_ReadImage(src1, scale, coord.xz, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        _viv_asm(COPY, scale_h, src1, 16); \
        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4); \
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertSecFp16Fp32_4x4); \
 \
        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); \
 \
        int4 tmpVal0 = convert_int4_rte(tmpData0 * scale_f0 * alpha + output_ZP); \
        int4 tmpVal1 = convert_int4_rte(tmpData1 * scale_f1 * alpha + output_ZP); \
        VXC_DP2x8(dst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
                uniConvertInt32toUint8_2x8); \
        VXC_WriteImage(output, coord.xy, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    } \
}

L2NORMSCALE_QINTF16TOQINT_AXIS0_2D(U8,  F16,  U8,   vxc_uchar8)
L2NORMSCALE_QINTF16TOQINT_AXIS0_2D(I8,  F16,  I8,   vxc_char8)
L2NORMSCALE_QINTF16TOQINT_AXIS0_2D(I16, F16,  I16,  vxc_short8)

#define L2NORMSCALE_QINTF16TOF16_AXIS0_2D(in0_name, in1_name, out_name, read_type) \
__kernel __attribute__((reqd_work_group_size(16, 1, 1))) \
     void l2normalizescale_axis0_##in0_name##_##in1_name##to##out_name##_2D( \
    __read_only image2d_t input, __read_only image2d_t scale, __write_only image2d_t output, int axis)\
{ \
    int4 coord = (int4)(get_global_id(0), get_global_id(1), 0, 0); \
    int lidx = get_local_id(0); \
    read_type src0; \
    vxc_short8 src1, dst; \
    vxc_half8 scale_h, tmpDst; \
    float sum = 0; \
    vxc_float4 scale_f0, scale_f1, sumsqr, tmpData0, tmpData1; \
    __local float lcl_sum[16]; \
    float4 one = (float4)(1, 1, 1, 1); \
    for(; coord.x < inputWidth;) \
    { \
        VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        coord.x += 128; \
        VXC_DP8x2(sumsqr, src0, src0, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), \
                    uniInt16SumSqr_dp8x2); \
        sum += sumsqr.y - zP2x * sumsqr.x + zpSqr8x; \
    } \
    sum *= e2InScale; \
    lcl_sum[lidx] = sum; \
    barrier(CLK_LOCAL_MEM_FENCE); \
    float4 *pLocalPtr = (float4 *)&lcl_sum[0]; \
    float4 data0; \
    data0 = pLocalPtr[0] + pLocalPtr[1] + pLocalPtr[2] + pLocalPtr[3]; \
    sum = dot(data0, one); \
    float alpha = rsqrt(sum) * inOutScale; \
    short zp = inputZP; \
    for(coord.x = get_global_id(0); coord.x < inputWidth; coord.x += 128) \
    { \
        VXC_ReadImage(src0, input, coord.xy, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        VXC_ReadImage(src1, scale, coord.xz, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
        _viv_asm(COPY, scale_h, src1, 16); \
        VXC_DP4x4(scale_f0, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            UniFP16toFP32Lo4_dp4x4); \
        VXC_DP4x4(scale_f1, scale_h, scale_h, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertSecFp16Fp32_4x4); \
 \
        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); \
 \
        half4 tmpVal0, tmpVal1; \
        tmpData0 *= scale_f0 * alpha; \
        tmpData1 *= scale_f1 * alpha; \
        _viv_asm(CONV, tmpVal0, tmpData0); \
        _viv_asm(CONV, tmpVal1, tmpData1); \
        VXC_DP2x8(tmpDst, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0),\
            uniConvertHalfToFp16_2x8); \
        _viv_asm(COPY, dst, tmpDst, 16); \
        VXC_WriteImage(output, coord.xy, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    } \
}

L2NORMSCALE_QINTF16TOF16_AXIS0_2D(U8,  F16,  F16,   vxc_uchar8)
L2NORMSCALE_QINTF16TOF16_AXIS0_2D(I8,  F16,  F16,   vxc_char8)
L2NORMSCALE_QINTF16TOF16_AXIS0_2D(I16, F16,  F16,   vxc_short8)
