Files
2026-07-13 13:33:03 +08:00

53 lines
1.3 KiB
Plaintext

#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;
}
}