wehub-resource-sync
87 lines
2.4 KiB
Plaintext
87 lines
2.4 KiB
Plaintext
layout(set=0, binding=0) writeonly buffer destBuffer{
|
|
FLOAT4 data[];
|
|
}uOutput;
|
|
|
|
layout(set=0, binding=1) readonly buffer sourceBuffer0{
|
|
FLOAT4 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;
|
|
|
|
FLOAT4 LoadSample(int positionX, int positionY, int c, int n) {
|
|
FLOAT4 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 = FLOAT4(0.0);
|
|
} else {
|
|
value = uInput.data[0
|
|
+ positionX
|
|
+ positionY * width
|
|
+ n * width * height
|
|
+ c * width * height * uGridSampleParam.inShape.w
|
|
];
|
|
}
|
|
#else
|
|
positionX = clamp(positionX, 0, width - 1);
|
|
positionY = clamp(positionY, 0, height - 1);
|
|
value = uInput.data[0
|
|
+ positionX
|
|
+ positionY * width
|
|
+ n * width * height
|
|
+ c * width * height * uGridSampleParam.inShape.w
|
|
];
|
|
#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[2 * gridPosition + 0];
|
|
FLOAT gridY = uGrid.data[2 * gridPosition + 1];
|
|
|
|
// compute position of input
|
|
FLOAT a = FLOAT(uGridSampleParam.alignCorners);
|
|
FLOAT b = FLOAT(1.0f) - a;
|
|
int positionX = int(floor(((FLOAT(1.0) + gridX) * (FLOAT(inputShape.x) - a) - b) * FLOAT(0.5f) + FLOAT(0.5f)));
|
|
int positionY = int(floor(((FLOAT(1.0) + gridY) * (FLOAT(inputShape.y) - a) - b) * FLOAT(0.5f) + FLOAT(0.5f))); FLOAT4 value = LoadSample(positionX, positionY, z, n);
|
|
|
|
uOutput.data[0
|
|
+ x
|
|
+ y * outputShape.x
|
|
+ z * outputShape.x * outputShape.y * outputShape.w
|
|
+ n * outputShape.x * outputShape.y
|
|
] = value;
|
|
}
|
|
}
|