MNN/source/backend/vulkan/execution/glsl/resizeNearest.comp

39 lines
1.0 KiB
Plaintext

#version 450 core
layout(std430) buffer;
layout(std430) uniform;
layout(set=0, binding=0) uniform mediump sampler3D uInput;
layout(set=0, binding=1, rgba16f) writeonly restrict mediump uniform image3D uOutput;
layout(set = 0, binding = 2) uniform reluBuffer{
ivec4 inImgSize;
ivec4 outImgSize;
vec2 scale;
}uResizeParam;
layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
void main()
{
ivec3 pos = ivec3(gl_GlobalInvocationID);
// input output layout is NC4HW4
ivec3 inputImgSize = uResizeParam.inImgSize.xyz;
ivec3 outputImgSize = uResizeParam.outImgSize.xyz;
if(pos.x < outputImgSize.x && pos.y < outputImgSize.y)
{
float srcX = float(pos.x) * uResizeParam.scale.x;
int x1 = int(floor(srcX));
int x11 = clamp(x1, 0, inputImgSize.x - 1);
float srcY = float(pos.y) * uResizeParam.scale.y;
int y1 = int(floor(srcY));
int y11 = clamp(y1, 0, inputImgSize.y - 1);
vec4 outValue = texelFetch(uInput, ivec3(x11, y11, pos.z), 0);
imageStore(uOutput, pos, outValue);
}
}