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

#version 450

layout (constant_id = 0) const int M = 0;
layout (constant_id = 1) const int N = 0;
layout (constant_id = 2) const int K = 0;
layout (constant_id = 3) const int B = 0;

layout (binding = 0) readonly buffer qkcross_blob { sfpvec4 qkcross_blob_data[]; };
layout (binding = 1) readonly buffer v_blob { sfp v_blob_data[]; };
layout (binding = 2) writeonly buffer qkvcross_blob { sfp qkvcross_blob_data[]; };

layout (push_constant) uniform parameter
{
    int M;
    int N;
    int K;
    int B;
} p;

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

    if (gx >= psc(N) || gy >= psc(M) || gz >= psc(B))
        return;

    afpvec4 sum = afpvec4(0.f);

    for (int k = 0; k < psc(K); k++)
    {
        const int ai = gz * psc(M) * psc(K) + gy * psc(K) + k;
        afpvec4 qk0 = buffer_ld4(qkcross_blob_data, ai);

        const int bi = gz * psc(N) * psc(K) + gx * psc(K) + k;
        afp v0 = buffer_ld1(v_blob_data, bi);

        sum += qk0 * v0;
    }

    const int gi = (gz * psc(M) * psc(N) + gx * psc(M) + gy) * 4;

    buffer_st1(qkvcross_blob_data, gi, sum.r);
    buffer_st1(qkvcross_blob_data, gi + 1, sum.g);
    buffer_st1(qkvcross_blob_data, gi + 2, sum.b);
    buffer_st1(qkvcross_blob_data, gi + 3, sum.a);
}
