#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniConvertInt32toUint8_2x8;

_viv_uniform VXC_512Bits uniGemmFp16toFp32Row0Lo_4x4;
_viv_uniform VXC_512Bits uniGemmFp16toFp32Row0Hi_4x4;
_viv_uniform VXC_512Bits uniGemmFp16toFp32Row1Lo_4x4;
_viv_uniform VXC_512Bits uniGemmFp16toFp32Row1Hi_4x4;

_viv_uniform int ac2zero;
_viv_uniform int bc2zero;

_viv_uniform VXC_512Bits uniGemmU8F16toF32Lo_4x4b;

#if (VX_VERSION==2)
__kernel void gemm_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);
    int4 coord_a = (int4)(0, gidy, (ac2zero ? 0 : get_global_id(2)), 0);
    int4 coord_b = (int4)(get_global_id(0), 0, (bc2zero ? 0 : get_global_id(2)), 0);

    half4 valC;
    vxc_short8 srcA0, srcA1, srcA2, srcA3, outC;
    vxc_half8 tmpA0, tmpA1, tmpA2, tmpA3;
    vxc_short16 srcB;
    vxc_half16 tmpB;
    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;
        VXC_OP4(img_load_3d, srcA0, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB.hi, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA1, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 1),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB.hi, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 1),
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA2, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 2),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB.lo, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 2),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA3, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 3),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB.lo, inputB, coord_b.xywz, 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, tmpA0, srcA0, 16);
        _viv_asm(COPY, tmpA1, srcA1, 16);
        _viv_asm(COPY, tmpA2, srcA2, 16);
        _viv_asm(COPY, tmpA3, srcA3, 16);
        _viv_asm(COPY, tmpB.hi, srcB.hi, 16);
        _viv_asm(COPY, tmpB.lo, srcB.lo, 16);
        VXC_DP4x4_b(tempA0, tmpB.hi, tmpB.lo, tmpA0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),
                    uniGemmU8F16toF32Lo_4x4b);
        VXC_DP4x4_b(tempA1, tmpB.hi, tmpB.lo, tmpA1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),
                    uniGemmU8F16toF32Lo_4x4b);
        VXC_DP4x4_b(tempA2, tmpB.hi, tmpB.lo, tmpA2, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),
                    uniGemmU8F16toF32Lo_4x4b);
        VXC_DP4x4_b(tempA3, tmpB.hi, tmpB.lo, tmpA3, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0),
                    uniGemmU8F16toF32Lo_4x4b);
        sum0 += (tempA0);
        sum1 += (tempA1);
        sum2 += (tempA2);
        sum3 += (tempA3);
    }
    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.xywz, 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.xywz, 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.xywz, 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.xywz, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));
}
#else
__kernel void gemm_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);
    int4 coord_a = (int4)(0, gidy, (ac2zero ? 0 : get_global_id(2)), 0);
    int4 coord_b = (int4)(get_global_id(0), 0, (bc2zero ? 0 : get_global_id(2)), 0);

    half4 valC;
    vxc_short8 srcA0, srcB0, srcA1, srcB1, outC;
    vxc_half8 tmpA0, tmpB0, tmpA1, tmpB1;
    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;
        vxc_float4 tempB0, tempB1, tempB2, tempB3;
        VXC_OP4(img_load_3d, srcA0, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB0, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA0, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 1),
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB0, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 1),
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA1, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 2),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB1, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 2),
                    VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA1, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 3),
                    VXC_MODIFIER(4, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB1, inputB, coord_b.xywz, 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, tmpA0, srcA0, 16);
        _viv_asm(COPY, tmpB0, srcB0, 16);
        _viv_asm(COPY, tmpA1, srcA1, 16);
        _viv_asm(COPY, tmpB1, srcB1, 16);

        VXC_DP4x4(tempA0, tmpA0, tmpB0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGemmFp16toFp32Row0Lo_4x4);
        VXC_DP4x4(tempB0, tmpA0, tmpB1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGemmFp16toFp32Row0Hi_4x4);
        VXC_DP4x4(tempA1, tmpA0, tmpB0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGemmFp16toFp32Row1Lo_4x4);
        VXC_DP4x4(tempB1, tmpA0, tmpB1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGemmFp16toFp32Row1Hi_4x4);
        VXC_DP4x4(tempA2, tmpA1, tmpB0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGemmFp16toFp32Row0Lo_4x4);
        VXC_DP4x4(tempB2, tmpA1, tmpB1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGemmFp16toFp32Row0Hi_4x4);
        VXC_DP4x4(tempA3, tmpA1, tmpB0, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGemmFp16toFp32Row1Lo_4x4);
        VXC_DP4x4(tempB3, tmpA1, tmpB1, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGemmFp16toFp32Row1Hi_4x4);
        sum0 += (tempA0 + tempB0);
        sum1 += (tempA1 + tempB1);
        sum2 += (tempA2 + tempB2);
        sum3 += (tempA3 + tempB3);
    }
    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.xywz, 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.xywz, 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.xywz, 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.xywz, outC.s0246, \
                VXC_MODIFIER(0, 3, 0,VXC_RM_TowardZero, 0));
}
#endif

__kernel void gemm_F32F32toF32(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 gidx = get_global_id(0);
    uint gidy = get_global_id(1);

    int4 coord_a = (int4)(0, gidy, (ac2zero ? 0 : get_global_id(2)), 0);
    int4 coord_b = (int4)(gidx, 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);

    vxc_int4 tmpOut0, tmpOut1;
    vxc_uchar16 outC;
    VXC_DP2x8(outC, tmpOut0, tmpOut1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1),\
        uniConvertInt32toUint8_2x8);

    for(int i = 0; i < K; i+=4)
    {
        vxc_float4 tempA0, tempA1, tempA2, tempA3;
        vxc_float4 tempB0, tempB1, tempB2, tempB3;

        coord_a.x = i;
        coord_a.y = gidy;

        coord_b.x = gidx;
        coord_b.y = i;

        tempA0 = read_imagef(inputA, coord_a);
        coord_a.y++;
        tempA1 = read_imagef(inputA, coord_a);
        coord_a.y++;
        tempA2 = read_imagef(inputA, coord_a);
        coord_a.y++;
        tempA3 = read_imagef(inputA, coord_a);

        tempB0 = read_imagef(inputB, coord_b);
        coord_b.y++;
        tempB1 = read_imagef(inputB, coord_b);
        coord_b.y++;
        tempB2 = read_imagef(inputB, coord_b);
        coord_b.y++;
        tempB3 = read_imagef(inputB, coord_b);

        sum0 += (tempA0.x * tempB0 + tempA0.y * tempB1 + tempA0.z * tempB2 + tempA0.w * tempB3);
        sum1 += (tempA1.x * tempB0 + tempA1.y * tempB1 + tempA1.z * tempB2 + tempA1.w * tempB3);
        sum2 += (tempA2.x * tempB0 + tempA2.y * tempB1 + tempA2.z * tempB2 + tempA2.w * tempB3);
        sum3 += (tempA3.x * tempB0 + tempA3.y * tempB1 + tempA3.z * tempB2 + tempA3.w * tempB3);
    }
    coord_b = (int4)(gidx, gidy, get_global_id(2), 0);
    write_imagef(output, coord_b, sum0);
    coord_b.y++;
    write_imagef(output, coord_b, sum1);
    coord_b.y++;
    write_imagef(output, coord_b, sum2);
    coord_b.y++;
    write_imagef(output, coord_b, sum3);
}