MNN/source/backend/vulkan/execution/glsl/normalizeChannel.comp

#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 imgSize;
	int channelDiv4;
	float eps;
}uNormalizeParam;

layout(local_size_x = 8, local_size_y = 8) in;

void main()
{
	ivec3 pos = ivec3(gl_GlobalInvocationID);

	if(all(lessThan(pos, uNormalizeParam.imgSize.xyz)))
	{
		vec4 color = texelFetch(uInput, ivec3(pos.x, pos.y, 0), 0);
		vec4 sum = color * color;
		for(int i = 1; i < uNormalizeParam.channelDiv4; ++i)
		{
			color = texelFetch(uInput, ivec3(pos.x, pos.y, i), 0);
			sum += color * color;
		}

		float summerResult = inversesqrt((sum.x + sum.y + sum.z + sum.w) + uNormalizeParam.eps);
		
		for(int i = 0; i < uNormalizeParam.channelDiv4; ++i)
		{
			vec4 tempSum = vec4(summerResult);
			ivec3 curPos = ivec3(pos.x, pos.y, i);
			imageStore(uOutput, curPos, texelFetch(uInput, curPos, 0) * tempSum);
		}

	}

}