wehub-resource-sync
77 lines
2.4 KiB
Plaintext
77 lines
2.4 KiB
Plaintext
layout(set=0, binding=0) writeonly buffer destBuffer{
|
|
FLOAT data[];
|
|
} uOutput;
|
|
|
|
|
|
layout(set=0, binding=1) readonly buffer sourceBuffer{
|
|
FLOAT4 data[];
|
|
}uInput;
|
|
|
|
|
|
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);
|
|
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
|
|
;
|
|
|
|
color = uInput.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
|
|
];
|
|
if(zDiv4 == lastZ)
|
|
{
|
|
if(cIndex == 1)
|
|
{
|
|
uOutput.data[basicOffset+0*uConstant.stride.y] = color.r;
|
|
}
|
|
else if(cIndex == 2)
|
|
{
|
|
uOutput.data[basicOffset+0*uConstant.stride.y] = color.r;
|
|
uOutput.data[basicOffset+1*uConstant.stride.y] = color.g;
|
|
}
|
|
else
|
|
{
|
|
uOutput.data[basicOffset+0*uConstant.stride.y] = color.r;
|
|
uOutput.data[basicOffset+1*uConstant.stride.y] = color.g;
|
|
uOutput.data[basicOffset+2*uConstant.stride.y] = color.b;
|
|
}
|
|
}
|
|
else
|
|
{
|
|
uOutput.data[basicOffset+0*uConstant.stride.y] = color.r;
|
|
uOutput.data[basicOffset+1*uConstant.stride.y] = color.g;
|
|
uOutput.data[basicOffset+2*uConstant.stride.y] = color.b;
|
|
uOutput.data[basicOffset+3*uConstant.stride.y] = color.a;
|
|
}
|
|
}
|
|
}
|