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

#version 450

layout (constant_id = 0) const int c = 0;

#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 cstep = 0;

layout (constant_id = shape_constant_id_offset + 3) const int outcstep = 0;

layout (constant_id = shape_constant_id_offset + 4) const int block_x = 0;
layout (constant_id = shape_constant_id_offset + 5) const int block_y = 0;

layout (binding = 0) readonly buffer bottom_blob { sfpvec8 bottom_blob_data[]; };
layout (binding = 1) writeonly buffer bottom_tm_blob { sfpvec8 bottom_tm_blob_data[]; };

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

    int outcstep;

    int block_x;
    int block_y;
} p;

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

    if (gx >= p.block_x || gy >= p.block_y || gz >= c)
        return;

    // load 6x6
    int sx = gx * 4;
    int sy = gy * 4;

    int v_offset_0 = gz * psc(cstep) + sy * psc(w) + sx;
    ivec4 v_offset = v_offset_0 + ivec4(0, 1, 2, 3) * psc(w);
    ivec2 v_offset45 = v_offset_0 + ivec2(4, 5) * psc(w);

    afpvec8 v00 = buffer_ld8(bottom_blob_data, v_offset.r + 0);
    afpvec8 v01 = sx + 1 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.r + 1) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v02 = sx + 2 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.r + 2) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v03 = sx + 3 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.r + 3) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v04 = sx + 4 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.r + 4) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v05 = sx + 5 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.r + 5) : afpvec8(afpvec4(0.f), afpvec4(0.f));

    afpvec8 v10 = sy + 1 < psc(h) ? buffer_ld8(bottom_blob_data, v_offset.g + 0) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v11 = sy + 1 < psc(h) && sx + 1 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.g + 1) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v12 = sy + 1 < psc(h) && sx + 2 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.g + 2) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v13 = sy + 1 < psc(h) && sx + 3 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.g + 3) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v14 = sy + 1 < psc(h) && sx + 4 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.g + 4) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v15 = sy + 1 < psc(h) && sx + 5 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.g + 5) : afpvec8(afpvec4(0.f), afpvec4(0.f));

    afpvec8 v20 = sy + 2 < psc(h) ? buffer_ld8(bottom_blob_data, v_offset.b + 0) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v21 = sy + 2 < psc(h) && sx + 1 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.b + 1) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v22 = sy + 2 < psc(h) && sx + 2 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.b + 2) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v23 = sy + 2 < psc(h) && sx + 3 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.b + 3) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v24 = sy + 2 < psc(h) && sx + 4 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.b + 4) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v25 = sy + 2 < psc(h) && sx + 5 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.b + 5) : afpvec8(afpvec4(0.f), afpvec4(0.f));

    afpvec8 v30 = sy + 3 < psc(h) ? buffer_ld8(bottom_blob_data, v_offset.a + 0) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v31 = sy + 3 < psc(h) && sx + 1 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.a + 1) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v32 = sy + 3 < psc(h) && sx + 2 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.a + 2) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v33 = sy + 3 < psc(h) && sx + 3 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.a + 3) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v34 = sy + 3 < psc(h) && sx + 4 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.a + 4) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v35 = sy + 3 < psc(h) && sx + 5 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset.a + 5) : afpvec8(afpvec4(0.f), afpvec4(0.f));

    afpvec8 v40 = sy + 4 < psc(h) ? buffer_ld8(bottom_blob_data, v_offset45.x + 0) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v41 = sy + 4 < psc(h) && sx + 1 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.x + 1) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v42 = sy + 4 < psc(h) && sx + 2 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.x + 2) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v43 = sy + 4 < psc(h) && sx + 3 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.x + 3) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v44 = sy + 4 < psc(h) && sx + 4 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.x + 4) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v45 = sy + 4 < psc(h) && sx + 5 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.x + 5) : afpvec8(afpvec4(0.f), afpvec4(0.f));

    afpvec8 v50 = sy + 5 < psc(h) ? buffer_ld8(bottom_blob_data, v_offset45.y + 0) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v51 = sy + 5 < psc(h) && sx + 1 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.y + 1) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v52 = sy + 5 < psc(h) && sx + 2 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.y + 2) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v53 = sy + 5 < psc(h) && sx + 3 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.y + 3) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v54 = sy + 5 < psc(h) && sx + 4 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.y + 4) : afpvec8(afpvec4(0.f), afpvec4(0.f));
    afpvec8 v55 = sy + 5 < psc(h) && sx + 5 < psc(w) ? buffer_ld8(bottom_blob_data, v_offset45.y + 5) : afpvec8(afpvec4(0.f), afpvec4(0.f));

