#version 450

#extension GL_EXT_control_flow_attributes : enable

#define CHUNK_SIZE 4

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

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

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_BATCHES]; };
layout(binding = 3) readonly buffer weightBuffer { float weight[]; };
layout(binding = 4) readonly uniform configBuffer {
    uint invRmsBufferIndex; // not used
    uint nColumns;
};
layout(binding = 5) readonly buffer invRmsBuffer { float invRms[]; };

void main() {
    const uint batchIndex = gl_WorkGroupID.y;
    const uint chunkIndex = gl_WorkGroupID.x;

    const BatchInfo info = infos[batchIndex];
    const uint dim = info.inputSizeX / nColumns;
    const uint offset = chunkIndex * CHUNK_SIZE;
    const uint colIndex = offset / dim;
    const float s = invRms[batchIndex * nColumns + colIndex];

    const uint xOffset = info.inputOffset + offset;
    const uint yOffset = info.outputOffset + offset;

    [[unroll]] for (uint i = 0; i < CHUNK_SIZE; i++) {
        y[yOffset + i] = (x[xOffset + i] * s) * weight[(offset + i) % dim];
    }
}
