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

#version 450

layout (constant_id = 0) const int maxk = 1;

#define shape_constant_id_offset 1
layout (constant_id = shape_constant_id_offset + 0) const int w = 0;
layout (constant_id = shape_constant_id_offset + 1) const int h = 0;
layout (constant_id = shape_constant_id_offset + 2) const int c = 0;
layout (constant_id = shape_constant_id_offset + 3) const int cstep = 0;

layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0;
layout (constant_id = shape_constant_id_offset + 5) const int outc = 0;

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

layout (push_constant) uniform parameter
{
    int w;
    int h;
    int c;
    int cstep;

    int outcstep;
    int outc;
} p;

void main()
{
    int gx = int(gl_GlobalInvocationID.x) * 4;
    int gy = int(gl_GlobalInvocationID.y);

    if (gx >= psc(outcstep) || gy >= psc(outc))
        return;

    afpvec8 sum0 = afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 sum1 = afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 sum2 = afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 sum3 = afpvec8(afpvec4(0.f), afpvec4(0.f));

    int v_offset = gx;
    int w_offset = gy * psc(c) * 8;

    for (int z = 0; z < psc(c); z++)
    {
        afpvec4 v0 = buffer_ld4(bottom_blob_data, v_offset + 0);
        afpvec4 v1 = buffer_ld4(bottom_blob_data, v_offset + 1);
        afpvec4 v2 = buffer_ld4(bottom_blob_data, v_offset + 2);
        afpvec4 v3 = buffer_ld4(bottom_blob_data, v_offset + 3);

        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;
        sum0[0].r += dot(v0, k0);
        sum0[0].g += dot(v0, k1);
        sum0[0].b += dot(v0, k2);
        sum0[0].a += dot(v0, k3);
        sum0[1].r += dot(v0, k4);
        sum0[1].g += dot(v0, k5);
        sum0[1].b += dot(v0, k6);
        sum0[1].a += dot(v0, k7);

        sum1[0].r += dot(v1, k0);
        sum1[0].g += dot(v1, k1);
        sum1[0].b += dot(v1, k2);
        sum1[0].a += dot(v1, k3);
        sum1[1].r += dot(v1, k4);
        sum1[1].g += dot(v1, k5);
        sum1[1].b += dot(v1, k6);
        sum1[1].a += dot(v1, k7);

        sum2[0].r += dot(v2, k0);
        sum2[0].g += dot(v2, k1);
        sum2[0].b += dot(v2, k2);
        sum2[0].a += dot(v2, k3);
        sum2[1].r += dot(v2, k4);
        sum2[1].g += dot(v2, k5);
        sum2[1].b += dot(v2, k6);
        sum2[1].a += dot(v2, k7);

        sum3[0].r += dot(v3, k0);
        sum3[0].g += dot(v3, k1);
        sum3[0].b += dot(v3, k2);
        sum3[0].a += dot(v3, k3);
        sum3[1].r += dot(v3, k4);
        sum3[1].g += dot(v3, k5);
        sum3[1].b += dot(v3, k6);
        sum3[1].a += dot(v3, k7);

        v_offset += psc(cstep);
        w_offset += 8;
    }

    const int gi = gy * psc(outcstep) + gx;

    buffer_st8(col_blob_data, gi, sum0);
    if (gx + 1 < psc(outcstep)) buffer_st8(col_blob_data, gi + 1, sum1);
    if (gx + 2 < psc(outcstep)) buffer_st8(col_blob_data, gi + 2, sum2);
    if (gx + 3 < psc(outcstep)) buffer_st8(col_blob_data, gi + 3, sum3);
}
