2020-03-14 22:36:08 +08:00
|
|
|
#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;
|
|
|
|
|
2020-03-14 22:36:08 +08:00
|
|
|
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()
|
|
|
|
{
|
2020-03-14 22:36:08 +08:00
|
|
|
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;
|
|
|
|
|
2020-03-14 22:36:08 +08:00
|
|
|
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);
|
|
|
|
}
|
|
|
|
}
|