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

#version 450

#extension GL_GOOGLE_include_directive: enable
#include "vulkan_activation.comp"

layout (constant_id = 0) const int scale_in_data_size = 0;
layout (constant_id = 1) const float scale_in_value = 1.f;
layout (constant_id = 2) const int scale_out_data_size = 0;
layout (constant_id = 3) const float scale_out_value = 1.f;
layout (constant_id = 4) const int bias_data_size = 0;
layout (constant_id = 5) const float bias_value = 0.f;
layout (constant_id = 6) const int activation_type = 0;
layout (constant_id = 7) const float activation_param_0 = 0;
layout (constant_id = 8) const float activation_param_1 = 0;

#define shape_constant_id_offset 9
layout (constant_id = shape_constant_id_offset + 0) const uint c = 0;
layout (constant_id = shape_constant_id_offset + 1) const uint in_stride = 0;
layout (constant_id = shape_constant_id_offset + 2) const uint out_stride = 0;

layout (binding = 0) readonly buffer bottom_blob { ivec8 bottom_blob_data[]; };
layout (binding = 1) writeonly buffer top_blob { sint8vec8 top_blob_data[]; };
layout (binding = 2) readonly buffer scale_in_blob { sfpvec8 scale_in_blob_data[]; };
layout (binding = 3) readonly buffer scale_out_blob { sfpvec8 scale_out_blob_data[]; };
layout (binding = 4) readonly buffer bias_blob { sfpvec8 bias_blob_data[]; };

layout (push_constant) uniform parameter
{
    uint c;
    uint in_stride;
    uint out_stride;
} p;

void main()
{
    const uint gi = gl_GlobalInvocationID.x;

    if (gi >= psc(in_stride) * psc(c))
        return;

    const uint gy = gi / psc(in_stride);
    const uint gx = gi % psc(in_stride);

    ivec8 v = bottom_blob_data[gi];

    afpvec8 scale_in;
    if (scale_in_data_size == 1)
    {
        scale_in = afpvec8(afpvec4(scale_in_value), afpvec4(scale_in_value));
    }
    else
    {
        scale_in = buffer_ld8(scale_in_blob_data, gy);
    }

    afpvec8 bias;
    if (bias_data_size == 0)
    {
        bias = afpvec8(afpvec4(0.f), afpvec4(0.f));
    }
    else if (bias_data_size == 1)
    {
        bias = afpvec8(afpvec4(bias_value), afpvec4(bias_value));
    }
    else
    {
        bias = buffer_ld8(bias_blob_data, gy);
    }

    afpvec8 v_fp;
    v_fp[0] = afpvec4(v.abcd) * scale_in[0] + bias[0];
    v_fp[1] = afpvec4(v.efgh) * scale_in[1] + bias[1];

    v_fp = activation_afpvec8(v_fp, activation_type, activation_param_0, activation_param_1);

    afpvec8 scale_out;
    if (scale_out_data_size == 1)
    {
        scale_out = afpvec8(afpvec4(scale_out_value), afpvec4(scale_out_value));
    }
    else
    {
        scale_out = buffer_ld8(scale_out_blob_data, gy);
    }

    ivec8 v_int;
    v_int.abcd = ivec4(round(clamp(v_fp[0] * scale_out[0], afp(-127.f), afp(127.f))));
    v_int.efgh = ivec4(round(clamp(v_fp[1] * scale_out[1], afp(-127.f), afp(127.f))));

    const uint outgi = gy * psc(out_stride) + gx;

    i8buffer_st8(top_blob_data, outgi, v_int);
}
