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

#version 450

layout (binding = 0) readonly buffer bottom_top_blob { sfp bottom_top_blob_data[]; };
layout (binding = 1) writeonly buffer sum_blob { float sum_blob_data[]; };

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

    int outw;
    int outh;
    int outc;
    int outcstep;
} p;

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

    if (gx >= p.outw || gy >= p.outh || gz >= p.outc)
        return;

    float sum;

    int sx = gx * 4;

    int v_offset = gz * p.cstep + sx;

    if (sx == p.w - 1)
    {
        float v0 = float(buffer_ld1(bottom_top_blob_data, v_offset));

        sum = v0;
    }
    else if (sx == p.w - 2)
    {
        float v0 = float(buffer_ld1(bottom_top_blob_data, v_offset));
        float v1 = float(buffer_ld1(bottom_top_blob_data, v_offset + 1));

        sum = v0 + v1;
    }
    else if (sx == p.w - 3)
    {
        float v0 = float(buffer_ld1(bottom_top_blob_data, v_offset));
        float v1 = float(buffer_ld1(bottom_top_blob_data, v_offset + 1));
        float v2 = float(buffer_ld1(bottom_top_blob_data, v_offset + 2));

        sum = v0 + v1 + v2;
    }
    else
    {
        float v0 = float(buffer_ld1(bottom_top_blob_data, v_offset));
        float v1 = float(buffer_ld1(bottom_top_blob_data, v_offset + 1));
        float v2 = float(buffer_ld1(bottom_top_blob_data, v_offset + 2));
        float v3 = float(buffer_ld1(bottom_top_blob_data, v_offset + 3));

        sum = v0 + v1 + v2 + v3;
    }

    int gi = gz * p.outcstep + gx;

    sum_blob_data[gi] = sum;
}
