#version 450

#define N_THREADS 128

layout(local_size_x = N_THREADS, local_size_y = 1, local_size_z = 1) in;

layout(constant_id = 0) const uint N_BATCHES = 32;
layout(constant_id = 1) const uint N_Z = 1;

struct BatchInfo {
    uint inputOffset;
    uint inputSizeX;
    uint outputOffset;
    uint outputSizeX;
};

layout(binding = 0) readonly buffer inputBuffer { float x[]; };
layout(binding = 1) writeonly buffer outputBuffer { float y[]; };
layout(binding = 2) readonly uniform batchInfosBuffer { BatchInfo infos[N_Z * N_BATCHES]; };
layout(binding = 3) readonly buffer weightBuffer { float weight[]; };
layout(binding = 4) readonly uniform opConfigBuffer {
    uint nExperts;
    uint nActiveExperts;
    uint activeExpertIndexesBufferIndex;
};
layout(binding = 5) readonly buffer activeExpertIndexesBuffer { float activeExpertIndexes[]; };

void main() {
    const uint threadIndex = gl_LocalInvocationID.x;
    const uint nWorkGroups = gl_NumWorkGroups.x;
    const uint workGroupIndex = gl_WorkGroupID.x;
    const uint batchIndex = gl_WorkGroupID.y;
    const uint zIndex = gl_WorkGroupID.z;

    BatchInfo info = infos[zIndex * N_BATCHES + batchIndex];
    const uint slice = info.outputSizeX / nWorkGroups;
    const uint rest = info.outputSizeX % nWorkGroups;
    const uint dim = slice + (workGroupIndex < rest ? 1 : 0);
    const uint offset = workGroupIndex * slice + min(workGroupIndex, rest);

    const uint expertIndex = nExperts == 0
        ? 0
        : uint(activeExpertIndexes[batchIndex * nActiveExperts + zIndex]);
    const uint expertOffset = expertIndex * info.inputSizeX * info.outputSizeX;

    const uint inputSizeX = info.inputSizeX;
    const uint xOffset = info.inputOffset;
    const uint yOffset = info.outputOffset;

    for (uint i = threadIndex; i < dim; i += N_THREADS) {
        const uint d = offset + i;
        const uint wOffset = expertOffset + d * inputSizeX;

        float sum = 0.0f;
        for (uint j = 0; j < inputSizeX; j++) {
            sum += x[xOffset + j] * weight[wOffset + j];
        }
        y[yOffset + d] = sum;
    }
}
