#version 450

#include "common.comp"

#extension GL_KHR_shader_subgroup_arithmetic : require

layout(local_size_x_id = 0) in;

layout (binding = 0) readonly buffer tensorInA { float16_t inA[]; };
layout (binding = 1) readonly buffer tensorInB { float inB[]; };
layout (binding = 2) writeonly buffer tensorOut { float out_[]; };

layout (push_constant) uniform parameter {
    uint inAOff;
    uint inBOff;
    uint outOff;
    int ne00;
    int ne01;
    int ne02;
    uint nb00;
    uint nb01;
    uint nb02;
    int ne10;
    int ne11;
    int ne12;
    uint nb10;
    uint nb11;
    uint nb12;
    int ne0;
    int ne1;
    uint r2;
    uint r3;
} pcs;

#define N_F16_F32 4

void main() {
    const uint r0 = gl_WorkGroupID.x;
    const uint rb = gl_WorkGroupID.y*N_F16_F32;
    const uint im = gl_WorkGroupID.z;

    const uint i12 = im%pcs.ne12;
    const uint i13 = im/pcs.ne12;

    const uint offset0 = r0*pcs.nb01 + (i12/pcs.r2)*pcs.nb02 + (i13/pcs.r3)*pcs.nb02*pcs.ne02;

    const uint x = offset0 / 2 + pcs.inAOff; // Based from inA

    for (uint row = 0; row < N_F16_F32; ++row) {
        uint r1 = rb + row;
        if (r1 >= pcs.ne11) {
            break;
        }

        const uint y = (r1*pcs.nb11 + im*pcs.nb12) / 4 + pcs.inBOff; // Based from inB

        float sumf = 0;
        for (uint i = gl_SubgroupInvocationID.x; i < pcs.ne00; i += gl_SubgroupSize) {
            sumf += float(inA[x+i]) * float(inB[y+i]);
        }

        const float all_sum = subgroupAdd(sumf);
        if (subgroupElect()) {
            out_[im*pcs.ne1*pcs.ne0 + r1*pcs.ne0 + r0 + pcs.outOff] = all_sum;
        }
    }
}