#include "cl_viv_vx_ext.h"

_viv_uniform float alpha;

float4 eltwise_unary_sin(float4 x)
{
    return native_sin(x);
}

#define logE        (1.44269502f)
#define twoLogE     (logE * 2.0f)
float4 eltwise_unary_exp(float4 x)
{
    x *= logE;
    x = exp2(x);
    return x;
}

#define rlogE    (0.693147182f)
float4 eltwise_unary_log(float4 x)
{
    x = log2(x);
    return x * rlogE;
}

float4 eltwise_unary_elu(float4 val)
{
    float4 x = val * logE;
    x = exp2(x) * alpha - alpha;

    return val < 0 ? x : val;
}

float4 eltwise_unary_neg(float4 x)
{
    return x * -1;
}

float4 eltwise_unary_hard_sigmoid(float4 x)
{
    x = 0.2 * x + 0.5;
    x = clamp(x, 0, 1);
    return x;
}

float4 _softrelu(float4 x)
{
    x *= logE;
    x = exp2(x);
    x += 1;
    x = log2(x);
    return x * rlogE;
}

float4 _tanh(float4 x)
{
    x *= -twoLogE;
    x = 1 + exp2(x);
    x = 1 / x;
    return (2 * x - 1);
}

float4 eltwise_unary_mish(float4 x)
{
    float4 y = _softrelu(x);
    x = x * _tanh(y);
    return x;
}

_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 ELTSISE_UNARY_3D(func_name, src_type_name, dst_type_name, src_type, \
                src_copy_type, convert_type, dst_type, dst_copy_type) \
__kernel void func_name##_##src_type_name##to##dst_type_name( \
    __read_only  image2d_array_t  input, \
    __write_only image2d_array_t  output, \
                 int              type, \
                 float            _alpha \
    ) \
{ \
    int4 coord = (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0); \
    src_type      src0; \
    src_copy_type src1; \
    VXC_ReadImage2DArray(src0, input, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    _viv_asm(COPY, src1, src0, 16); \
 \
    float4 vecA; \
    float4 vecB; \
    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 * inputScale + inputTail; \
    vecB = vecB * inputScale + inputTail; \
    vecA = eltwise_unary_##func_name(vecA); \
    vecB = eltwise_unary_##func_name(vecB); \
    vecA = vecA * outputScale + outputZP; \
    vecB = vecB * outputScale + outputZP; \
 \
    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); \
    VXC_WriteImage2DArray(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}
//EXP
ELTSISE_UNARY_3D(exp, F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(exp, F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(exp, F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(exp, F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(exp, I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(exp, I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(exp, U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(exp, U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(exp, I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(exp, I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8)
//SIN
ELTSISE_UNARY_3D(sin, F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(sin, F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(sin, F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(sin, F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(sin, I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(sin, I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(sin, U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(sin, U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(sin, I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(sin, I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8)
//LOG
ELTSISE_UNARY_3D(log, F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(log, F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(log, F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(log, F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(log, I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(log, I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(log, U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(log, U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(log, I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(log, I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8)
//ELU
ELTSISE_UNARY_3D(elu, F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(elu, F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(elu, F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(elu, F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(elu, I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(elu, I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(elu, U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(elu, U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(elu, I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(elu, I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8)
//NEG
ELTSISE_UNARY_3D(neg, F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(neg, F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(neg, F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(neg, F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(neg, I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(neg, I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(neg, U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(neg, U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(neg, I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(neg, I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8)
//MISH
ELTSISE_UNARY_3D(mish, F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(mish, F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(mish, F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(mish, F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(mish, I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(mish, I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(mish, U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(mish, U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(mish, I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(mish, I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8)
//HARD_SIGMOID
ELTSISE_UNARY_3D(hard_sigmoid, F16, F16, vxc_short8, vxc_half8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(hard_sigmoid, F16, I8,  vxc_short8, vxc_half8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(hard_sigmoid, F16, U8,  vxc_short8, vxc_half8,  int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(hard_sigmoid, F16, I16, vxc_short8, vxc_half8,  int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(hard_sigmoid, I8,  I8,  vxc_char8,  vxc_char8,  int4,  vxc_char8,  vxc_char8)
ELTSISE_UNARY_3D(hard_sigmoid, I8,  F16, vxc_char8,  vxc_char8,  half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(hard_sigmoid, U8,  U8,  vxc_uchar8, vxc_uchar8, int4,  vxc_uchar8, vxc_uchar8)
ELTSISE_UNARY_3D(hard_sigmoid, U8,  F16, vxc_uchar8, vxc_uchar8, half4, vxc_half8,  vxc_short8)
ELTSISE_UNARY_3D(hard_sigmoid, I16, I16, vxc_short8, vxc_short8, int4,  vxc_short8, vxc_short8)
ELTSISE_UNARY_3D(hard_sigmoid, I16, F16, vxc_short8, vxc_short8, half4, vxc_half8,  vxc_short8)

_viv_uniform VXC_512Bits uniConvBF16toF32_Part0_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_Part1_2x8;
_viv_uniform VXC_512Bits uniExtractOddData_2x8;
#define ELTSISE_UNARY_BF16(func_name) \
    __kernel void func_name##_BF16toBF16( \
    __read_only  image2d_array_t  input, \
    __write_only image2d_array_t  output, \
                 int              type, \
                 float            _alpha \
    ) \
{ \
    int4 coord = (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0); \
    vxc_ushort8   src0, src1, dst; \
    VXC_ReadImage2DArray(src0, input, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
 \
    float4 vecA; \
    float4 vecB; \
    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); \
    vecA = eltwise_unary_##func_name(vecA); \
    vecB = eltwise_unary_##func_name(vecB); \
 \
    _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); \
    VXC_WriteImage2DArray(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}
//EXP
ELTSISE_UNARY_BF16(exp)
//SIN
ELTSISE_UNARY_BF16(sin)
//LOG
ELTSISE_UNARY_BF16(log)
//ELU
ELTSISE_UNARY_BF16(elu)
//NEG
ELTSISE_UNARY_BF16(neg)
//MISH
ELTSISE_UNARY_BF16(mish)
//HARD_SIGMOID
ELTSISE_UNARY_BF16(hard_sigmoid)