#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniExtract8Data_2x8;

_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;

#define INSTANCE_NORM_8_TO_F16_IMPL(name, src_type) \
__kernel void instance_norm_##name( \
    __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_para = (int4)(0, gidz, 0, 0); \
    src_type src0; \
    vxc_short8 outval; \
    vxc_half8 dst; \
    float4 coef; \
 \
    coef = read_imagef(means, coord_para.xy); \
    float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    half4 tmpVal0, tmpVal1; \
 \
    coord_para = coord; \
    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.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_para.z, baseAddr); \
    for(coord.y = 0; coord.y < height;) \
    { \
    VXC_OP4(img_load_3d, src0, input, coord, 0, \
                    VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
    coord_para.xy = coord.xy; \
    coord.y++; \
    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); \
    VXC_DP4x4(tmpData2, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_2_4x4); \
    VXC_DP4x4(tmpData3, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_3_4x4); \
    norm = tmpData0 * coef.x + coef.y; \
    _viv_asm(CONV, tmpVal0, norm); \
    norm = tmpData1 * coef.x + coef.y; \
    _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_para, outval, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0)); \
    coord_para.x += 8; \
    norm = tmpData2 * coef.x + coef.y; \
    _viv_asm(CONV, tmpVal0, norm); \
    norm = tmpData3 * coef.x + coef.y; \
    _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_para, outval, VXC_MODIFIER(0, 7, 0,VXC_RM_TowardZero, 0)); \
    } \
}
INSTANCE_NORM_8_TO_F16_IMPL(U8_F32toF16, vxc_uchar16)
INSTANCE_NORM_8_TO_F16_IMPL(I8_F32toF16, vxc_char16)

#define INSTANCE_NORM_8_TO_F16_IMPL_2D(name, src_type) \
__kernel void instance_norm_##name##_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; \
    int4 coord = (int4)(get_global_id(0), gidy, 0, 0); \
    int4 coord_para = (int4)(0, gidz, 0, 0); \
    int endH = gidy + height; \
    src_type src0; \
    vxc_short8 outval; \
    vxc_half8 dst; \
    float4 coef; \
 \
    coef = read_imagef(means, coord_para.xy); \
    float4  tmpData0, tmpData1, tmpData2, tmpData3, norm; \
    half4 tmpVal0, tmpVal1; \
    for(; coord.y < endH;) \
    { \
    VXC_ReadImage(src0, input, coord.xy, 0, VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0)); \
    coord_para = coord; \
    coord.y++; \
    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); \
    VXC_DP4x4(tmpData2, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_2_4x4); \
    VXC_DP4x4(tmpData3, src0, src0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDataToFP32_3_4x4); \
    norm = tmpData0 * coef.x + coef.y; \
    _viv_asm(CONV, tmpVal0, norm); \
    norm = tmpData1 * coef.x + coef.y; \
    _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_WriteImage(output, coord_para.xy, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    coord_para.x += 8; \
    norm = tmpData2 * coef.x + coef.y; \
    _viv_asm(CONV, tmpVal0, norm); \
    norm = tmpData3 * coef.x + coef.y; \
    _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_WriteImage(output, coord_para.xy, outval, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    } \
}
INSTANCE_NORM_8_TO_F16_IMPL_2D(U8_F32toF16, vxc_uchar16)
INSTANCE_NORM_8_TO_F16_IMPL_2D(I8_F32toF16, vxc_char16)
