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

#version 450

layout (constant_id = 0) const int across_spatial = 0;
layout (constant_id = 1) const int across_channel = 0;

layout (binding = 0) readonly buffer bottom_top_blob { sfpvec8 bottom_top_blob_data[]; };
layout (binding = 1) writeonly buffer sqsum_blob { mat2x4 sqsum_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;

    mat2x4 sqsum;

    if (across_spatial == 1 && across_channel == 1)
    {
        int sz = gz * 2;
        int sx = gx * 2;

        ivec2 v_offset0 = sz * p.cstep + sx + ivec2(0, 1);
        ivec2 v_offset1 = v_offset0 + p.cstep;

        if (sz == p.c - 1)
        {
            if (sx == p.w - 1)
            {
                mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset0.r));

                sqsum[0] = v0[0] * v0[0];
                sqsum[1] = v0[1] * v0[1];
            }
            else
            {
                mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset0.r));
                mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset0.g));

                sqsum[0] = v0[0] * v0[0] + v1[0] * v1[0];
                sqsum[1] = v0[1] * v0[1] + v1[1] * v1[1];
            }
        }
        else
        {
            if (sx == p.w - 1)
            {
                mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset0.r));
                mat2x4 v2 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset1.r));

                sqsum[0] = v0[0] * v0[0] + v2[0] * v2[0];
                sqsum[1] = v0[1] * v0[1] + v2[1] * v2[1];
            }
            else
            {
                mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset0.r));
                mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset0.g));
                mat2x4 v2 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset1.r));
                mat2x4 v3 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset1.g));

                sqsum[0] = v0[0] * v0[0] + v1[0] * v1[0] + v2[0] * v2[0] + v3[0] * v3[0];
                sqsum[1] = v0[1] * v0[1] + v1[1] * v1[1] + v2[1] * v2[1] + v3[1] * v3[1];
            }
        }
    }

    if (across_spatial == 1 && across_channel == 0)
    {
        int sz = gz;
        int sx = gx * 4;

        ivec4 v_offset = sz * p.cstep + sx + ivec4(0, 1, 2, 3);

        if (sx == p.w - 1)
        {
            mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.r));

            sqsum[0] = v0[0] * v0[0];
            sqsum[1] = v0[1] * v0[1];
        }
        else if (sx == p.w - 2)
        {
            mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.r));
            mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.g));

            sqsum[0] = v0[0] * v0[0] + v1[0] * v1[0];
            sqsum[1] = v0[1] * v0[1] + v1[1] * v1[1];
        }
        else if (sx == p.w - 3)
        {
            mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.r));
            mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.g));
            mat2x4 v2 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.b));

            sqsum[0] = v0[0] * v0[0] + v1[0] * v1[0] + v2[0] * v2[0];
            sqsum[1] = v0[1] * v0[1] + v1[1] * v1[1] + v2[1] * v2[1];
        }
        else
        {
            mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.r));
            mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.g));
            mat2x4 v2 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.b));
            mat2x4 v3 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.a));

            sqsum[0] = v0[0] * v0[0] + v1[0] * v1[0] + v2[0] * v2[0] + v3[0] * v3[0];
            sqsum[1] = v0[1] * v0[1] + v1[1] * v1[1] + v2[1] * v2[1] + v3[1] * v3[1];
        }
    }

    if (across_spatial == 0 && across_channel == 1)
    {
        int sz = gz * 4;
        int sx = gx;

        ivec4 v_offset = (sz + ivec4(0, 1, 2, 3)) * p.cstep + sx;

        if (sz == p.c - 1)
        {
            mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.r));

            sqsum[0] = v0[0] * v0[0];
            sqsum[1] = v0[1] * v0[1];
        }
        else if (sz == p.c - 2)
        {
            mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.r));
            mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.g));

            sqsum[0] = v0[0] * v0[0] + v1[0] * v1[0];
            sqsum[1] = v0[1] * v0[1] + v1[1] * v1[1];
        }
        else if (sz == p.c - 3)
        {
            mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.r));
            mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.g));
            mat2x4 v2 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.b));

            sqsum[0] = v0[0] * v0[0] + v1[0] * v1[0] + v2[0] * v2[0];
            sqsum[1] = v0[1] * v0[1] + v1[1] * v1[1] + v2[1] * v2[1];
        }
        else
        {
            mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.r));
            mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.g));
            mat2x4 v2 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.b));
            mat2x4 v3 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset.a));

            sqsum[0] = v0[0] * v0[0] + v1[0] * v1[0] + v2[0] * v2[0] + v3[0] * v3[0];
            sqsum[1] = v0[1] * v0[1] + v1[1] * v1[1] + v2[1] * v2[1] + v3[1] * v3[1];
        }
    }

    int gi = gz * p.outcstep + gx;

    sqsum_blob_data[gi] = sqsum;
}
