// Copyright 2020 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 { sfpvec8 coeffs_blob_data[]; };
layout (binding = 1) readonly buffer mean_blob { mat2x4 mean_data[]; };
layout (binding = 2) readonly buffer var_blob { mat2x4 var_data[]; };
layout (binding = 3) readonly buffer gamma_blob { sfpvec8 gamma_data[]; };
layout (binding = 4) readonly buffer beta_blob { sfpvec8 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;

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

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

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

    buffer_st8(coeffs_blob_data, gx*2, afpvec8(a));
    buffer_st8(coeffs_blob_data, gx*2 +1, afpvec8(b));
}
