#version 440 core
layout(std430) buffer;

layout(set=0, binding=0) readonly buffer destBuffer{
    float data[];
} uDx;

layout(set=0, binding=1) readonly buffer destBuffer1{
    float data[];
} uDy;

layout(set=0, binding=2) writeonly buffer sourceBuffer0{
    float data[];
} uInput;

layout(set=0, binding=3) uniform constBuffer {
    ivec4 size;
} uConstant;

layout (local_size_x = 256, local_size_y = 1, local_size_z = 1) in;

void main()
{
	int pos = int(gl_GlobalInvocationID.x);
	// input output grid layout is NC4HW4

	ivec4 size = uConstant.size;
	int total = size.x * size.y * size.z * size.w;

	if(pos < total)
	{
		// get nchw num of output 
		int x = pos % size.x;
		int tmp = pos / size.x;
		int y = tmp % size.y;
		tmp = tmp / size.y;
		int z = tmp % size.z;
		int n = tmp / size.z;

        int outPos = n * size.x * size.y * size.z + x * size.z + y * size.x * size.z + z;
        int xDPos = n * size.x * size.y * size.z + (x-1) * size.z + y * size.x * size.z + z;
        int yDPos = n * size.x * size.y * size.z + x * size.z + (y-1) * size.x * size.z + z;
        float summer = 0.0;
        if (x > 0) {
            summer = summer + uDx.data[xDPos] - uDx.data[outPos];
        }
        if (y > 0) {
            summer = summer + uDy.data[yDPos] - uDy.data[outPos];
        }
        uInput.data[outPos] = summer;
    }
}