#version 440 core
layout(std140) buffer;

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

layout(set=0, binding=2) 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 inputSize = uConstant.inputSize.xyz;
    ivec3 outputSize = uConstant.outputSize.xyz;
    int fz = pos.z % uConstant.inputSize.z;
    int sb = pos.z / uConstant.inputSize.z;
    int batchOffset = uConstant.offset.x * uConstant.inputSize.z;

    if (pos.x < outputSize.x && pos.y < outputSize.y && sb < uConstant.outputSize.w)
    {
        ivec2 s0 = (pos.xy + ivec2(0, uConstant.offset.y))*uConstant.stride-uConstant.pad;
        int destYOrigin = pos.x + (pos.y)*uConstant.outputSize.x + sb*uConstant.outputSize.x*uConstant.outputSize.y;
        int destY = destYOrigin / 4;
        int destXOffset = destYOrigin % 4;
        int destX = fz*uConstant.kernelSize.x * uConstant.kernelSize.y;
        int szb_offset = (pos.z + batchOffset);
        int sz_offset = szb_offset % uConstant.inputSize.z;
        int sb_offset = szb_offset / uConstant.inputSize.z;
        vec4 color = texelFetch(uInput, ivec2(s0.x+sz_offset*uConstant.inputSize.x, s0.y+sb_offset*uConstant.inputSize.y), 0);
        imageStore(uOutput, ivec2(4*destX+destXOffset, destY), color);
    }
}