#include "cl_viv_vx_ext.h"

_viv_uniform float input0Scale;
_viv_uniform VXC_512Bits uniConvertUint8SubZpToFp32_4x4;
_viv_uniform VXC_512Bits uniConvert1stFp16ToFp32_4x4;

_viv_uniform int ac2zero;
_viv_uniform int bc2zero;
_viv_uniform VXC_512Bits uniGemmU8F16toF32Lo_4x4b;
_viv_uniform VXC_512Bits uniGemmU8F16toF32Hi_4x4b;
_viv_uniform VXC_512Bits uniGemmFp16MulZptoFp32_4x4;

_viv_uniform VXC_512Bits uniGemm1stU8F16toF32Lo_4x4;
_viv_uniform VXC_512Bits uniGemm2ndU8F16toF32Lo_4x4;
_viv_uniform VXC_512Bits uniGemm1stU8F16toF32Hi_4x4;
_viv_uniform VXC_512Bits uniGemm2ndU8F16toF32Hi_4x4;

#if (VX_VERSION==2)
#define GEMM_QINT_F16_TO_F16(src0_type_name, read_type) \
__kernel void gemm_##src0_type_name##F16toF16( \
        image2d_array_t inputA, image2d_array_t inputB, image2d_array_t output, \
        int transposeA, int transposeB, int adjointA, int adjointB, uint M, uint K, uint N) \
{ \
    uint gidy = get_global_id(1); \
    uint gidz = get_global_id(2); \
    vxc_short16 srcB; \
    vxc_half16 tmpB; \
    half4 valC; \
    read_type srcA0, srcA1; \
    vxc_short8 outC; \
    int4 coord_a = (int4)(0, gidy, (ac2zero ? 0 : gidz), 0); \
    int4 coord_b = (int4)(get_global_id(0), 0, (bc2zero ? 0 : gidz), 0); \
    vxc_float4 sum0 = (vxc_float4)(0), sum1 = (vxc_float4)(0); \
    vxc_float4 sum2 = (vxc_float4)(0), sum3 = (vxc_float4)(0); \
 \
    int8 inputA_desc, inputB_desc, output_desc; \
    _viv_asm(COPY, inputA_desc, inputA, sizeof(inputA_desc)); \
    int baseAddr_a = (int)coord_a.z * inputA_desc.s4 + inputA_desc.s0; \
    _viv_asm(MOV, coord_a.w, baseAddr_a);  \
    _viv_asm(COPY, inputB_desc, inputB, sizeof(inputB_desc)); \
    int baseAddr_b = (int)coord_b.z * inputB_desc.s4 + inputB_desc.s0; \
    _viv_asm(MOV, coord_b.w, baseAddr_b);  \
 \
    for(coord_a.x = 0, coord_b.y = 0; coord_a.x < K;) \
    { \
        vxc_float4 tempA0, tempA1, tempA2, tempA3, tmpZpScale; \
        vxc_float4 tempB0, tempB1, tempB2, tempB3; \
        VXC_OP4(img_load_3d, srcA0, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB.hi, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcA0, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 1), \
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB.hi, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 1), \
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcA1, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 2), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB.lo, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 2), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcA1, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 3), \
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB.lo, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 3), \
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0)); \
        coord_a.x += 4; \
        coord_b.y += 4; \
        _viv_asm(COPY, tmpB.hi, srcB.hi, 16); \
        _viv_asm(COPY, tmpB.lo, srcB.lo, 16); \
        VXC_DP4x4_b(tempA0, tmpB.hi, tmpB.lo, srcA0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemmU8F16toF32Lo_4x4b); \
        VXC_DP4x4_b(tempA1, tmpB.hi, tmpB.lo, srcA0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemmU8F16toF32Hi_4x4b); \
        VXC_DP4x4_b(tempA2, tmpB.hi, tmpB.lo, srcA1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemmU8F16toF32Lo_4x4b); \
        VXC_DP4x4_b(tempA3, tmpB.hi, tmpB.lo, srcA1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemmU8F16toF32Hi_4x4b); \
        VXC_DP4x4(tmpZpScale, tmpB.hi, tmpB.lo, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemmFp16MulZptoFp32_4x4); \
        sum0 += tempA0 + tmpZpScale; \
        sum1 += tempA1 + tmpZpScale; \
        sum2 += tempA2 + tmpZpScale; \
        sum3 += tempA3 + tmpZpScale; \
    } \
    sum0 *= input0Scale; \
    sum1 *= input0Scale; \
    sum2 *= input0Scale; \
    sum3 *= input0Scale; \
    coord_b.y = gidy; \
    _viv_asm(COPY, output_desc, output, sizeof(output_desc)); \
    int baseAddr = (int)gidz * output_desc.s4 + output_desc.s0; \
    _viv_asm(MOV, coord_b.w, baseAddr);  \
 \
    _viv_asm(CONV, valC, sum0); \
    _viv_asm(COPY, outC, valC, 16); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));  \
    coord_b.y++; \
    _viv_asm(CONV, valC, sum1); \
    _viv_asm(COPY, outC, valC, 16); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));  \
    coord_b.y++; \
    _viv_asm(CONV, valC, sum2); \
    _viv_asm(COPY, outC, valC, 16); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));  \
    coord_b.y++; \
    _viv_asm(CONV, valC, sum3); \
    _viv_asm(COPY, outC, valC, 16); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));  \
}
GEMM_QINT_F16_TO_F16(U8, vxc_uchar16)
GEMM_QINT_F16_TO_F16(I8, vxc_char16)
GEMM_QINT_F16_TO_F16(I16, vxc_short8)
#else
#define GEMM_QINT_F16_TO_F16(src0_type_name, read_type) \
__kernel void gemm_##src0_type_name##F16toF16( \
        image2d_array_t inputA, image2d_array_t inputB, image2d_array_t output, \
        int transposeA, int transposeB, int adjointA, int adjointB, uint M, uint K, uint N) \
{ \
    uint gidy = get_global_id(1); \
    uint gidz = get_global_id(2); \
    vxc_short16 srcB; \
    vxc_half16 tmpB; \
    half4 valC; \
    read_type srcA0, srcA1; \
    vxc_short8 outC; \
    int4 coord_a = (int4)(0, gidy, (ac2zero ? 0 : gidz), 0); \
    int4 coord_b = (int4)(get_global_id(0), 0, (bc2zero ? 0 : gidz), 0); \
    vxc_float4 sum0 = (vxc_float4)(0), sum1 = (vxc_float4)(0); \
    vxc_float4 sum2 = (vxc_float4)(0), sum3 = (vxc_float4)(0); \
 \
    int8 inputA_desc, inputB_desc, output_desc; \
    _viv_asm(COPY, inputA_desc, inputA, sizeof(inputA_desc)); \
    int baseAddr_a = (int)coord_a.z * inputA_desc.s4 + inputA_desc.s0; \
    _viv_asm(MOV, coord_a.w, baseAddr_a);  \
    _viv_asm(COPY, inputB_desc, inputB, sizeof(inputB_desc)); \
    int baseAddr_b = (int)coord_b.z * inputB_desc.s4 + inputB_desc.s0; \
    _viv_asm(MOV, coord_b.w, baseAddr_b);  \
 \
    for(coord_a.x = 0, coord_b.y = 0; coord_a.x < K;) \
    { \
        vxc_float4 tempA0, tempA1, tempA2, tempA3, tmpZpScale; \
        vxc_float4 tempB0, tempB1, tempB2, tempB3; \
        VXC_OP4(img_load_3d, srcA0, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB.hi, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcA0, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 1), \
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB.hi, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 1), \
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcA1, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 2), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB.lo, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 2), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcA1, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 3), \
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB.lo, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 3), \
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0)); \
        coord_a.x += 4; \
        coord_b.y += 4; \
        _viv_asm(COPY, tmpB.hi, srcB.hi, 16); \
        _viv_asm(COPY, tmpB.lo, srcB.lo, 16); \
        VXC_DP4x4(tempA0, srcA0, tmpB.hi, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemm1stU8F16toF32Lo_4x4); \
        VXC_DP4x4(tempB0, srcA0, tmpB.lo, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemm2ndU8F16toF32Lo_4x4); \
        VXC_DP4x4(tempA1, srcA0, tmpB.hi, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemm1stU8F16toF32Hi_4x4); \
        VXC_DP4x4(tempB1, srcA0, tmpB.lo, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemm2ndU8F16toF32Hi_4x4); \
        VXC_DP4x4(tempA2, srcA1, tmpB.hi, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemm1stU8F16toF32Lo_4x4); \
        VXC_DP4x4(tempB2, srcA1, tmpB.lo, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemm2ndU8F16toF32Lo_4x4); \
        VXC_DP4x4(tempA3, srcA1, tmpB.hi, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemm1stU8F16toF32Hi_4x4); \
        VXC_DP4x4(tempB3, srcA1, tmpB.lo, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemm2ndU8F16toF32Hi_4x4); \
        VXC_DP4x4(tmpZpScale, tmpB.hi, tmpB.lo, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), \
                    uniGemmFp16MulZptoFp32_4x4); \
        sum0 += tempA0 + tempB0 + tmpZpScale; \
        sum1 += tempA1 + tempB1 + tmpZpScale; \
        sum2 += tempA2 + tempB2 + tmpZpScale; \
        sum3 += tempA3 + tempB3 + tmpZpScale; \
    } \
    sum0 *= input0Scale; \
    sum1 *= input0Scale; \
    sum2 *= input0Scale; \
    sum3 *= input0Scale; \
    coord_b.y = gidy; \
    _viv_asm(COPY, output_desc, output, sizeof(output_desc)); \
    int baseAddr = (int)gidz * output_desc.s4 + output_desc.s0; \
    _viv_asm(MOV, coord_b.w, baseAddr);  \
 \
    _viv_asm(CONV, valC, sum0); \
    _viv_asm(COPY, outC, valC, 16); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));  \
    coord_b.y++; \
    _viv_asm(CONV, valC, sum1); \
    _viv_asm(COPY, outC, valC, 16); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));  \
    coord_b.y++; \
    _viv_asm(CONV, valC, sum2); \
    _viv_asm(COPY, outC, valC, 16); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));  \
    coord_b.y++; \
    _viv_asm(CONV, valC, sum3); \
    _viv_asm(COPY, outC, valC, 16); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));  \
}
GEMM_QINT_F16_TO_F16(U8, vxc_uchar16)
GEMM_QINT_F16_TO_F16(I8, vxc_char16)
GEMM_QINT_F16_TO_F16(I16, vxc_short8)
#endif
