#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniConvertFstDataToFp32_4x4;
_viv_uniform VXC_512Bits uniConvertSecDataToFp32_4x4;
_viv_uniform VXC_512Bits uniExtact8Bit_2x8;
_viv_uniform float input0_scale;
_viv_uniform float input1_scale;
_viv_uniform float input0_tail;
_viv_uniform float input1_tail;
_viv_uniform float output_scale;
_viv_uniform float output_zp;

#define POW_SH_IMPL(name, src0_type, copy0_type, src1_type, copy1_type, dst_type, copy2_type, conv_type) \
__kernel void pow_##name \
    ( \
    __read_only  image2d_array_t input0, \
    __read_only  image2d_array_t input1, \
    __write_only image2d_array_t output \
    ) \
{ \
    int4 coord = (int4)(get_global_id(0), get_global_id(1), get_global_id(2), get_global_id(2)); \
 \
    src0_type src0; \
    copy0_type data0; \
    src1_type src1; \
    copy1_type data1; \
    VXC_ReadImage2DArray(src0, input0, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    _viv_asm(COPY, data0, src0, 16); \
    VXC_ReadImage2DArray(src1, input1, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    _viv_asm(COPY, data1, src1, 16); \
    float4 x0, x1; \
    float4 y0, y1; \
    VXC_DP4x4(x0, data0, data0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertFstDataToFp32_4x4); \
    VXC_DP4x4(x1, data0, data0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertSecDataToFp32_4x4); \
    VXC_DP4x4(y0, data1, data1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertFstDataToFp32_4x4); \
    VXC_DP4x4(y1, data1, data1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertSecDataToFp32_4x4); \
    x0 = x0 * input0_scale + input0_tail; \
    x1 = x1 * input0_scale + input0_tail; \
    y0 = y0 * input1_scale + input1_tail; \
    y1 = y1 * input1_scale + input1_tail; \
    float4 s0 = sign(x0); \
    float4 s1 = sign(x1); \
    int4 t0 = convert_int4(y0) & 1; \
    int4 t1 = convert_int4(y1) & 1; \
    s0 = s0 == -1 ? convert_float4(t0) == 1.0f ? -1.0f : 1.0f : s0; \
    s1 = s1 == -1 ? convert_float4(t1) == 1.0f ? -1.0f : 1.0f : s1; \
    x0 = s0 * exp2(y0 * log2(fabs(x0))); \
    x1 = s1 * exp2(y1 * log2(fabs(x1))); \
    x0 = x0 * output_scale + output_zp; \
    x1 = x1 * output_scale + output_zp; \
 \
    conv_type tmpVal0, tmpVal1; \
    _viv_asm(CONV_RTE, tmpVal0, x0); \
    _viv_asm(CONV_RTE, tmpVal1, x1); \
    dst_type dst0; \
 \
    copy2_type dst; \
    VXC_DP2x8(dst0, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_ToNearestEven, 1), uniExtact8Bit_2x8); \
    _viv_asm(COPY, dst, dst0, 16); \
    VXC_WriteImage2DArray(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}
POW_SH_IMPL(F16_F16toF16, vxc_short8,  vxc_half8,   vxc_short8,  vxc_half8,   vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(F16_F16toI16, vxc_short8,  vxc_half8,   vxc_short8,  vxc_half8,   vxc_short8,  vxc_short8,  int4)
POW_SH_IMPL(F16_F16toI8,  vxc_short8,  vxc_half8,   vxc_short8,  vxc_half8,   vxc_char16,  vxc_char16,  int4)
POW_SH_IMPL(F16_F16toU8,  vxc_short8,  vxc_half8,   vxc_short8,  vxc_half8,   vxc_uchar16, vxc_uchar16, int4)
POW_SH_IMPL(F16_I16toF16, vxc_short8,  vxc_half8,   vxc_short8,  vxc_short8,  vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(F16_I16toI16, vxc_short8,  vxc_half8,   vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  int4)
POW_SH_IMPL(I16_F16toF16, vxc_short8,  vxc_short8,  vxc_short8,  vxc_half8,   vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(I16_F16toI16, vxc_short8,  vxc_short8,  vxc_short8,  vxc_half8,   vxc_short8,  vxc_short8,  int4)
POW_SH_IMPL(I16_I16toF16, vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(I16_I16toI16, vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  int4)
POW_SH_IMPL(F16_I8toF16,  vxc_short8,  vxc_half8,   vxc_char16,  vxc_char16,  vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(F16_I8toI8,   vxc_short8,  vxc_half8,   vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  int4)
POW_SH_IMPL(I8_F16toF16,  vxc_char16,  vxc_char16,  vxc_short8,  vxc_half8,   vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(I8_F16toI8,   vxc_char16,  vxc_char16,  vxc_short8,  vxc_half8,   vxc_char16,  vxc_char16,  int4)
POW_SH_IMPL(I8_I8toF16,   vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(I8_I8toI8,    vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  int4)
POW_SH_IMPL(F16_U8toF16,  vxc_short8,  vxc_half8,   vxc_uchar16, vxc_uchar16, vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(F16_U8toU8,   vxc_short8,  vxc_half8,   vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, int4)
POW_SH_IMPL(U8_F16toF16,  vxc_uchar16, vxc_uchar16, vxc_short8,  vxc_half8,   vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(U8_F16toU8,   vxc_uchar16, vxc_uchar16, vxc_short8,  vxc_half8,   vxc_uchar16, vxc_uchar16, int4)
POW_SH_IMPL(U8_U8toF16,   vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL(U8_U8toU8,    vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, int4)

#define POW_SH_IMPL_2D(name, src0_type, copy0_type, src1_type, copy1_type, dst_type, copy2_type, conv_type) \
__kernel void pow_##name##_2D \
    ( \
    __read_only  image2d_array_t input0, \
    __read_only  image2d_array_t input1, \
    __write_only image2d_array_t output \
    ) \
{ \
    int2 coord = (int2)(get_global_id(0), get_global_id(1)); \
 \
    src0_type src0; \
    copy0_type data0; \
    src1_type src1; \
    copy1_type data1; \
    VXC_ReadImage(src0, input0, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    _viv_asm(COPY, data0, src0, 16); \
    VXC_ReadImage(src1, input1, coord, 0, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
    _viv_asm(COPY, data1, src1, 16); \
    float4 x0, x1; \
    float4 y0, y1; \
    VXC_DP4x4(x0, data0, data0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertFstDataToFp32_4x4); \
    VXC_DP4x4(x1, data0, data0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertSecDataToFp32_4x4); \
    VXC_DP4x4(y0, data1, data1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertFstDataToFp32_4x4); \
    VXC_DP4x4(y1, data1, data1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertSecDataToFp32_4x4); \
    x0 = x0 * input0_scale + input0_tail; \
    x1 = x1 * input0_scale + input0_tail; \
    y0 = y0 * input1_scale + input1_tail; \
    y1 = y1 * input1_scale + input1_tail; \
    float4 s0 = sign(x0); \
    float4 s1 = sign(x1); \
    int4 t0 = convert_int4(y0) & 1; \
    int4 t1 = convert_int4(y1) & 1; \
    s0 = s0 == -1 ? convert_float4(t0) == 1.0f ? -1.0f : 1.0f : s0; \
    s1 = s1 == -1 ? convert_float4(t1) == 1.0f ? -1.0f : 1.0f : s1; \
    x0 = s0 * exp2(y0 * log2(fabs(x0))); \
    x1 = s1 * exp2(y1 * log2(fabs(x1))); \
    x0 = x0 * output_scale + output_zp; \
    x1 = x1 * output_scale + output_zp; \
 \
    conv_type tmpVal0, tmpVal1; \
    _viv_asm(CONV_RTE, tmpVal0, x0); \
    _viv_asm(CONV_RTE, tmpVal1, x1); \
    dst_type dst0; \
 \
    copy2_type dst; \
    VXC_DP2x8(dst0, tmpVal0, tmpVal1, VXC_MODIFIER(0, 7, 0, VXC_RM_ToNearestEven, 1), uniExtact8Bit_2x8); \
    _viv_asm(COPY, dst, dst0, 16); \
    VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0)); \
}
POW_SH_IMPL_2D(F16_F16toF16, vxc_short8,  vxc_half8,   vxc_short8,  vxc_half8,   vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(F16_F16toI16, vxc_short8,  vxc_half8,   vxc_short8,  vxc_half8,   vxc_short8,  vxc_short8,  int4)
POW_SH_IMPL_2D(F16_F16toI8,  vxc_short8,  vxc_half8,   vxc_short8,  vxc_half8,   vxc_char16,  vxc_char16,  int4)
POW_SH_IMPL_2D(F16_F16toU8,  vxc_short8,  vxc_half8,   vxc_short8,  vxc_half8,   vxc_uchar16, vxc_uchar16, int4)
POW_SH_IMPL_2D(F16_I16toF16, vxc_short8,  vxc_half8,   vxc_short8,  vxc_short8,  vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(F16_I16toI16, vxc_short8,  vxc_half8,   vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  int4)
POW_SH_IMPL_2D(I16_F16toF16, vxc_short8,  vxc_short8,  vxc_short8,  vxc_half8,   vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(I16_F16toI16, vxc_short8,  vxc_short8,  vxc_short8,  vxc_half8,   vxc_short8,  vxc_short8,  int4)
POW_SH_IMPL_2D(I16_I16toF16, vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(I16_I16toI16, vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  vxc_short8,  int4)
POW_SH_IMPL_2D(F16_I8toF16,  vxc_short8,  vxc_half8,   vxc_char16,  vxc_char16,  vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(F16_I8toI8,   vxc_short8,  vxc_half8,   vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  int4)
POW_SH_IMPL_2D(I8_F16toF16,  vxc_char16,  vxc_char16,  vxc_short8,  vxc_half8,   vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(I8_F16toI8,   vxc_char16,  vxc_char16,  vxc_short8,  vxc_half8,   vxc_char16,  vxc_char16,  int4)
POW_SH_IMPL_2D(I8_I8toF16,   vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(I8_I8toI8,    vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  vxc_char16,  int4)
POW_SH_IMPL_2D(F16_U8toF16,  vxc_short8,  vxc_half8,   vxc_uchar16, vxc_uchar16, vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(F16_U8toU8,   vxc_short8,  vxc_half8,   vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, int4)
POW_SH_IMPL_2D(U8_F16toF16,  vxc_uchar16, vxc_uchar16, vxc_short8,  vxc_half8,   vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(U8_F16toU8,   vxc_uchar16, vxc_uchar16, vxc_short8,  vxc_half8,   vxc_uchar16, vxc_uchar16, int4)
POW_SH_IMPL_2D(U8_U8toF16,   vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_half8,   vxc_short8,  half4)
POW_SH_IMPL_2D(U8_U8toU8,    vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, vxc_uchar16, int4)

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

__kernel void pow_BF16_BF16toBF16
    (
    __read_only  image2d_array_t input0,
    __read_only  image2d_array_t input1,
    __write_only image2d_array_t output
    )
{
    int4 coord = (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0);

    vxc_ushort8 src0, src1, dst, tmpData;
    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    VXC_ReadImage2DArray(src0, input0, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    VXC_ReadImage2DArray(src1, input1, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    float4 x0, x1;
    float4 y0, y1;
    float4 tmpDst0, tmpDst1;

    VXC_DP2x8(tmpData, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part0_2x8);
    _viv_asm(COPY, x0, tmpData, 16);
    VXC_DP2x8(tmpData, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part1_2x8);
    _viv_asm(COPY, x1, tmpData, 16);

    VXC_DP2x8(tmpData, src1, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part0_2x8);
    _viv_asm(COPY, y0, tmpData, 16);
    VXC_DP2x8(tmpData, src1, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part1_2x8);
    _viv_asm(COPY, y1, tmpData, 16);

    float4  s0 = sign(x0);
    float4  s1 = sign(x1);
    int4 t0 = convert_int4(y0) & 1;
    int4 t1 = convert_int4(y1) & 1;
    s0 = s0 == -1 ? convert_float4(t0) == 1.0f ? -1.0f : 1.0f : s0;
    s1 = s1 == -1 ? convert_float4(t1) == 1.0f ? -1.0f : 1.0f : s1;
    tmpDst0 = s0 * exp2(y0 * log2(fabs(x0)));
    tmpDst1 = s1 * exp2(y1 * log2(fabs(x1)));

    _viv_asm(COPY, src0, tmpDst0, 16);
    _viv_asm(COPY, src1, tmpDst1, 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));
}

__kernel void pow_BF16_BF16toBF16_2D
    (
    __read_only  image2d_array_t input0,
    __read_only  image2d_array_t input1,
    __write_only image2d_array_t output
    )
{
    int2 coord = (int2)(get_global_id(0), get_global_id(1));

    vxc_ushort8 src0, src1, dst, tmpData;
    vxc_short8 zero = (vxc_short8)(0, 0, 0, 0, 0, 0, 0, 0);
    VXC_ReadImage(src0, input0, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    VXC_ReadImage(src1, input1, coord, VXC_5BITOFFSET_XY(0, 0),
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    float4 x0, x1;
    float4 y0, y1;
    float4 tmpDst0, tmpDst1;

    VXC_DP2x8(tmpData, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part0_2x8);
    _viv_asm(COPY, x0, tmpData, 16);
    VXC_DP2x8(tmpData, src0, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part1_2x8);
    _viv_asm(COPY, x1, tmpData, 16);

    VXC_DP2x8(tmpData, src1, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part0_2x8);
    _viv_asm(COPY, y0, tmpData, 16);
    VXC_DP2x8(tmpData, src1, zero, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvBF16toF32_Part1_2x8);
    _viv_asm(COPY, y1, tmpData, 16);

    float4  s0 = sign(x0);
    float4  s1 = sign(x1);
    int4 t0 = convert_int4(y0) & 1;
    int4 t1 = convert_int4(y1) & 1;
    s0 = s0 == -1 ? convert_float4(t0) == 1.0f ? -1.0f : 1.0f : s0;
    s1 = s1 == -1 ? convert_float4(t1) == 1.0f ? -1.0f : 1.0f : s1;
    tmpDst0 = s0 * exp2(y0 * log2(fabs(x0)));
    tmpDst1 = s1 * exp2(y1 * log2(fabs(x1)));

    _viv_asm(COPY, src0, tmpDst0, 16);
    _viv_asm(COPY, src1, tmpDst1, 16);
    VXC_DP2x8(dst, src0, src1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniExtractOddData_2x8);
    VXC_WriteImage(output, coord, dst, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
}