mirror of https://github.com/alibaba/MNN.git
27 lines
841 B
Plaintext
27 lines
841 B
Plaintext
|
#version 440 core
|
||
|
|
layout(std140) buffer;
|
||
|
|
layout(set=0, binding=0, rgba16f) writeonly restrict uniform image3D uOutput;
|
||
|
|
layout(set=0, binding=1) uniform sampler3D uInput;
|
||
|
|
|
||
|
|
layout(set=0, binding=2) uniform constBuffer {
|
||
|
|
ivec4 inImageSize;
|
||
|
|
ivec4 outImageSize;
|
||
|
|
ivec4 offset; // w, h, c, 0
|
||
|
|
}uConst;
|
||
|
|
|
||
|
|
layout (local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
|
||
|
|
|
||
|
|
void main()
|
||
|
|
{
|
||
|
|
ivec3 pos = ivec3(gl_GlobalInvocationID);
|
||
|
|
ivec3 imgSize = uConst.inImageSize.xyz;
|
||
|
|
if (pos.x < imgSize.x && pos.y < imgSize.y)
|
||
|
|
{
|
||
|
|
int batchIndex = pos.z / imgSize.z;
|
||
|
|
int c4Index = pos.z % imgSize.z;
|
||
|
|
int outPosZ = batchIndex * uConst.outImageSize.z + uConst.offset.z + c4Index;
|
||
|
|
vec4 color = texelFetch(uInput, pos, 0);
|
||
|
|
imageStore(uOutput, ivec3(pos.xy + uConst.offset.xy, outPosZ), color);
|
||
|
|
}
|
||
|
|
}
|