mirror of https://github.com/alibaba/MNN.git
42 lines
1.3 KiB
Plaintext
42 lines
1.3 KiB
Plaintext
#version 440 core
|
|
layout(std140) buffer;
|
|
layout(std430) uniform;
|
|
|
|
layout(set=0, binding=0, rgba16f) writeonly mediump uniform image2D uOutput;
|
|
layout(set=0, binding=1) mediump uniform sampler3D uInput;
|
|
|
|
layout(set=0, binding=2) readonly uniform constBuffer {
|
|
ivec2 pad;
|
|
ivec2 kernelSize;
|
|
ivec2 stride;
|
|
ivec2 dilate;
|
|
ivec4 inputSize;
|
|
ivec4 outputSize;
|
|
int batch;
|
|
int group;
|
|
} uConstant;
|
|
|
|
layout (local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
|
|
|
|
#define UP_DIV(x, y) (((x)+(y)-1)/(y))
|
|
|
|
void main()
|
|
{
|
|
ivec3 pos = ivec3(gl_GlobalInvocationID);
|
|
ivec3 inputSize = uConstant.inputSize.xyz;
|
|
ivec3 outputSize = uConstant.outputSize.xyz;
|
|
|
|
if (pos.x < outputSize.x && pos.y < outputSize.y)
|
|
{
|
|
ivec2 s0 = pos.xy*uConstant.stride-uConstant.pad;
|
|
int fz = pos.z % uConstant.inputSize.z;
|
|
int sb = pos.z / uConstant.inputSize.z;
|
|
int destYOrigin = pos.x + pos.y*uConstant.outputSize.x + sb*uConstant.outputSize.x*uConstant.outputSize.y;
|
|
int destY = destYOrigin / 4;
|
|
int destXOffset = destYOrigin % 4;
|
|
int destX = fz*uConstant.kernelSize.x * uConstant.kernelSize.y;
|
|
vec4 color = texelFetch(uInput, ivec3(s0.x, s0.y, pos.z), 0);
|
|
imageStore(uOutput, ivec2(4*destX+destXOffset, destY), color);
|
|
}
|
|
}
|