// Copyright 2019 Tencent
// SPDX-License-Identifier: BSD-3-Clause

#version 450

layout (constant_id = 0) const float eps = 0;
layout (constant_id = 1) const int affine = 0;
layout (constant_id = 2) const int w = 0;

layout (binding = 0) writeonly buffer coeffs_blob { sfp coeffs_blob_data[]; };
layout (binding = 1) readonly buffer mean_blob { float mean_data[]; };
layout (binding = 2) readonly buffer var_blob { float var_data[]; };
layout (binding = 3) readonly buffer gamma_blob { sfp gamma_data[]; };
layout (binding = 4) readonly buffer beta_blob { sfp beta_data[]; };

layout (push_constant) uniform parameter
{
    int w;
} p;

void main()
{
    int gx = int(gl_GlobalInvocationID.x);
    int gy = int(gl_GlobalInvocationID.y);
    int gz = int(gl_GlobalInvocationID.z);

    if (gx >= psc(w) || gy >= 1 || gz >= 1)
        return;

    float mean = mean_data[gx];
    float var = var_data[gx];

    float a;
    float b;
    if (affine == 0)
    {
        a = 1.f / (sqrt(var + eps));
        b = - mean * a;
    }
    else
    {
        float gamma = float(buffer_ld1(gamma_data, gx));
        float beta = float(buffer_ld1(beta_data, gx));

        a = gamma / (sqrt(var + eps));
        b = - mean * a + beta;
    }

    buffer_st1(coeffs_blob_data, gx*2, afp(a));
    buffer_st1(coeffs_blob_data, gx*2 +1, afp(b));
}
