wehub-resource-sync
84 lines
2.6 KiB
Plaintext
84 lines
2.6 KiB
Plaintext
layout(set=0, binding=0) readonly buffer sourceBuffer{
|
|
FLOAT data[];
|
|
} uInBuffer;
|
|
|
|
|
|
layout(set=0, binding=1) writeonly buffer destBuffer{
|
|
FLOAT4 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 posTmp = ivec3(gl_GlobalInvocationID);
|
|
int channelC4 = (uConstant.size.y + 3) / 4;
|
|
if (posTmp.x < channelC4 * uConstant.size.x * uConstant.size.z * uConstant.size.w)
|
|
{
|
|
ivec3 pos;
|
|
pos.z = posTmp.x % (channelC4 * uConstant.size.x);
|
|
int subIndex = posTmp.x / (channelC4 * uConstant.size.x);
|
|
pos.x = subIndex % uConstant.size.w;
|
|
pos.y = subIndex / uConstant.size.w;
|
|
int batchIndex = pos.z / channelC4;
|
|
int zDiv4 = pos.z % channelC4;
|
|
|
|
int lastZ = uConstant.size.y / 4;
|
|
int cIndex = uConstant.size.y % 4;
|
|
|
|
FLOAT4 color = FLOAT4(0.0);
|
|
int z = zDiv4*4;
|
|
int basicOffset = 0
|
|
+ batchIndex*uConstant.stride.x
|
|
+ z * uConstant.stride.y
|
|
+ pos.y * uConstant.stride.z
|
|
+ pos.x * uConstant.stride.w
|
|
;
|
|
|
|
if(zDiv4 == lastZ)
|
|
{
|
|
if(cIndex == 1)
|
|
{
|
|
color.r = uInBuffer.data[basicOffset+0];
|
|
color.g = FLOAT(0.0);
|
|
color.b = FLOAT(0.0);
|
|
color.a = FLOAT(0.0);
|
|
}
|
|
else if(cIndex == 2)
|
|
{
|
|
color.r = uInBuffer.data[basicOffset+0];
|
|
color.g = uInBuffer.data[basicOffset+1*uConstant.stride.y];
|
|
color.b = FLOAT(0.0);
|
|
color.a = FLOAT(0.0);
|
|
}
|
|
else
|
|
{
|
|
color.r = uInBuffer.data[basicOffset+0];
|
|
color.g = uInBuffer.data[basicOffset+1*uConstant.stride.y];
|
|
color.b = uInBuffer.data[basicOffset+2*uConstant.stride.y];
|
|
color.a = FLOAT(0.0);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
color.r = uInBuffer.data[basicOffset+0];
|
|
color.g = uInBuffer.data[basicOffset+1*uConstant.stride.y];
|
|
color.b = uInBuffer.data[basicOffset+2*uConstant.stride.y];
|
|
color.a = uInBuffer.data[basicOffset+3*uConstant.stride.y];
|
|
}
|
|
|
|
uOutBuffer.data[0
|
|
+ pos.x
|
|
+ uConstant.size.w*pos.y
|
|
+ uConstant.size.x * uConstant.size.z*uConstant.size.w*zDiv4
|
|
+ batchIndex*uConstant.size.z*uConstant.size.w
|
|
] = color;
|
|
}
|
|
}
|