wehub-resource-sync
59 lines
1.9 KiB
Plaintext
59 lines
1.9 KiB
Plaintext
#version 440 core
|
|
layout(std430) buffer;
|
|
layout(set=0, binding=0) uniform sampler2D uInput;
|
|
|
|
layout(set=0, binding=1) writeonly buffer sourceBuffer{
|
|
vec4 data[];
|
|
} uOutBuffer;
|
|
|
|
layout(set=0, binding=2) uniform constBuffer{
|
|
ivec4 size; // NCHW
|
|
ivec4 stride; // NCHW
|
|
} uConstant;
|
|
|
|
layout (local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
void main()
|
|
{
|
|
ivec3 posOrigin = ivec3(gl_GlobalInvocationID);
|
|
int channelC4 = (uConstant.size.y + 3) / 4;
|
|
int batchC4 = (uConstant.size.x + 3) / 4;
|
|
int totalChannelC4 = channelC4 * uConstant.size.z * uConstant.size.w;
|
|
if (posOrigin.x < totalChannelC4 * batchC4) {
|
|
ivec2 pos;
|
|
pos.x = posOrigin.x % totalChannelC4;
|
|
pos.y = posOrigin.x / totalChannelC4;
|
|
int startIndex = pos.y * 4 * totalChannelC4 + pos.x;
|
|
vec4 u0 = texelFetch(uInput, ivec2(4 * pos.x + 0, pos.y), 0);
|
|
vec4 u1 = texelFetch(uInput, ivec2(4 * pos.x + 1, pos.y), 0);
|
|
vec4 u2 = texelFetch(uInput, ivec2(4 * pos.x + 2, pos.y), 0);
|
|
vec4 u3 = texelFetch(uInput, ivec2(4 * pos.x + 3, pos.y), 0);
|
|
#ifdef TRANSPOSE
|
|
vec4 v0 = vec4(u0.x, u1.x, u2.x, u3.x);
|
|
vec4 v1 = vec4(u0.y, u1.y, u2.y, u3.y);
|
|
vec4 v2 = vec4(u0.z, u1.z, u2.z, u3.z);
|
|
vec4 v3 = vec4(u0.w, u1.w, u2.w, u3.w);
|
|
#else
|
|
vec4 v0 = u0;
|
|
vec4 v1 = u1;
|
|
vec4 v2 = u2;
|
|
vec4 v3 = u3;
|
|
#endif
|
|
// Debug
|
|
//v0 = vec4(1);
|
|
//v1 = vec4(1);
|
|
//v2 = vec4(1);
|
|
//v3 = vec4(1);
|
|
uOutBuffer.data[startIndex + 0 * totalChannelC4] = v0;
|
|
if (pos.y * 4 + 1 < uConstant.size.x) {
|
|
uOutBuffer.data[startIndex + 1 * totalChannelC4] = v1;
|
|
}
|
|
if (pos.y * 4 + 2 < uConstant.size.x) {
|
|
uOutBuffer.data[startIndex + 2 * totalChannelC4] = v2;
|
|
}
|
|
if (pos.y * 4 + 3 < uConstant.size.x) {
|
|
uOutBuffer.data[startIndex + 3 * totalChannelC4] = v3;
|
|
}
|
|
}
|
|
}
|