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;
    }
}