
#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniConvFP16toFP32_Lo_4x4;
_viv_uniform VXC_512Bits uniConvFP16toFP32_Hi_4x4;
_viv_uniform VXC_512Bits uniExtractHalf8_2x8;
_viv_uniform VXC_512Bits uniExtractInteger_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_Part0_2x8;
_viv_uniform VXC_512Bits uniConvBF16toF32_Part1_2x8;
_viv_uniform VXC_512Bits uniPackedBF16_2x8;
_viv_uniform VXC_512Bits uniConvIntegertoFP32_Lo_4x4;
_viv_uniform VXC_512Bits uniConvIntegertoFP32_Hi_4x4;
_viv_uniform float offset;
_viv_uniform float input_scale;
_viv_uniform float inputTail;
_viv_uniform float output_scale;
_viv_uniform float outputZP;

float4 I8toF32_Lo(vxc_char8 src)
{
    float4 dst;

    VXC_DP4x4(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvIntegertoFP32_Lo_4x4);
    dst *= input_scale;
    return dst;
}

float4 I8toF32_Hi(vxc_char8 src)
{
    float4 dst;

    VXC_DP4x4(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvIntegertoFP32_Hi_4x4);
    dst *= input_scale;
    return dst;
}

float4 U8toF32_Lo(vxc_uchar8 src)
{
    float4 dst;

    VXC_DP4x4(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvIntegertoFP32_Lo_4x4);
    dst = dst * input_scale + inputTail;
    return dst;
}

float4 U8toF32_Hi(vxc_uchar8 src)
{
    float4 dst;

    VXC_DP4x4(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvIntegertoFP32_Hi_4x4);
    dst = dst * input_scale + inputTail;
    return dst;
}

float4 I16toF32_Lo(vxc_short8 src)
{
    float4 dst;

    VXC_DP4x4(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvIntegertoFP32_Lo_4x4);
    dst *= input_scale;
    return dst;
}

float4 I16toF32_Hi(vxc_short8 src)
{
    float4 dst;

    VXC_DP4x4(dst, src, src, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvIntegertoFP32_Hi_4x4);
    dst *= input_scale;
    return dst;
}

float4 F16toF32_Lo(vxc_ushort8 src)
{
    vxc_half8 srcHalf;
    float4 dst;

    _viv_asm(COPY, srcHalf, src, 16);
    VXC_DP4x4(dst, srcHalf, srcHalf, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvFP16toFP32_Lo_4x4);
    return dst;
}

float4 F16toF32_Hi(vxc_ushort8 src)
{
    vxc_half8 srcHalf;
    float4 dst;

    _viv_asm(COPY, srcHalf, src, 16);
    VXC_DP4x4(dst, srcHalf, srcHalf, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvFP16toFP32_Hi_4x4);
    return dst;
}

float4 BF16toF32_Lo(vxc_ushort8 src)
{
    vxc_ushort8 srcA;
    float4 dst;

    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    VXC_DP2x8(srcA, src, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part0_2x8);
    _viv_asm(COPY, dst, srcA, 16);

    return dst;
}

float4 BF16toF32_Hi(vxc_ushort8 src)
{
    vxc_ushort8 srcA;
    float4 dst;

    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    VXC_DP2x8(srcA, src, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part1_2x8);
    _viv_asm(COPY, dst, srcA, 16);

    return dst;
}

vxc_ushort8 F32toF16(float4 src0, float4 src1)
{
    half4 srcHalf0, srcHalf1;
    vxc_half8 dst0;
    vxc_ushort8 dst;

    _viv_asm(CONV, srcHalf0, src0);
    _viv_asm(CONV, srcHalf1, src1);

    VXC_DP2x8(dst0, srcHalf0, srcHalf1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniExtractHalf8_2x8);
    _viv_asm(COPY, dst, dst0, 16);
    return dst;
}

vxc_ushort8 F32toBF16(float4 src0, float4 src1)
{
    vxc_ushort8 srcA, srcB;
    vxc_ushort8 dst;

    _viv_asm(COPY, srcA, src0, 16);
    _viv_asm(COPY, srcB, src1, 16);
    VXC_DP2x8(dst, srcA, srcB, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniPackedBF16_2x8);
    return dst;
}

vxc_char8 F32toI8(float4 src0, float4 src1)
{
    int4 srcInt0, srcInt1;
    vxc_char8 dst;

    src0 *= output_scale;
    src1 *= output_scale;
    _viv_asm(CONV_RTE, srcInt0, src0);
    _viv_asm(CONV_RTE, srcInt1, src1);

    VXC_DP2x8(dst, srcInt0, srcInt1, VXC_MODIFIER(0, 7, 0, VXC_RM_ToNearestEven, 1), uniExtractInteger_2x8);
    return dst;
}

vxc_short8 F32toI16(float4 src0, float4 src1)
{
    int4 srcInt0, srcInt1;
    vxc_short8 dst;

    src0 *= output_scale;
    src1 *= output_scale;
    _viv_asm(CONV_RTE, srcInt0, src0);
    _viv_asm(CONV_RTE, srcInt1, src1);

    VXC_DP2x8(dst, srcInt0, srcInt1, VXC_MODIFIER(0, 7, 0, VXC_RM_ToNearestEven, 1), uniExtractInteger_2x8);
    return dst;
}

