#version 440 core

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

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

layout (local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
#define UP_DIV(x, y) (((x)+(y)-1)/(y))

void main()
{
    ivec3 posTmp = ivec3(gl_GlobalInvocationID);
    ivec3 pos;
    pos.x = posTmp.x % uConstant.outputSize.x;
    int subIndex = posTmp.x / uConstant.outputSize.x;
    pos.y = subIndex % uConstant.outputSize.y;
    pos.z = subIndex / uConstant.outputSize.y;
    ivec3 outputSize = uConstant.outputSize.xyz;
    int oz = pos.z % uConstant.outputSize.z;
    int ob = pos.z / uConstant.outputSize.z;
    int batchOffset = uConstant.offset.x * uConstant.outputSize.z;
    if (pos.x < outputSize.x && pos.y < outputSize.y && ob < uConstant.outputSize.w)
    {
        int sourceXIndex = pos.x + pos.y*uConstant.outputSize.x + ob*uConstant.outputSize.x*uConstant.outputSize.y;
        int sourceY = sourceXIndex / 4;
        int sourceX = 4*oz + sourceXIndex % 4;

        vec4 color = texelFetch(uInput, ivec2(sourceX, sourceY), 0) + texelFetch(uBias, ivec2(oz, 0), 0);
#ifdef RELU
        color = max(color, vec4(0));
#endif
#ifdef RELU6
        color = clamp(color, vec4(0), vec4(6));
#endif
        int dzb_offset = pos.z + batchOffset;
        int dz_offset = dzb_offset % uConstant.outputSize.z;
        int db_offset = dzb_offset / uConstant.outputSize.z;
        imageStore(uOutput, ivec2(pos.x + dz_offset * uConstant.outputSize.x, pos.y + uConstant.offset.z * db_offset + uConstant.offset.y), color);
    }
}