// Copyright 2023 Tencent
// SPDX-License-Identifier: BSD-3-Clause

#version 450

#define LOCAL_MEMORY_UNROLL_INCH 8

layout (constant_id = 0) const float alpha = 1.f;
layout (constant_id = 1) const float beta = 1.f;
layout (constant_id = 2) const int transA = 0;
layout (constant_id = 3) const int transB = 0;
layout (constant_id = 4) const int constantA = 0;
layout (constant_id = 5) const int constantB = 0;
layout (constant_id = 6) const int constantC = 0;
layout (constant_id = 7) const int M = 0;
layout (constant_id = 8) const int N = 0;
layout (constant_id = 9) const int K = 0;
layout (constant_id = 10) const int constant_broadcast_type_C = 0;
layout (constant_id = 11) const int output_N1M = 0;
layout (constant_id = 12) const int output_elempack = 0;
layout (constant_id = 13) const int output_elemtype = 0;
layout (constant_id = 14) const int output_transpose = 0;

// TODO psc more

layout (binding = 0) writeonly buffer top_blob { sfp top_blob_data[]; };
layout (binding = 1) readonly buffer A_blob { sfp A_blob_data[]; };
layout (binding = 2) readonly buffer B_blob { sfp B_blob_data[]; };
layout (binding = 3) readonly buffer C_blob { sfp C_blob_data[]; };

layout (push_constant) uniform parameter
{
    int M;
    int N;
    int K;
    int broadcast_type_C;
    int A_dims;
    int A_hstep;
    int B_dims;
    int B_hstep;
    int outdims;
    int outhstep;
} p;

#if NCNN_shader_local_memory
shared lfp tmp_a[8][LOCAL_MEMORY_UNROLL_INCH][2];
shared lfp tmp_b[8][LOCAL_MEMORY_UNROLL_INCH][2];
#endif

