#version 450

#include "common.comp"

layout(local_size_x = 256) in;

layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; };
layout(binding = 1) buffer restrict tensorOut { float out_[]; };

layout(push_constant) uniform PushConstants {
    uint inOff;
    uint outOff;
    uint ne00;
    uint nb01;
    float eps;
} pcs;

shared float sum[gl_WorkGroupSize.x];

void main() {
    const uint x = (gl_WorkGroupID.x*pcs.nb01/4) + pcs.inOff; // Based from in_
    // MEAN
    // parallel sum
    sum[gl_LocalInvocationID.x] = 0.0;
    for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) {
        sum[gl_LocalInvocationID.x] += in_[x+i00];
    }

    // reduce
    barrier();
    memoryBarrierShared();
    [[unroll]] for (uint i = gl_WorkGroupSize.x/2; i > 0; i /= 2) {
        if (gl_LocalInvocationID.x < i) {
            sum[gl_LocalInvocationID.x] += sum[gl_LocalInvocationID.x + i];
        }
        barrier();
        memoryBarrierShared();
    }

    // broadcast
    if (gl_LocalInvocationID.x == 0) {
        sum[0] /= float(pcs.ne00);
    }
    barrier();
    memoryBarrierShared();
    const float mean = sum[0];

    // recenter
    const uint y = (gl_WorkGroupID.x*pcs.ne00) + pcs.outOff; // Based from out_
    for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) {
        out_[y+i00] = in_[x+i00] - mean;
    }

    // VARIANCE
    // parallel sum
    sum[gl_LocalInvocationID.x] = 0.0;
    for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) {
        sum[gl_LocalInvocationID.x] += out_[y+i00] * out_[y+i00];
    }

    // reduce
    barrier();
    memoryBarrierShared();
    [[unroll]] for (uint i = gl_WorkGroupSize.x/2; i > 0; i /= 2) {
        if (gl_LocalInvocationID.x < i) {
            sum[gl_LocalInvocationID.x] += sum[gl_LocalInvocationID.x + i];
        }
        barrier();
        memoryBarrierShared();
    }

    // broadcast
    if (gl_LocalInvocationID.x == 0) {
        sum[0] /= float(pcs.ne00);
    }
    barrier();
    memoryBarrierShared();
    const float variance = sum[0];

    const float scale = 1.0f/sqrt(variance + pcs.eps);
    for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) {
        out_[y+i00] *= scale;
    }
}