// 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 { sfpvec8 bottom_blob_data[]; };
layout (binding = 1) writeonly buffer top_blob { sfp top_blob_data[]; };
layout (binding = 2) readonly buffer weight_blob { sfpvec8 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));

    afp sum = afp(0.f);

    int w_offset = gy * psc(w);

    for (int i = start; i < end; i++)
    {
        afpvec8 v = buffer_ld8(bottom_blob_data, i);

        afpvec8 k = buffer_ld8(weight_data, w_offset + i);

        // sum += dot(v, k);
        sum += dot(v[0], k[0]) + dot(v[1], k[1]);
    }

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