#include "cl_viv_vx_ext.h"

_viv_uniform int input0_ZP;
_viv_uniform int input1_ZP;
_viv_uniform float output_ZP;
_viv_uniform float outputScale;
_viv_uniform VXC_512Bits uniConvertUint8SubZpToFp32_4x4;
_viv_uniform VXC_512Bits uniConvertUint8SubZpToFp32B_4x4;
_viv_uniform VXC_512Bits uniConvertInt32toUint8_2x8;
_viv_uniform VXC_512Bits uniConvert1stFp16ToFp32_4x4;

_viv_uniform int ac2zero;
_viv_uniform int bc2zero;

#define GEMM_TRANSA_QINT(src0_type_name, src1_type_name, dst_type_name, read0_type, read1_type, write_type) \
__kernel void gemm_transa_##src0_type_name##src1_type_name##to##dst_type_name( \
            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); \
    read0_type srcA; \
    read1_type srcB; \
    write_type outC; \
 \
    int4 coord_a = (int4)(gidy, 0, (ac2zero ? 0 : get_global_id(2)), 0); \
    int4 coord_b = (int4)(get_global_id(0), 0, (bc2zero ? 0 : get_global_id(2)), 0); \
 \
    vxc_float4 sum0 = (vxc_float4)(0); \
    vxc_float4 sum1 = (vxc_float4)(0); \
    vxc_float4 sum2 = (vxc_float4)(0); \
    vxc_float4 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);  \
 \
    short in0_zp, in1_zp; \
    _viv_asm(COPY, in0_zp, input0_ZP, 4); \
    _viv_asm(COPY, in1_zp, input1_ZP, 4); \
 \
    vxc_float4 tempA0; \
    vxc_float4 tempB0; \
 \
    for(coord_a.y = 0, coord_b.y = 0; coord_a.y < K;) \
    { \
        VXC_OP4(img_load_3d, srcA, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        coord_a.y++; \
        coord_b.y++; \
        VXC_DP4x4(tempA0, srcA, in0_zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertUint8SubZpToFp32_4x4); \
        VXC_DP4x4(tempB0, srcB, in1_zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertUint8SubZpToFp32B_4x4); \
        sum0 = (sum0 + tempA0.x * tempB0); \
        sum1 = (sum1 + tempA0.y * tempB0); \
        sum2 = (sum2 + tempA0.z * tempB0); \
        sum3 = (sum3 + tempA0.w * tempB0); \
    } \
    vxc_int4 tmpOut0, tmpOut1; \
    coord_b.y = gidy; \
    _viv_asm(COPY, output_desc, output, sizeof(output_desc)); \
    int baseAddr = (int)get_global_id(2) * output_desc.s4 + output_desc.s0; \
    _viv_asm(MOV, coord_b.w, baseAddr); \
    tmpOut0 = convert_int4_rte(sum0 * outputScale + output_ZP); \
    tmpOut1 = convert_int4_rte(sum1 * outputScale + output_ZP); \
    VXC_DP2x8(outC, tmpOut0, tmpOut1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
        uniConvertInt32toUint8_2x8); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0123, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
    coord_b.y++; \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s4567, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
 \
    coord_b.y++; \
    tmpOut0 = convert_int4_rte(sum2 * outputScale + output_ZP); \
    tmpOut1 = convert_int4_rte(sum3 * outputScale + output_ZP); \
    VXC_DP2x8(outC, tmpOut0, tmpOut1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
        uniConvertInt32toUint8_2x8); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0123, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
    coord_b.y++; \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s4567, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
}
GEMM_TRANSA_QINT(U8, U8, U8, vxc_uchar16, vxc_uchar16, vxc_uchar16)
GEMM_TRANSA_QINT(I8, I8, I8, vxc_char16, vxc_char16, vxc_char16)
GEMM_TRANSA_QINT(I16, I16, I16, vxc_short8, vxc_short8, vxc_short8)

