anki-3d-engine/AnKi/Shaders/RtShadowsDenoise.ankiprog

// Copyright (C) 2009-2022, Panagiotis Christopoulos Charitos and contributors.
// All rights reserved.
// Code licensed under the BSD License.
// http://www.anki3d.org/LICENSE

#pragma anki mutator BLUR_ORIENTATION 0 1

#pragma anki start comp

ANKI_SPECIALIZATION_CONSTANT_UVEC2(OUT_IMAGE_SIZE, 0u);
ANKI_SPECIALIZATION_CONSTANT_U32(MIN_SAMPLE_COUNT, 2u);
ANKI_SPECIALIZATION_CONSTANT_U32(MAX_SAMPLE_COUNT, 3u);

#include <AnKi/Shaders/BilateralFilter.glsl>
#include <AnKi/Shaders/PackFunctions.glsl>
#include <AnKi/Shaders/RtShadows.glsl>
#include <AnKi/Shaders/Functions.glsl>

const UVec2 WORKGROUP_SIZE = UVec2(8u, 8u);
layout(local_size_x = WORKGROUP_SIZE.x, local_size_y = WORKGROUP_SIZE.y, local_size_z = 1) in;

layout(set = 0, binding = 0) uniform sampler u_nearestAnyClampSampler;
layout(set = 0, binding = 1) uniform sampler u_linearAnyClampSampler;
layout(set = 0, binding = 2) uniform utexture2D u_inTex;
layout(set = 0, binding = 3) uniform texture2D u_depthTex;
layout(set = 0, binding = 4) uniform texture2D u_gbuffer2Tex;
layout(set = 0, binding = 5) uniform texture2D u_momentsTex;
layout(set = 0, binding = 6) uniform texture2D u_historyLengthTex;
layout(set = 0, binding = 7) writeonly uniform uimage2D u_outImg;

layout(push_constant, row_major, std140) uniform b_pc
{
	RtShadowsDenoiseUniforms u_unis;
};

Vec3 unproject(Vec2 ndc, F32 depth)
{
	const Vec4 worldPos4 = u_unis.invViewProjMat * Vec4(ndc, depth, 1.0);
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;
	return worldPos;
}

F32 computeVarianceCenter(Vec2 uv)
{
	const F32 kernel[2][2] = {{1.0 / 4.0, 1.0 / 8.0}, {1.0 / 8.0, 1.0 / 16.0}};
	const I32 radius = 1;
	const Vec2 texelSize = 1.0 / Vec2(textureSize(u_momentsTex, 0).xy);
	Vec2 sumMoments = Vec2(0.0);

	for(I32 yy = -radius; yy <= radius; yy++)
	{
		for(I32 xx = -radius; xx <= radius; xx++)
		{
			const Vec2 newUv = uv + Vec2(xx, yy) * texelSize;
			const F32 k = kernel[abs(xx)][abs(yy)];
			sumMoments += textureLod(u_momentsTex, u_linearAnyClampSampler, newUv, 0.0).xy * k;
		}
	}

	return abs(sumMoments.y - sumMoments.x * sumMoments.x);
}

void main()
{
	if(skipOutOfBoundsInvocations(WORKGROUP_SIZE, OUT_IMAGE_SIZE))
	{
		return;
	}

	const Vec2 uv = (Vec2(gl_GlobalInvocationID.xy) + 0.5) / Vec2(OUT_IMAGE_SIZE);

	// Reference
	const F32 depthCenter = textureLod(u_depthTex, u_linearAnyClampSampler, uv, 0.0).r;
	if(depthCenter == 1.0)
	{
		imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy), UVec4(0));
		return;
	}

	const Vec3 positionCenter = unproject(UV_TO_NDC(uv), depthCenter);
	const Vec3 normalCenter = unpackNormalFromGBuffer(textureLod(u_gbuffer2Tex, u_linearAnyClampSampler, uv, 0.0));

	F32 shadowFactors[MAX_RT_SHADOW_LAYERS];
	unpackRtShadows(textureLod(u_inTex, u_nearestAnyClampSampler, uv, 0.0), shadowFactors);

	// Decide the amount of blurring
	const F32 varianceCenter = computeVarianceCenter(uv);
	const F32 historyLength =
		textureLod(u_historyLengthTex, u_linearAnyClampSampler, uv, 0.0).x * RT_SHADOWS_MAX_HISTORY_LENGTH;

	U32 sampleCount;
	if(historyLength < 2.0)
	{
		// Worst case
		sampleCount = MAX_SAMPLE_COUNT;
	}
	else if(historyLength > 4.0 && varianceCenter < 0.0001)
	{
		// Best case
		sampleCount = MIN_SAMPLE_COUNT;
	}
	else
	{
		// Every other case

		F32 blur = varianceCenter * 100.0;
		blur = min(1.0, blur);

		const F32 sampleCountf = mix(F32(MIN_SAMPLE_COUNT), F32(MAX_SAMPLE_COUNT), blur);

		sampleCount = U32(sampleCountf);
	}

	sampleCount = sampleCount / 2u;

	// Sample
	F32 weight = 1.0; // TODO fix that

	for(I32 i = -I32(sampleCount); i < I32(sampleCount); ++i)
	{
		const Vec2 texelSize = 1.0 / Vec2(OUT_IMAGE_SIZE);
#if BLUR_ORIENTATION == 0
		const Vec2 sampleUv = Vec2(uv.x + F32(i) * texelSize.x, uv.y);
#else
		const Vec2 sampleUv = Vec2(uv.x, uv.y + F32(i) * texelSize.y);
#endif

		F32 localShadowFactors[MAX_RT_SHADOW_LAYERS];
		unpackRtShadows(textureLod(u_inTex, u_nearestAnyClampSampler, sampleUv, 0.0), localShadowFactors);

		const F32 depthTap = textureLod(u_depthTex, u_linearAnyClampSampler, sampleUv, 0.0).r;
		const Vec3 positionTap = unproject(UV_TO_NDC(sampleUv), depthTap);
		const Vec3 normalTap =
			unpackNormalFromGBuffer(textureLod(u_gbuffer2Tex, u_linearAnyClampSampler, sampleUv, 0.0));

		// F32 w = calculateBilateralWeighPlane(depthCenter, depthTap, 1.0);
		F32 w = calculateBilateralWeightPlane(positionCenter, normalCenter, positionTap, normalTap, 1.0);

		w *= gaussianWeight(0.4, abs(F32(i)) / F32(sampleCount + 1u));

		ANKI_UNROLL for(U32 i = 0u; i < MAX_RT_SHADOW_LAYERS; ++i)
		{
			shadowFactors[i] += localShadowFactors[i] * w;
		}

		weight += w;
	}

	// Write value
	ANKI_UNROLL for(U32 i = 0u; i < MAX_RT_SHADOW_LAYERS; ++i)
	{
		shadowFactors[i] /= weight;
	}

	imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy), packRtShadows(shadowFactors));
}

#pragma anki end