#define sq2 1.41421356237
#define sq2_d2 1.41421356237/2

    // const float itm[6][6] = {
    //     {1.0f,  0.0f,  -2.5f,  0.0f,  1.0f, 0.0f},
    //     {0.0f, -sq2,   -2.0f,  sq2/2, 1.0f, 0.0f},
    //     {0.0f,  sq2,   -2.0f, -sq2/2, 1.0f, 0.0f},
    //     {0.0f, -sq2/2, -0.5f,  sq2,   1.0f, 0.0f},
    //     {0.0f,  sq2/2, -0.5f, -sq2,   1.0f, 0.0f},
    //     {0.0f,  1.0f,   0.0f,  -2.5f, 0.0f, 1.0f}
    // };

    // implicit transpose
    afpvec8 m00 = v00 - v02 * afp(2.5) + v04;
    afpvec8 m01 = v10 - v12 * afp(2.5) + v14;
    afpvec8 m02 = v20 - v22 * afp(2.5) + v24;
    afpvec8 m03 = v30 - v32 * afp(2.5) + v34;
    afpvec8 m04 = v40 - v42 * afp(2.5) + v44;
    afpvec8 m05 = v50 - v52 * afp(2.5) + v54;

    afpvec8 m10 = v04 - v02 * afp(2) + v03 * afp(sq2_d2) - v01 * afp(sq2);
    afpvec8 m11 = v14 - v12 * afp(2) + v13 * afp(sq2_d2) - v11 * afp(sq2);
    afpvec8 m12 = v24 - v22 * afp(2) + v23 * afp(sq2_d2) - v21 * afp(sq2);
    afpvec8 m13 = v34 - v32 * afp(2) + v33 * afp(sq2_d2) - v31 * afp(sq2);
    afpvec8 m14 = v44 - v42 * afp(2) + v43 * afp(sq2_d2) - v41 * afp(sq2);
    afpvec8 m15 = v54 - v52 * afp(2) + v53 * afp(sq2_d2) - v51 * afp(sq2);

    afpvec8 m20 = v04 - v02 * afp(2) - v03 * afp(sq2_d2) + v01 * afp(sq2);
    afpvec8 m21 = v14 - v12 * afp(2) - v13 * afp(sq2_d2) + v11 * afp(sq2);
    afpvec8 m22 = v24 - v22 * afp(2) - v23 * afp(sq2_d2) + v21 * afp(sq2);
    afpvec8 m23 = v34 - v32 * afp(2) - v33 * afp(sq2_d2) + v31 * afp(sq2);
    afpvec8 m24 = v44 - v42 * afp(2) - v43 * afp(sq2_d2) + v41 * afp(sq2);
    afpvec8 m25 = v54 - v52 * afp(2) - v53 * afp(sq2_d2) + v51 * afp(sq2);

    afpvec8 m30 = v04 - v02 * afp(0.5) + v03 * afp(sq2) - v01 * afp(sq2_d2);
    afpvec8 m31 = v14 - v12 * afp(0.5) + v13 * afp(sq2) - v11 * afp(sq2_d2);
    afpvec8 m32 = v24 - v22 * afp(0.5) + v23 * afp(sq2) - v21 * afp(sq2_d2);
    afpvec8 m33 = v34 - v32 * afp(0.5) + v33 * afp(sq2) - v31 * afp(sq2_d2);
    afpvec8 m34 = v44 - v42 * afp(0.5) + v43 * afp(sq2) - v41 * afp(sq2_d2);
    afpvec8 m35 = v54 - v52 * afp(0.5) + v53 * afp(sq2) - v51 * afp(sq2_d2);

    afpvec8 m40 = v04 - v02 * afp(0.5) - v03 * afp(sq2) + v01 * afp(sq2_d2);
    afpvec8 m41 = v14 - v12 * afp(0.5) - v13 * afp(sq2) + v11 * afp(sq2_d2);
    afpvec8 m42 = v24 - v22 * afp(0.5) - v23 * afp(sq2) + v21 * afp(sq2_d2);
    afpvec8 m43 = v34 - v32 * afp(0.5) - v33 * afp(sq2) + v31 * afp(sq2_d2);
    afpvec8 m44 = v44 - v42 * afp(0.5) - v43 * afp(sq2) + v41 * afp(sq2_d2);
    afpvec8 m45 = v54 - v52 * afp(0.5) - v53 * afp(sq2) + v51 * afp(sq2_d2);

    afpvec8 m50 = v01 - v03 * afp(2.5) + v05;
    afpvec8 m51 = v11 - v13 * afp(2.5) + v15;
    afpvec8 m52 = v21 - v23 * afp(2.5) + v25;
    afpvec8 m53 = v31 - v33 * afp(2.5) + v35;
    afpvec8 m54 = v41 - v43 * afp(2.5) + v45;
    afpvec8 m55 = v51 - v53 * afp(2.5) + v55;

    v00 = m00 - m02 * afp(2.5) + m04;
    v10 = m10 - m12 * afp(2.5) + m14;
    v20 = m20 - m22 * afp(2.5) + m24;
    v30 = m30 - m32 * afp(2.5) + m34;
    v40 = m40 - m42 * afp(2.5) + m44;
    v50 = m50 - m52 * afp(2.5) + m54;

    v01 = m04 - m02 * afp(2) + m03 * afp(sq2_d2) - m01 * afp(sq2);
    v11 = m14 - m12 * afp(2) + m13 * afp(sq2_d2) - m11 * afp(sq2);
    v21 = m24 - m22 * afp(2) + m23 * afp(sq2_d2) - m21 * afp(sq2);
    v31 = m34 - m32 * afp(2) + m33 * afp(sq2_d2) - m31 * afp(sq2);
    v41 = m44 - m42 * afp(2) + m43 * afp(sq2_d2) - m41 * afp(sq2);
    v51 = m54 - m52 * afp(2) + m53 * afp(sq2_d2) - m51 * afp(sq2);

    v02 = m04 - m02 * afp(2) - m03 * afp(sq2_d2) + m01 * afp(sq2);
    v12 = m14 - m12 * afp(2) - m13 * afp(sq2_d2) + m11 * afp(sq2);
    v22 = m24 - m22 * afp(2) - m23 * afp(sq2_d2) + m21 * afp(sq2);
    v32 = m34 - m32 * afp(2) - m33 * afp(sq2_d2) + m31 * afp(sq2);
    v42 = m44 - m42 * afp(2) - m43 * afp(sq2_d2) + m41 * afp(sq2);
    v52 = m54 - m52 * afp(2) - m53 * afp(sq2_d2) + m51 * afp(sq2);

    v03 = m04 - m02 * afp(0.5) + m03 * afp(sq2) - m01 * afp(sq2_d2);
    v13 = m14 - m12 * afp(0.5) + m13 * afp(sq2) - m11 * afp(sq2_d2);
    v23 = m24 - m22 * afp(0.5) + m23 * afp(sq2) - m21 * afp(sq2_d2);
    v33 = m34 - m32 * afp(0.5) + m33 * afp(sq2) - m31 * afp(sq2_d2);
    v43 = m44 - m42 * afp(0.5) + m43 * afp(sq2) - m41 * afp(sq2_d2);
    v53 = m54 - m52 * afp(0.5) + m53 * afp(sq2) - m51 * afp(sq2_d2);

    v04 = m04 - m02 * afp(0.5) - m03 * afp(sq2) + m01 * afp(sq2_d2);
    v14 = m14 - m12 * afp(0.5) - m13 * afp(sq2) + m11 * afp(sq2_d2);
    v24 = m24 - m22 * afp(0.5) - m23 * afp(sq2) + m21 * afp(sq2_d2);
    v34 = m34 - m32 * afp(0.5) - m33 * afp(sq2) + m31 * afp(sq2_d2);
    v44 = m44 - m42 * afp(0.5) - m43 * afp(sq2) + m41 * afp(sq2_d2);
    v54 = m54 - m52 * afp(0.5) - m53 * afp(sq2) + m51 * afp(sq2_d2);

    v05 = m01 - m03 * afp(2.5) + m05;
    v15 = m11 - m13 * afp(2.5) + m15;
    v25 = m21 - m23 * afp(2.5) + m25;
    v35 = m31 - m33 * afp(2.5) + m35;
    v45 = m41 - m43 * afp(2.5) + m45;
    v55 = m51 - m53 * afp(2.5) + m55;

    // store 36
    int v_tm_offset = gz * psc(outcstep) + gy * psc(block_x) + gx;

    buffer_st8(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep) * c, v00);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep) * c, v01);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep) * c, v02);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep) * c, v03);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep) * c, v04);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep) * c, v05);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep) * c, v10);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep) * c, v11);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep) * c, v12);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep) * c, v13);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep) * c, v14);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep) * c, v15);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep) * c, v20);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep) * c, v21);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep) * c, v22);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep) * c, v23);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 16 * psc(outcstep) * c, v24);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 17 * psc(outcstep) * c, v25);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 18 * psc(outcstep) * c, v30);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 19 * psc(outcstep) * c, v31);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 20 * psc(outcstep) * c, v32);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 21 * psc(outcstep) * c, v33);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 22 * psc(outcstep) * c, v34);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 23 * psc(outcstep) * c, v35);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 24 * psc(outcstep) * c, v40);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 25 * psc(outcstep) * c, v41);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 26 * psc(outcstep) * c, v42);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 27 * psc(outcstep) * c, v43);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 28 * psc(outcstep) * c, v44);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 29 * psc(outcstep) * c, v45);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 30 * psc(outcstep) * c, v50);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 31 * psc(outcstep) * c, v51);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 32 * psc(outcstep) * c, v52);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 33 * psc(outcstep) * c, v53);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 34 * psc(outcstep) * c, v54);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 35 * psc(outcstep) * c, v55);
}