#define GEMM_TRANSA_INPUTB_F16(src0_type_name, read0_type) \
__kernel void gemm_transa_##src0_type_name##F16to##src0_type_name( \
                        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); \
    read0_type srcA, outC; \
    vxc_short8 srcB; \
    vxc_half8 tmpB; \
 \
    int4 coord_a = (int4)(gidy, 0, (ac2zero ? 0 : get_global_id(2)), 0); \
    int4 coord_b = (int4)(get_global_id(0), 0, (bc2zero ? 0 : get_global_id(2)), 0); \
 \
    vxc_float4 sum0 = (vxc_float4)(0); \
    vxc_float4 sum1 = (vxc_float4)(0); \
    vxc_float4 sum2 = (vxc_float4)(0); \
    vxc_float4 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);  \
 \
    short in0_zp; \
    _viv_asm(COPY, in0_zp, input0_ZP, 4); \
 \
    vxc_float4 tempA0; \
    vxc_float4 tempB0; \
 \
    for(coord_a.y = 0, coord_b.y = 0; coord_a.y < K;) \
    { \
        VXC_OP4(img_load_3d, srcA, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        VXC_OP4(img_load_3d, srcB, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 0), \
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0)); \
        coord_a.y++; \
        coord_b.y++; \
        VXC_DP4x4(tempA0, srcA, in0_zp, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),\
            uniConvertUint8SubZpToFp32_4x4); \
        _viv_asm(COPY, tmpB, srcB, 16); \
        VXC_DP4x4(tempB0,tmpB,tmpB,VXC_MODIFIER(0,3,0,VXC_RM_TowardZero,0),uniConvert1stFp16ToFp32_4x4); \
        sum0 = (sum0 + tempA0.x * tempB0); \
        sum1 = (sum1 + tempA0.y * tempB0); \
        sum2 = (sum2 + tempA0.z * tempB0); \
        sum3 = (sum3 + tempA0.w * tempB0); \
    } \
    vxc_int4 tmpOut0, tmpOut1; \
    coord_b.y = gidy; \
    _viv_asm(COPY, output_desc, output, sizeof(output_desc)); \
    int baseAddr = (int)get_global_id(2) * output_desc.s4 + output_desc.s0; \
    _viv_asm(MOV, coord_b.w, baseAddr); \
    tmpOut0 = convert_int4_rte(sum0 * outputScale + output_ZP); \
    tmpOut1 = convert_int4_rte(sum1 * outputScale + output_ZP); \
    VXC_DP2x8(outC, tmpOut0, tmpOut1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
        uniConvertInt32toUint8_2x8); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0123, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
    coord_b.y++; \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s4567, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
 \
    coord_b.y++; \
    tmpOut0 = convert_int4_rte(sum2 * outputScale + output_ZP); \
    tmpOut1 = convert_int4_rte(sum3 * outputScale + output_ZP); \
    VXC_DP2x8(outC, tmpOut0, tmpOut1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
        uniConvertInt32toUint8_2x8); \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s0123, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
    coord_b.y++; \
    VXC_OP4_NoDest(img_store_3d, output, coord_b.xyww, outC.s4567, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0)); \
}
GEMM_TRANSA_INPUTB_F16(U8, vxc_uchar16)
GEMM_TRANSA_INPUTB_F16(I8, vxc_char16)
GEMM_TRANSA_INPUTB_F16(I16, vxc_short8)

__kernel void gemm_transa_F16F16toF16(
                        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);

    half4 valC;
    vxc_short8 srcA, srcB, outC;
    vxc_half8 tmpA, tmpB;

    int4 coord_a = (int4)(gidy, 0, (ac2zero ? 0 : get_global_id(2)), 0);
    int4 coord_b = (int4)(get_global_id(0), 0, (bc2zero ? 0 : get_global_id(2)), 0);

    vxc_float4 sum0 = (vxc_float4)(0);
    vxc_float4 sum1 = (vxc_float4)(0);
    vxc_float4 sum2 = (vxc_float4)(0);
    vxc_float4 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);

    vxc_float4 tempA0;
    vxc_float4 tempB0;

    for(coord_a.y = 0, coord_b.y = 0; coord_a.y < K;)
    {
        VXC_OP4(img_load_3d, srcA, inputA, coord_a.xyww, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB, inputB, coord_b.xyww, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        coord_a.y++;
        coord_b.y++;
        _viv_asm(COPY, tmpA, srcA, 16);
        VXC_DP4x4(tempA0,tmpA,tmpA,VXC_MODIFIER(0,3,0,VXC_RM_TowardZero,0),uniConvert1stFp16ToFp32_4x4);
        _viv_asm(COPY, tmpB, srcB, 16);
        VXC_DP4x4(tempB0,tmpB,tmpB,VXC_MODIFIER(0,3,0,VXC_RM_TowardZero,0),uniConvert1stFp16ToFp32_4x4);

        sum0 = (sum0 + tempA0.x * tempB0);
        sum1 = (sum1 + tempA0.y * tempB0);
        sum2 = (sum2 + tempA0.z * tempB0);
        sum3 = (sum3 + tempA0.w * tempB0);
    }
    coord_b.y = gidy;
    _viv_asm(COPY, output_desc, output, sizeof(output_desc));
    int baseAddr = (int)get_global_id(2) * 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));
}