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

44 lines
1.3 KiB
Plaintext
Raw Normal View History

#version 440 core
2019-04-17 10:49:11 +08:00
layout(std430) uniform;
layout(set=0, binding=0) uniform mediump sampler3D uInput;
layout(set=0, rgba16f, binding=1) writeonly uniform mediump image2D uOutput;
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 = 256, local_size_y = 1, local_size_z = 1) in;
2019-04-17 10:49:11 +08:00
#define UP_DIV(x, y) (((x)+(y)-1)/(y))
void main()
{
ivec3 posTmp = ivec3(gl_GlobalInvocationID);
ivec3 pos;
pos.x = posTmp.x % uConstant.inputSize.x;
int subIndex = posTmp.x / uConstant.inputSize.x;
pos.y = subIndex % uConstant.inputSize.y;
pos.z = subIndex / uConstant.inputSize.y;
2019-04-17 10:49:11 +08:00
ivec3 inputSize = uConstant.inputSize.xyz;
int oz = pos.z % uConstant.inputSize.z;
int ob = pos.z / uConstant.inputSize.z;
if (pos.x < uConstant.inputSize.x && pos.y < uConstant.inputSize.y && ob < uConstant.outputSize.w)
2019-04-17 10:49:11 +08:00
{
int sourceXIndex = pos.x + pos.y*uConstant.inputSize.x + ob*uConstant.inputSize.x*uConstant.inputSize.y;
int sourceX = sourceXIndex / 4;
int sourceY = 4*oz + sourceXIndex % 4;
vec4 color = texelFetch(uInput, pos, 0);
//Tranpose
imageStore(uOutput, ivec2(sourceY, sourceX), color);
}
}