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

#version 450

#extension GL_GOOGLE_include_directive: enable
#include "vulkan_activation.comp"

layout (constant_id = 0) const int bias_term = 0;
layout (constant_id = 1) const int activation_type = 0;
layout (constant_id = 2) const float activation_param_0 = 0;
layout (constant_id = 3) const float activation_param_1 = 0;
layout (constant_id = 4) const int c = 0;

#define shape_constant_id_offset 5
layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0;

layout (constant_id = shape_constant_id_offset + 1) const int block_x = 0;
layout (constant_id = shape_constant_id_offset + 2) const int block_y = 0;

layout (constant_id = shape_constant_id_offset + 3) const int outw = 0;
layout (constant_id = shape_constant_id_offset + 4) const int outh = 0;
layout (constant_id = shape_constant_id_offset + 5) const int outcstep = 0;

layout (binding = 0) readonly buffer top_tm_blob { sfpvec4 top_tm_blob_data[]; };
layout (binding = 1) writeonly buffer top_blob { sfpvec4 top_blob_data[]; };
layout (binding = 2) readonly buffer bias_blob { sfpvec4 bias_data[]; };

layout (push_constant) uniform parameter
{
    int cstep;

    int block_x;
    int block_y;

    int outw;
    int outh;
    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 >= psc(block_x) || gy >= psc(block_y) || gz >= c)
        return;

    // load 16
    int v_tm_offset = gz * psc(cstep) + gy * psc(block_x) + gx;

    afpvec4 v00 = buffer_ld4(top_tm_blob_data, v_tm_offset + 0 * psc(cstep) * c);
    afpvec4 v01 = buffer_ld4(top_tm_blob_data, v_tm_offset + 1 * psc(cstep) * c);
    afpvec4 v02 = buffer_ld4(top_tm_blob_data, v_tm_offset + 2 * psc(cstep) * c);
    afpvec4 v03 = buffer_ld4(top_tm_blob_data, v_tm_offset + 3 * psc(cstep) * c);
    afpvec4 v10 = buffer_ld4(top_tm_blob_data, v_tm_offset + 4 * psc(cstep) * c);
    afpvec4 v11 = buffer_ld4(top_tm_blob_data, v_tm_offset + 5 * psc(cstep) * c);
    afpvec4 v12 = buffer_ld4(top_tm_blob_data, v_tm_offset + 6 * psc(cstep) * c);
    afpvec4 v13 = buffer_ld4(top_tm_blob_data, v_tm_offset + 7 * psc(cstep) * c);
    afpvec4 v20 = buffer_ld4(top_tm_blob_data, v_tm_offset + 8 * psc(cstep) * c);
    afpvec4 v21 = buffer_ld4(top_tm_blob_data, v_tm_offset + 9 * psc(cstep) * c);
    afpvec4 v22 = buffer_ld4(top_tm_blob_data, v_tm_offset + 10 * psc(cstep) * c);
    afpvec4 v23 = buffer_ld4(top_tm_blob_data, v_tm_offset + 11 * psc(cstep) * c);
    afpvec4 v30 = buffer_ld4(top_tm_blob_data, v_tm_offset + 12 * psc(cstep) * c);
    afpvec4 v31 = buffer_ld4(top_tm_blob_data, v_tm_offset + 13 * psc(cstep) * c);
    afpvec4 v32 = buffer_ld4(top_tm_blob_data, v_tm_offset + 14 * psc(cstep) * c);
    afpvec4 v33 = buffer_ld4(top_tm_blob_data, v_tm_offset + 15 * psc(cstep) * c);

    // const float itm[2][4] = {
    //     {1.0f,  1.0f,  1.0f,  0.0f},
    //     {0.0f,  1.0f, -1.0f,  1.0f}
    // };

    // implicit transpose
    afpvec4 m00 = v00 + v01 + v02;
    afpvec4 m01 = v10 + v11 + v12;
    afpvec4 m02 = v20 + v21 + v22;
    afpvec4 m03 = v30 + v31 + v32;

    afpvec4 m10 = v01 - v02 + v03;
    afpvec4 m11 = v11 - v12 + v13;
    afpvec4 m12 = v21 - v22 + v23;
    afpvec4 m13 = v31 - v32 + v33;

    if (bias_term == 1)
    {
        const afpvec4 bias_value = buffer_ld4(bias_data, gz);

        v00 = bias_value + m00 + m01 + m02;
        v10 = bias_value + m10 + m11 + m12;

        v01 = bias_value + m01 - m02 + m03;
        v11 = bias_value + m11 - m12 + m13;
    }
    else
    {
        v00 = m00 + m01 + m02;
        v10 = m10 + m11 + m12;

        v01 = m01 - m02 + m03;
        v11 = m11 - m12 + m13;
    }

    v00 = activation_afpvec4(v00, activation_type, activation_param_0, activation_param_1);
    v01 = activation_afpvec4(v01, activation_type, activation_param_0, activation_param_1);
    v10 = activation_afpvec4(v10, activation_type, activation_param_0, activation_param_1);
    v11 = activation_afpvec4(v11, activation_type, activation_param_0, activation_param_1);

    // store 2x2
    int x = gx * 2;
    int y = gy * 2;

    int v_offset_0 = gz * psc(outcstep) + y * psc(outw) + x;
    int v_offset_1 = v_offset_0 + psc(outw);

    buffer_st4(top_blob_data, v_offset_0 + 0, v00);
    if (x + 1 < psc(outw)) buffer_st4(top_blob_data, v_offset_0 + 1, v01);
    if (y + 1 < psc(outh)) buffer_st4(top_blob_data, v_offset_1 + 0, v10);
    if (y + 1 < psc(outh) && x + 1 < psc(outw)) buffer_st4(top_blob_data, v_offset_1 + 1, v11);
}
