mirror of https://github.com/alibaba/MNN.git
46 lines
1.0 KiB
Plaintext
46 lines
1.0 KiB
Plaintext
#version 440 core
|
|
layout(std140) buffer;
|
|
layout(set = 0, binding = 0, rgba16f) uniform image3D uOutput;
|
|
layout(set = 0, binding = 1) uniform sampler3D uInput;
|
|
|
|
layout(set = 0, binding = 2) uniform constBuffer {
|
|
int w;
|
|
int h;
|
|
int c;
|
|
}uConst;
|
|
|
|
layout(local_size_x = 1, local_size_y = 8, local_size_z = 1) in;
|
|
|
|
|
|
void main()
|
|
{
|
|
// input tensor's layout is NC4HW4
|
|
ivec3 pos = ivec3(gl_GlobalInvocationID);
|
|
|
|
int channelDiv4 = (uConst.c + 3) / 4;
|
|
int HW = uConst.w * uConst.h;
|
|
|
|
if(pos.y < uConst.h && pos.z < channelDiv4)
|
|
{
|
|
// get the max value
|
|
vec4 maxValue = vec4(-1000.0);
|
|
for(int i = 0; i < uConst.w; ++i)
|
|
{
|
|
maxValue = max(maxValue, texelFetch(uInput, ivec3(i, pos.y, pos.z), 0));
|
|
}
|
|
|
|
// sum
|
|
vec4 sum = vec4(0.0);
|
|
for(int i = 0; i < uConst.w; ++i)
|
|
{
|
|
sum += exp(texelFetch(uInput, ivec3(i, pos.y, pos.z), 0) - maxValue);
|
|
}
|
|
// div
|
|
for(int i = 0; i < uConst.w; ++i)
|
|
{
|
|
ivec3 curPos = ivec3(i, pos.y, pos.z);
|
|
imageStore(uOutput, curPos, exp(texelFetch(uInput, ivec3(i, pos.y, pos.z), 0) - maxValue) / sum);
|
|
}
|
|
|
|
}
|
|
} |