void main()
{
    int gx = int(gl_GlobalInvocationID.x) * 2;
    int gy = int(gl_GlobalInvocationID.y) * 2;
    int gz = int(gl_GlobalInvocationID.z);

#if !NCNN_shader_local_memory
    if (gx >= psc(N) || gy >= psc(M) || gz >= 1)
        return;
#endif

    afp sum0 = afp(0.f);
    afp sum1 = afp(0.f);
    afp sum2 = afp(0.f);
    afp sum3 = afp(0.f);

    const int broadcast_type_C = constantC == 1 ? constant_broadcast_type_C : p.broadcast_type_C;

    if (broadcast_type_C == 0)
    {
        sum0 = buffer_ld1(C_blob_data, 0);
        sum1 = sum0;
        sum2 = sum0;
        sum3 = sum0;
    }
    if (broadcast_type_C == 1 || broadcast_type_C == 2)
    {
        sum0 = buffer_ld1(C_blob_data, gy);
        sum1 = sum0;
        sum2 = buffer_ld1(C_blob_data, gy + 1);
        sum3 = sum2;
    }
    if (broadcast_type_C == 3)
    {
        const int ci = gy * psc(N) + gx;
        sum0 = buffer_ld1(C_blob_data, ci);
        sum1 = buffer_ld1(C_blob_data, ci + 1);
        sum2 = buffer_ld1(C_blob_data, ci + psc(N));
        sum3 = buffer_ld1(C_blob_data, ci + psc(N) + 1);
    }
    if (broadcast_type_C == 4)
    {
        sum0 = buffer_ld1(C_blob_data, gx);
        sum1 = buffer_ld1(C_blob_data, gx + 1);
        sum2 = sum0;
        sum3 = sum1;
    }

    sum0 *= afp(beta);
    sum1 *= afp(beta);
    sum2 *= afp(beta);
    sum3 *= afp(beta);

#if NCNN_shader_local_memory
    const int NN = psc(K);

    const int lx = int(gl_LocalInvocationID.x);
    const int ly = int(gl_LocalInvocationID.y);

    int k = 0;
    for (; k + (LOCAL_MEMORY_UNROLL_INCH - 1) < NN; k += LOCAL_MEMORY_UNROLL_INCH)
    {
        {
            if (transA == 1)
            {
                const int ai = (k + lx) * p.A_hstep + gy;
                tmp_a[ly][lx][0] = sfp2lfp(buffer_ld1(A_blob_data, ai));
                tmp_a[ly][lx][1] = sfp2lfp(buffer_ld1(A_blob_data, ai + 1));
            }
            else
            {
                const int ai = gy * p.A_hstep + (k + lx);
                tmp_a[ly][lx][0] = sfp2lfp(buffer_ld1(A_blob_data, ai));
                tmp_a[ly][lx][1] = sfp2lfp(buffer_ld1(A_blob_data, ai + p.A_hstep));
            }

            if (transB == 1)
            {
                const int bi = gx * p.B_hstep + (k + ly);
                tmp_b[lx][ly][0] = sfp2lfp(buffer_ld1(B_blob_data, bi));
                tmp_b[lx][ly][1] = sfp2lfp(buffer_ld1(B_blob_data, bi + p.B_hstep));
            }
            else
            {
                const int bi = (k + ly) * p.B_hstep + gx;
                tmp_b[lx][ly][0] = sfp2lfp(buffer_ld1(B_blob_data, bi));
                tmp_b[lx][ly][1] = sfp2lfp(buffer_ld1(B_blob_data, bi + 1));
            }
        }

        barrier();

        for (int k4 = 0; k4 < LOCAL_MEMORY_UNROLL_INCH; k4++)
        {
            afp a0 = lfp2afp(tmp_a[ly][k4][0]);
            afp a1 = lfp2afp(tmp_a[ly][k4][1]);

            afp b0 = lfp2afp(tmp_b[lx][k4][0]);
            afp b1 = lfp2afp(tmp_b[lx][k4][1]);

            sum0 += a0 * b0;
            sum1 += a0 * b1;
            sum2 += a1 * b0;
            sum3 += a1 * b1;
        }

        barrier();
    }

    if (k < NN)
    {
        const int remain = NN - k;

        if (lx < remain)
        {
            if (transA == 1)
            {
                const int ai = (k + lx) * p.A_hstep + gy;
                tmp_a[ly][lx][0] = sfp2lfp(buffer_ld1(A_blob_data, ai));
                tmp_a[ly][lx][1] = sfp2lfp(buffer_ld1(A_blob_data, ai + 1));
            }
            else
            {
                const int ai = gy * p.A_hstep + (k + lx);
                tmp_a[ly][lx][0] = sfp2lfp(buffer_ld1(A_blob_data, ai));
                tmp_a[ly][lx][1] = sfp2lfp(buffer_ld1(A_blob_data, ai + p.A_hstep));
            }
        }

        if (ly < remain)
        {
            if (transB == 1)
            {
                const int bi = gx * p.B_hstep + (k + ly);
                tmp_b[lx][ly][0] = sfp2lfp(buffer_ld1(B_blob_data, bi));
                tmp_b[lx][ly][1] = sfp2lfp(buffer_ld1(B_blob_data, bi + p.B_hstep));
            }
            else
            {
                const int bi = (k + ly) * p.B_hstep + gx;
                tmp_b[lx][ly][0] = sfp2lfp(buffer_ld1(B_blob_data, bi));
                tmp_b[lx][ly][1] = sfp2lfp(buffer_ld1(B_blob_data, bi + 1));
            }
        }

        barrier();

        for (int k4 = 0; k4 < remain; k4++)
        {
            afp a0 = lfp2afp(tmp_a[ly][k4][0]);
            afp a1 = lfp2afp(tmp_a[ly][k4][1]);

            afp b0 = lfp2afp(tmp_b[lx][k4][0]);
            afp b1 = lfp2afp(tmp_b[lx][k4][1]);

            sum0 += a0 * b0;
            sum1 += a0 * b1;
            sum2 += a1 * b0;
            sum3 += a1 * b1;
        }
    }
#else
    for (int k = 0; k < psc(K); k++)
    {
        afp a0;
        afp a1;
        afp b0;
        afp b1;
        if (transA == 1)
        {
            const int ai = k * p.A_hstep + gy;
            a0 = buffer_ld1(A_blob_data, ai);
            a1 = buffer_ld1(A_blob_data, ai + 1);
        }
        else
        {
            const int ai = gy * p.A_hstep + k;
            a0 = buffer_ld1(A_blob_data, ai);
            a1 = buffer_ld1(A_blob_data, ai + p.A_hstep);
        }

        if (transB == 1)
        {
            const int bi = gx * p.B_hstep + k;
            b0 = buffer_ld1(B_blob_data, bi);
            b1 = buffer_ld1(B_blob_data, bi + p.B_hstep);
        }
        else
        {
            const int bi = k * p.B_hstep + gx;
            b0 = buffer_ld1(B_blob_data, bi);
            b1 = buffer_ld1(B_blob_data, bi + 1);
        }

        sum0 += a0 * b0;
        sum1 += a0 * b1;
        sum2 += a1 * b0;
        sum3 += a1 * b1;
    }
#endif

#if NCNN_shader_local_memory
    if (gx >= psc(N) || gy >= psc(M) || gz >= 1)
        return;
#endif

    sum0 *= afp(alpha);
    sum1 *= afp(alpha);
    sum2 *= afp(alpha);
    sum3 *= afp(alpha);

    if (output_transpose == 1)
    {
        const int gi = gx * p.outhstep + gy;

        buffer_st1(top_blob_data, gi, sum0);
        if (gy + 1 < psc(M)) buffer_st1(top_blob_data, gi + 1, sum2);
        if (gx + 1 < psc(N))
        {
            buffer_st1(top_blob_data, gi + p.outhstep, sum1);
            if (gy + 1 < psc(M)) buffer_st1(top_blob_data, gi + p.outhstep + 1, sum3);
        }
    }
    else
    {
        const int gi = gy * p.outhstep + gx;

        buffer_st1(top_blob_data, gi, sum0);
        if (gx + 1 < psc(N)) buffer_st1(top_blob_data, gi + 1, sum1);
        if (gy + 1 < psc(M))
        {
            buffer_st1(top_blob_data, gi + p.outhstep, sum2);
            if (gx + 1 < psc(N)) buffer_st1(top_blob_data, gi + p.outhstep + 1, sum3);
        }
    }
}
