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

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

ANKI_SPECIALIZATION_CONSTANT_UVEC2(FB_SIZE, 0u);

#pragma anki start comp

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

const UVec2 WORKGROUP_SIZE = UVec2(8, 8);
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_shadowsTex;
layout(set = 0, binding = 3) uniform texture2D u_momentsTex;
layout(set = 0, binding = 4) uniform texture2D u_historyLengthTex;
layout(set = 0, binding = 5) uniform texture2D u_depthTex;

layout(set = 0, binding = 6) uniform uimage2D u_shadowsImage;
layout(set = 0, binding = 7) uniform image2D u_varianceImage;

layout(std430, push_constant, row_major) uniform b_pc
{
	Mat4 u_invProjMat;
};

const I32 CONVOLUTION_RADIUS = 1;

Vec3 toViewspace(Vec2 uv, F32 depth)
{
	const Vec4 pos4 = u_invProjMat * Vec4(UV_TO_NDC(uv), depth, 1.0);
	const Vec3 pos = pos4.xyz / pos4.w;
	return pos;
}

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

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

	const F32 depth = textureLod(u_depthTex, u_linearAnyClampSampler, uv, 0.0).r;
	if(depth == 1.0)
	{
		// Sky
		imageStore(u_shadowsImage, IVec2(gl_GlobalInvocationID), UVec4(0));
		imageStore(u_varianceImage, IVec2(gl_GlobalInvocationID), Vec4(0.0));
		return;
	}

	const F32 historyLength = textureLod(u_historyLengthTex, u_linearAnyClampSampler, uv, 0.0).r;

	UVec4 outPackedShadowLayers;
	F32 outVariance;

	if(historyLength < 4.0 / RT_SHADOWS_MAX_HISTORY_LENGTH)
	{
		// It's been stable less than 4 frames, need to do some work

		const Vec2 texelSize = 1.0 / Vec2(FB_SIZE);

		// Set the reference sample
		const F32 depthCenter = depth;
		const Vec3 positionCenter = toViewspace(uv, depthCenter);

		// Init the sums
		Vec2 sumMoments = Vec2(0.0);
		F32 sumWeight = 0.0;
		F32 sumShadowLayers[MAX_RT_SHADOW_LAYERS];
		zeroRtShadowLayers(sumShadowLayers);

		// Convolve
		for(I32 offsetx = -CONVOLUTION_RADIUS; offsetx <= CONVOLUTION_RADIUS; offsetx++)
		{
			for(I32 offsety = -CONVOLUTION_RADIUS; offsety <= CONVOLUTION_RADIUS; offsety++)
			{
				const Vec2 sampleUv = uv + Vec2(offsetx, offsety) * texelSize;

				// Set the current sample
				const F32 depthTap = textureLod(u_depthTex, u_linearAnyClampSampler, sampleUv, 0.0).r;
				const Vec3 positionTap = toViewspace(sampleUv, depthTap);

				// Do bilateral
				const F32 w = calculateBilateralWeightViewspacePosition(positionCenter, positionTap, 0.5);

				// Sum
				const Vec2 moments = textureLod(u_momentsTex, u_linearAnyClampSampler, sampleUv, 0.0).xy;
				sumMoments += moments * w;

				F32 shadowLayers[MAX_RT_SHADOW_LAYERS];
				unpackRtShadows(textureLod(u_shadowsTex, u_nearestAnyClampSampler, sampleUv, 0.0), shadowLayers);
				ANKI_UNROLL for(U32 i = 0u; i < MAX_RT_SHADOW_LAYERS; ++i)
				{
					sumShadowLayers[i] += shadowLayers[i] * w;
				}

				sumWeight += w;
			}
		}

		sumWeight += EPSILON;

		ANKI_UNROLL for(U32 i = 0u; i < MAX_RT_SHADOW_LAYERS; ++i)
		{
			sumShadowLayers[i] /= sumWeight;
		}
		sumMoments /= sumWeight;

		outPackedShadowLayers = packRtShadows(sumShadowLayers);

		outVariance = max(0.0, sumMoments.y - sumMoments.x * sumMoments.x);

		// Give the variance a boost for the first frames
		outVariance *= 4.0 / (historyLength * RT_SHADOWS_MAX_HISTORY_LENGTH);
	}
	else
	{
		// Stable for more that 4 frames, passthrough

		outPackedShadowLayers = textureLod(u_shadowsTex, u_nearestAnyClampSampler, uv, 0.0);

		const Vec2 moments = textureLod(u_momentsTex, u_linearAnyClampSampler, uv, 0.0).xy;
		outVariance = max(0.0, moments.y - moments.x * moments.x);
	}

	// Store
	imageStore(u_shadowsImage, IVec2(gl_GlobalInvocationID), outPackedShadowLayers);
	imageStore(u_varianceImage, IVec2(gl_GlobalInvocationID), Vec4(outVariance, 0.0, 0.0, 0.0));
}
#pragma anki end