#version 440 core

#if defined(FP16) || defined(RELU_FP16) || defined(RELU6_FP16)
#extension GL_AMD_gpu_shader_half_float: enable
#define FLOAT4 f16vec4
#else
#define FLOAT4 vec4
#endif

layout(set=0, binding=0) writeonly uniform image2D uOutput;
layout(set=0, binding=1) uniform sampler2D uInput;

layout(set=0, binding=2) uniform sampler2D uKernel;

layout(set=0, binding=3) uniform sampler2D uBias;

layout(set=0, binding=4) readonly uniform constBuffer {
    ivec2 pad;
    ivec2 kernelSize;
    ivec2 stride;
    ivec2 dilate;
    ivec4 inputSize;
    ivec4 outputSize;
    ivec4 offset;//batchOffset, hOffset, outputHeight, other
} uConstant;

#define UP_DIV(x, y) (((x)+(y)-1)/(y))

layout (local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;

void main()
{
    ivec3 pos = ivec3(gl_GlobalInvocationID);
    ivec3 outputSize = uConstant.outputSize.xyz;
    int oz = pos.z % uConstant.outputSize.z;
    int ob = pos.z / uConstant.outputSize.z;

    if (all(lessThan(pos.xy, outputSize.xy)))
    {
        ivec3 inputSize = uConstant.inputSize.xyz;
        ivec2 s0 = pos.xy*uConstant.stride-uConstant.pad;
        ivec2 sta = max(ivec2(0, 0), (UP_DIV(-s0, uConstant.dilate)));
        ivec2 end = min(uConstant.kernelSize, UP_DIV(uConstant.inputSize.xy - s0, uConstant.dilate));
        int fx, fy, fz;
        FLOAT4 color = FLOAT4(texelFetch(uBias, ivec2(oz, 0), 0));
        for (fy=sta.y; fy<end.y; ++fy)
        {
            int sy = fy*uConstant.dilate.y + s0.y;
            for (fx=sta.x; fx<end.x; ++fx)
            {
                int sx = fx*uConstant.dilate.x + s0.x;
                FLOAT4 inputValue = FLOAT4(texelFetch(uInput, ivec2(sx+oz*uConstant.inputSize.x, sy+ob*uConstant.inputSize.y), 0));

                FLOAT4 k = FLOAT4(texelFetch(uKernel, ivec2(fx+fy*uConstant.kernelSize.x, oz), 0));

                color  += k*inputValue;
            }
        }

#if defined(RELU_FP32) || defined(RELU_FP16)
        color = FLOAT4(max(color, FLOAT4(0)));
#endif
#if defined(RELU6_FP32) || defined(RELU6_FP16)
        color = FLOAT4(clamp(color, FLOAT4(0), FLOAT4(6)));
#endif
        imageStore(uOutput, ivec2(pos.x+oz*uConstant.outputSize.x, pos.y+ob*uConstant.outputSize.y), color);
    }

}