vxc_uchar8 F32toU8(float4 src0, float4 src1)
{
    int4 srcInt0, srcInt1;
    vxc_uchar8 dst;

    src0 = src0 * output_scale + outputZP;
    src1 = src1 * output_scale + outputZP;
    _viv_asm(CONV_RTE, srcInt0, src0);
    _viv_asm(CONV_RTE, srcInt1, src1);

    VXC_DP2x8(dst, srcInt0, srcInt1, VXC_MODIFIER(0, 7, 0, VXC_RM_ToNearestEven, 1), uniExtractInteger_2x8);
    return dst;
}


#define TENSOR_KERAS_RELU(src_type_name, dst_type_name, tensor_dims, image_type, \
    convert2FP32_Func, convert2DstType_Func, src_type, dst_type) \
__kernel void relu_keras_##src_type_name##to##dst_type_name##tensor_dims( \
__read_only  image2d_array_t  input, \
__write_only image2d_array_t  output, \
                    float     alpha, \
                    float     max_value, \
                    float     threshold \
                    ) \
{ \
    int4 coord = (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0); \
    src_type src; \
    VXC_Read##image_type(src, input, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    float4 dataA = convert2FP32_Func##_Lo(src); \
    float4 dataB = convert2FP32_Func##_Hi(src); \
    float4 dataA0 = dataA < threshold ? threshold : dataA; \
    dataA0 = dataA0 > max_value ? max_value : dataA0; \
    float4 dataB0 = dataB < threshold ? threshold : dataB; \
    dataB0 = dataB0 > max_value ? max_value : dataB0; \
    float4 dataA1 = dataA * alpha - offset; \
    float4 dataB1 = dataB * alpha - offset; \
    float4 dst0 = dataA  < threshold ? dataA1 : dataA0; \
    float4 dst1 = dataB  < threshold ? dataB1 : dataB0; \
    dst_type result = convert2DstType_Func(dst0, dst1); \
    VXC_Write##image_type(output, coord, result, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}

TENSOR_KERAS_RELU(F16,  F16,  _3D, Image2DArray, F16toF32,  F32toF16,  vxc_ushort8, vxc_ushort8)
TENSOR_KERAS_RELU(F16,  I16,  _3D, Image2DArray, F16toF32,  F32toI16,  vxc_ushort8, vxc_short8)
TENSOR_KERAS_RELU(F16,  I8,   _3D, Image2DArray, F16toF32,  F32toI8,   vxc_ushort8, vxc_char8)
TENSOR_KERAS_RELU(F16,  U8,   _3D, Image2DArray, F16toF32,  F32toU8,   vxc_ushort8, vxc_uchar8)
TENSOR_KERAS_RELU(BF16, BF16, _3D, Image2DArray, BF16toF32, F32toBF16, vxc_ushort8, vxc_ushort8)

TENSOR_KERAS_RELU(I16,  I16,  _3D, Image2DArray, I16toF32,  F32toI16,  vxc_short8,  vxc_short8)
TENSOR_KERAS_RELU(I16,  F16,  _3D, Image2DArray, I16toF32,  F32toF16,  vxc_short8,  vxc_ushort8)
TENSOR_KERAS_RELU(I8,   I8,   _3D, Image2DArray, I8toF32,   F32toI8,   vxc_char8,   vxc_char8)
TENSOR_KERAS_RELU(I8,   F16,  _3D, Image2DArray, I8toF32,   F32toF16,  vxc_char8,   vxc_ushort8)
TENSOR_KERAS_RELU(U8,   U8,   _3D, Image2DArray, U8toF32,   F32toU8,   vxc_uchar8,  vxc_uchar8)
TENSOR_KERAS_RELU(U8,   F16,  _3D, Image2DArray, U8toF32,   F32toF16,  vxc_uchar8,  vxc_ushort8)

TENSOR_KERAS_RELU(F16,  F16,  _2D, Image,        F16toF32,  F32toF16,  vxc_ushort8, vxc_ushort8)
TENSOR_KERAS_RELU(F16,  I16,  _2D, Image,        F16toF32,  F32toI16,  vxc_ushort8, vxc_short8)
TENSOR_KERAS_RELU(F16,  I8,   _2D, Image,        F16toF32,  F32toI8,   vxc_ushort8, vxc_char8)
TENSOR_KERAS_RELU(F16,  U8,   _2D, Image,        F16toF32,  F32toU8,   vxc_ushort8, vxc_uchar8)
TENSOR_KERAS_RELU(BF16, BF16, _2D, Image,        BF16toF32, F32toBF16, vxc_ushort8, vxc_ushort8)
TENSOR_KERAS_RELU(I16,  I16,  _2D, Image,        I16toF32,  F32toI16,  vxc_short8,  vxc_short8)
TENSOR_KERAS_RELU(I16,  F16,  _2D, Image,        I16toF32,  F32toF16,  vxc_short8,  vxc_ushort8)
TENSOR_KERAS_RELU(I8,   I8,   _2D, Image,        I8toF32,   F32toI8,   vxc_char8,   vxc_char8)
TENSOR_KERAS_RELU(I8,   F16,  _2D, Image,        I8toF32,   F32toF16,  vxc_char8,   vxc_ushort8)
TENSOR_KERAS_RELU(U8,   U8,   _2D, Image,        U8toF32,   F32toU8,   vxc_uchar8,  vxc_uchar8)
TENSOR_KERAS_RELU(U8,   F16,  _2D, Image,        U8toF32,   F32toF16,  vxc_uchar8,  vxc_ushort8)
