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

117 lines
3.0 KiB
Plaintext

#version 450 core
layout(std430) buffer;
layout(set=0, binding=0) writeonly buffer destBuffer{
float data[];
}uOutput;
layout(set=0, binding=1) readonly buffer sourceBuffer0{
float data[];
} uInput;
layout(set=0, binding=2) readonly buffer sourceBuffer1{
float data[];
} uGrid;
layout(set=0, binding=3) uniform gridSampleBuffer{
ivec4 inShape; // inW, inH
ivec4 outShape; // outW, outH
bool alignCorners;
}uGridSampleParam;
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
float LoadSample(int positionX, int positionY, int c, int n) {
float value;
int width = uGridSampleParam.inShape.x;
int height = uGridSampleParam.inShape.y;
#ifdef PAD_MODE_ZEROS
if (positionX < 0 || positionX >= width || positionY < 0 || positionY >= height) {
value = 0.0;
} else {
value = uInput.data[0
+ positionX * uGridSampleParam.inShape.z
+ positionY * width * uGridSampleParam.inShape.z
+ n * width * height * uGridSampleParam.inShape.z
+ c
];
}
#else
positionX = clamp(positionX, 0, width - 1);
positionY = clamp(positionY, 0, height - 1);
value = uInput.data[0
+ positionX * uGridSampleParam.inShape.z
+ positionY * width * uGridSampleParam.inShape.z
+ n * width * height * uGridSampleParam.inShape.z
+ c
];
#endif
return value;
}
void main()
{
int pos = int(gl_GlobalInvocationID.x);
// input output grid layout is NC4HW4
ivec4 inputShape = uGridSampleParam.inShape;
ivec4 outputShape = uGridSampleParam.outShape;
int total = outputShape.x * outputShape.y * outputShape.z * outputShape.w;
if(pos < total)
{
// get nchw num of output
int x = pos % outputShape.x;
int tmp = pos / outputShape.x;
int y = tmp % outputShape.y;
tmp = tmp / outputShape.y;
int z = tmp % outputShape.z;
int n = tmp / outputShape.z;
// get position in grid
int gridPosition = n * outputShape.x * outputShape.y + y * outputShape.x + x;
float gridX = uGrid.data[inputShape.w * gridPosition + 0];
float gridY = uGrid.data[inputShape.w * gridPosition + 1];
// compute position of input
#ifdef NEAREST
float cordH = (gridY) * (inputShape.y);
float cordW = (gridX) * (inputShape.x);
int positionX = int(floor(cordW));
int positionY = int(floor(cordH));
float value = LoadSample(positionX, positionY, z, n);
#else
float cordH = (gridY) * (inputShape.y) - 0.5;
float cordW = (gridX) * (inputShape.x) - 0.5;
int w0_h = int(floor(cordH));
int w0_w = int(floor(cordW));
int w1_h = w0_h + 1;
int w1_w = w0_w + 1;
float oneV = float(1.0);
float i00 = LoadSample(w0_w, w0_h, z, n);
float i01 = LoadSample(w1_w, w0_h, z, n);
float i10 = LoadSample(w0_w, w1_h, z, n);
float i11 = LoadSample(w1_w, w1_h, z, n);
float f0 = float(float(w1_w) - cordW);
float f1 = oneV - f0;
float h0 = float(float(w1_h) - cordH);
float h1 = oneV - h0;
float i0 = i00 * f0 + i01 * f1;
float i1 = i10 * f0 + i11 * f1;
float value = i0 * h0 + i1 * h1;
#endif
uOutput.data[0
+ x * outputShape.z
+ y * outputShape.x * outputShape.z
+ z
+ n * outputShape.x * outputShape.y * outputShape.z
] = value;
}
}