#version 450

#extension GL_EXT_control_flow_attributes : enable

#define N_THREADS 256

layout(local_size_x = N_THREADS, 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 uniform opConfigBuffer {
    float epsilon;
    uint nColumns;
};

shared float sums[N_THREADS];

void main() {
    const uint threadIndex = gl_LocalInvocationID.x;
    const uint batchIndex = gl_WorkGroupID.y;
    const uint colIndex = gl_WorkGroupID.x;

    const BatchInfo info = infos[batchIndex];
    const uint dim = info.inputSizeX / nColumns;
    const uint offset = info.inputOffset + dim * colIndex;

    float sum = 0.0f;
    for (uint i = threadIndex; i < dim; i += N_THREADS) {
        float v = x[offset + i];
        sum += v * v;
    }
    sums[threadIndex] = sum;

    barrier();

    [[unroll]] for (uint i = N_THREADS / 2; i > 0; i >>= 1) {
        if (threadIndex < i)
            sums[threadIndex] += sums[threadIndex + i];
        barrier();
    }

    if (threadIndex == 0) {
        y[batchIndex * nColumns + colIndex] = inversesqrt((sums[0] / float(dim)) + epsilon);
    }
}
