#include "cl_viv_vx_ext.h"

_viv_uniform float inputScale;
_viv_uniform float inputTail;
_viv_uniform float outputScale;
_viv_uniform float outputZP;
_viv_uniform VXC_512Bits uniExtract8Data_2x8;
_viv_uniform VXC_512Bits uniDatatoFp32Part0_4x4;
_viv_uniform VXC_512Bits uniDatatoFp32Part1_4x4;

#define HSWISH_PROCESS(read_fun, write_fun, src_type, src_copy_type, convert_type, dst_type, dst_copy_type, \
                     INSCALE, INTAIL, OUTSCALE, OUTZP) \
    src_type      src0; \
    src_copy_type src1; \
    read_fun(src0, input, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    _viv_asm(COPY, src1, src0, 16); \
    float4 vecA, vecB, vecC, vecD, vecE, vecDstA, vecDstB; \
    VXC_DP4x4(vecA, src1, src1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDatatoFp32Part0_4x4); \
    VXC_DP4x4(vecB, src1, src1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniDatatoFp32Part1_4x4); \
    vecA = vecA * INSCALE + INTAIL; \
    vecB = vecB * INSCALE + INTAIL; \
    vecC = vecA + 3.0f; \
    vecD = vecB + 3.0f; \
    vecE = 6.0f; \
    _viv_asm(CLAMP0MAX, vecDstA, vecC, vecE); \
    _viv_asm(CLAMP0MAX, vecDstB, vecD, vecE); \
    vecA = vecA * vecDstA; \
    vecB = vecB * vecDstB; \
    vecA = vecA / 6.0f; \
    vecB = vecB / 6.0f; \
    vecA = vecA * OUTSCALE + OUTZP; \
    vecB = vecB * OUTSCALE + OUTZP; \
    convert_type dst0, dst1; \
    _viv_asm(CONV_RTE, dst0, vecA); \
    _viv_asm(CONV_RTE, dst1, vecB); \
    dst_type dst2; \
    VXC_DP2x8(dst2, dst0, dst1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniExtract8Data_2x8); \
    dst_copy_type dst; \
    _viv_asm(COPY, dst, dst2, 16); \
    write_fun(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));


#define HSWISH_FUNC(src_type_name, dst_type_name, src_type, src_copy_type, convert_type, dst_type, \
                   dst_copy_type, INSCALE, INTAIL, OUTSCALE, OUTZP) \
    __kernel void hswish_##src_type_name##to##dst_type_name( \
    __read_only  image2d_array_t  input, \
    __write_only image2d_array_t  output, \
                           float  beta \
    ) \
{ \
    int4 coord = (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0); \
    HSWISH_PROCESS(VXC_ReadImage2DArray, VXC_WriteImage2DArray, src_type, \
                 src_copy_type, convert_type, dst_type, dst_copy_type, \
                 INSCALE, INTAIL, OUTSCALE, OUTZP) \
}

HSWISH_FUNC(F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8, 1, 0, 1, 0)
HSWISH_FUNC(F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8,  1, 0, outputScale, 0)
HSWISH_FUNC(F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8, 1, 0, outputScale, outputZP)
HSWISH_FUNC(F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8, 1, 0, outputScale, 0)
HSWISH_FUNC(I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8,  inputScale, 0, outputScale, 0)
HSWISH_FUNC(I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8, inputScale, 0, 1, 0)
HSWISH_FUNC(U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8, \
            inputScale, inputTail, outputScale, outputZP)
HSWISH_FUNC(U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8, inputScale, inputTail, 1, 0)
HSWISH_FUNC(I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8, inputScale, 0, outputScale, 0)
HSWISH_FUNC(I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8, inputScale, 0, 1, 0)


#define HSWISH_FUNC_2D(src_type_name, dst_type_name, src_type, src_copy_type, convert_type, dst_type, \
                      dst_copy_type, INSCALE, INTAIL, OUTSCALE, OUTZP) \
    __kernel void hswish_##src_type_name##to##dst_type_name##_2D( \
    __read_only  image2d_array_t  input, \
    __write_only image2d_array_t  output, \
                           float  beta \
    ) \
{ \
    int2 coord = (int2)(get_global_id(0), get_global_id(1)); \
    HSWISH_PROCESS(VXC_ReadImage, VXC_WriteImage, src_type, src_copy_type, convert_type, dst_type, \
                   dst_copy_type, INSCALE, INTAIL, OUTSCALE, OUTZP) \
}

HSWISH_FUNC_2D(F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8, 1, 0, 1, 0)
HSWISH_FUNC_2D(F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8,  1, 0, outputScale, 0)
HSWISH_FUNC_2D(F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8, 1, 0, outputScale, outputZP)
HSWISH_FUNC_2D(F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8, 1, 0, outputScale, 0)
HSWISH_FUNC_2D(I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8,  inputScale, 0, outputScale, 0)
HSWISH_FUNC_2D(I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8, inputScale, 0, 1, 0)
HSWISH_FUNC_2D(U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8, inputScale, \
               inputTail, outputScale, outputZP)
HSWISH_FUNC_2D(U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8, inputScale, inputTail, 1, 0)
HSWISH_FUNC_2D(I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8, inputScale, 0, outputScale, 0)
HSWISH_FUNC_2D(I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8, inputScale, 0, 1, 0)


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

#define HSWISH_BF16_PROCESS(read_fun, write_fun) \
    vxc_ushort8   src0, src1, dst; \
    float4 vecA, vecB, vecC, vecD, vecE, vecDstA, vecDstB; \
    read_fun(src0, input, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0); \
    VXC_DP2x8(src1, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part0_2x8); \
    _viv_asm(COPY, vecA, src1, 16); \
    VXC_DP2x8(src1, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part1_2x8); \
    _viv_asm(COPY, vecB, src1, 16); \
    vecC = vecA + 3.0f; \
    vecD = vecB + 3.0f; \
    vecE = 6.0f; \
    _viv_asm(CLAMP0MAX, vecDstA, vecC, vecE); \
    _viv_asm(CLAMP0MAX, vecDstB, vecD, vecE); \
    vecA = vecA * vecDstA; \
    vecB = vecB * vecDstB; \
    vecA = vecA / 6.0f; \
    vecB = vecB / 6.0f; \
    _viv_asm(COPY, src0, vecA, 16); \
    _viv_asm(COPY, src1, vecB, 16); \
    VXC_DP2x8(dst, src0, src1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniExtractOddData_2x8); \
    write_fun(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

__kernel void hswish_BF16toBF16(
    __read_only  image2d_array_t  input,
    __write_only image2d_array_t  output,
                           float  beta
    )
{
    int4 coord = (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0);
    HSWISH_BF16_PROCESS(VXC_ReadImage2DArray, VXC_WriteImage2DArray);
}

__kernel void hswish_BF16toBF16_2D(
    __read_only  image2d_array_t  input,
    __write_only image2d_array_t  output,
                           float  beta
    )
{
    int2 coord = (int2)(get_global_id(0), get_global_id(1));
    HSWISH_BF16_PROCESS(VXC_ReadImage, VXC_WriteImage);
}
