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

#version 450

#define shape_constant_id_offset 0
layout (constant_id = shape_constant_id_offset + 0) const int w = 0;

layout (constant_id = shape_constant_id_offset + 1) const int outw = 0;
layout (constant_id = shape_constant_id_offset + 2) const int outh = 0;

layout (binding = 0) readonly buffer bottom_blob { sfpvec4 bottom_blob_data[]; };
layout (binding = 1) writeonly buffer top_blob { sfpvec8 top_blob_data[]; };
layout (binding = 2) readonly buffer weight_blob { sfpvec4 weight_data[]; };

layout (push_constant) uniform parameter
{
    int w;

    int outw;
    int outh;
} p;

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

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

    const int start = gx * 8;
    const int end = min(gx * 8 + 8, psc(w));

    afpvec8 sum = afpvec8(afpvec4(0.f), afpvec4(0.f));

    int w_offset = (gy * psc(w) + start) * 8;

    for (int i = start; i < end; i++)
    {
        afpvec4 v = buffer_ld4(bottom_blob_data, i);

        afpvec4 k0 = buffer_ld4(weight_data, w_offset + 0);
        afpvec4 k1 = buffer_ld4(weight_data, w_offset + 1);
        afpvec4 k2 = buffer_ld4(weight_data, w_offset + 2);
        afpvec4 k3 = buffer_ld4(weight_data, w_offset + 3);
        afpvec4 k4 = buffer_ld4(weight_data, w_offset + 4);
        afpvec4 k5 = buffer_ld4(weight_data, w_offset + 5);
        afpvec4 k6 = buffer_ld4(weight_data, w_offset + 6);
        afpvec4 k7 = buffer_ld4(weight_data, w_offset + 7);

        // sum += v * k;
        sum[0].r += dot(v, k0);
        sum[0].g += dot(v, k1);
        sum[0].b += dot(v, k2);
        sum[0].a += dot(v, k3);
        sum[1].r += dot(v, k4);
        sum[1].g += dot(v, k5);
        sum[1].b += dot(v, k6);
        sum[1].a += dot(v, k7);

        w_offset += 8;
    }

    const int gi = gy * psc(outw) + gx;
    buffer_st8(top_blob_data, gi, sum);
}
