#include "cl_viv_vx_ext.h"

/********************gemm transposeB fp16 fp16 to fp16*************************/
_viv_uniform VXC_512Bits uniFp16MulFp16AddtoFp32_dp8x2;

_viv_uniform int ac2zero;
_viv_uniform int bc2zero;

__kernel void gemm_transb_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)
{
    int4 coord_out = (int4)(get_global_id(0), get_global_id(1), get_global_id(2), 0);
    int4 coord_a = (int4)(0, coord_out.y, (ac2zero ? 0 : get_global_id(2)), 0);
    int4 coord_b = (int4)(0, coord_out.x, (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);

    for(coord_a.x = 0, coord_b.x = 0; coord_a.x < K;)
    {
        vxc_short8 srcA0,srcA1,srcA2,srcA3;
        vxc_short8 srcB0,srcB1,srcB2,srcB3;
        VXC_OP4(img_load_3d, srcA0, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA1, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 1),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA2, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 2),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcA3, inputA, coord_a.xywz, VXC_5BITOFFSET_XY(0, 3),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB0, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 0),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB1, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 1),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB2, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 2),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        VXC_OP4(img_load_3d, srcB3, inputB, coord_b.xywz, VXC_5BITOFFSET_XY(0, 3),
                    VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord_a.x += 8;
        coord_b.x += 8;

        vxc_half8 halfB0,halfB1,halfB2,halfB3;
        _viv_asm(COPY, halfB0, srcB0, 16);
        _viv_asm(COPY, halfB1, srcB1, 16);
        _viv_asm(COPY, halfB2, srcB2, 16);
        _viv_asm(COPY, halfB3, srcB3, 16);
        vxc_half8 halfA0,halfA1,halfA2,halfA3;
        _viv_asm(COPY, halfA0, srcA0, 16);
        _viv_asm(COPY, halfA1, srcA1, 16);
        _viv_asm(COPY, halfA2, srcA2, 16);
        _viv_asm(COPY, halfA3, srcA3, 16);
        vxc_float4 fpVal;
        VXC_DP8x2(fpVal, halfA0, halfB0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA0, halfB1, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA0, halfB2, VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA0, halfB3, VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        sum0 += fpVal;
        VXC_DP8x2(fpVal, halfA1, halfB0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA1, halfB1, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA1, halfB2, VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA1, halfB3, VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        sum1 += fpVal;
        VXC_DP8x2(fpVal, halfA2, halfB0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA2, halfB1, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA2, halfB2, VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA2, halfB3, VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        sum2 += fpVal;
        VXC_DP8x2(fpVal, halfA3, halfB0, VXC_MODIFIER(0, 0, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA3, halfB1, VXC_MODIFIER(1, 1, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA3, halfB2, VXC_MODIFIER(2, 2, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        VXC_DP8x2(fpVal, halfA3, halfB3, VXC_MODIFIER(3, 3, 0, VXC_RM_TowardZero, 0),
            uniFp16MulFp16AddtoFp32_dp8x2);
        sum3 += fpVal;
    }
    half4 halfDst;
    vxc_short8 valDst;
    _viv_asm(CONV, halfDst, sum0);
    _viv_asm(COPY, valDst, halfDst, 16);
    VXC_WriteImage2DArray(output, coord_out, valDst.s0246, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
    coord_out.y++;
    _viv_asm(CONV, halfDst, sum1);
    _viv_asm(COPY, valDst, halfDst, 16);
    VXC_WriteImage2DArray(output, coord_out, valDst.s0246, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
    coord_out.y++;
    _viv_asm(CONV, halfDst, sum2);
    _viv_asm(COPY, valDst, halfDst, 16);
    VXC_WriteImage2DArray(output, coord_out, valDst.s0246, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
    coord_out.y++;
    _viv_asm(CONV, halfDst, sum3);
    _viv_asm(COPY, valDst, halfDst, 16);
    VXC_WriteImage2DArray(output, coord_out, valDst.s0246, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
}