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

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

#pragma anki mutator LAST_PASS 0 1

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 texture2D u_depthTex;
layout(set = 0, binding = 3) uniform utexture2D u_shadowsTex;
layout(set = 0, binding = 4) uniform texture2D u_varianceTex;

layout(set = 0, binding = 5) uniform uimage2D u_shadowsImage;
#if !LAST_PASS
layout(set = 0, binding = 6) uniform image2D u_varianceImage;
#endif

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

const I32 CONVOLUTION_RADIUS = 2;
const F32 KERNEL_WEIGHTS[CONVOLUTION_RADIUS + 1] = F32[3](1.0, 2.0 / 3.0, 1.0 / 6.0);

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;
}

F32 computeShadowsLuma(F32 shadowLayers[MAX_RT_SHADOW_LAYERS])
{
	F32 l = 0.0;
	ANKI_UNROLL for(U32 i = 0u; i < MAX_RT_SHADOW_LAYERS; ++i)
	{
		l += shadowLayers[i];
	}
	return l;
}

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_varianceTex, 0).xy);
	F32 sum = 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)];
			sum += textureLod(u_varianceTex, u_linearAnyClampSampler, newUv, 0.0).r * k;
		}
	}

	return sum;
}

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));
#if !LAST_PASS
		imageStore(u_varianceImage, IVec2(gl_GlobalInvocationID), Vec4(0.0));
#endif
		return;
	}

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

	// Read center luma
	F32 shadowLayers[MAX_RT_SHADOW_LAYERS];
	unpackRtShadows(textureLod(u_shadowsTex, u_nearestAnyClampSampler, uv, 0.0), shadowLayers);
	const F32 refLuma = computeShadowsLuma(shadowLayers);

	// Center variance
	const F32 varianceCenter = computeVarianceCenter(uv);

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

	// Convolve
	const Vec2 texelSize = 1.0 / Vec2(textureSize(u_shadowsTex, 0).xy);
	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;

			// Read shadows
			F32 shadowLayers[MAX_RT_SHADOW_LAYERS];
			unpackRtShadows(textureLod(u_shadowsTex, u_nearestAnyClampSampler, sampleUv, 0.0), shadowLayers);

			// Compute luma weight
			const F32 luma = computeShadowsLuma(shadowLayers);
			const F32 variance = textureLod(u_varianceTex, u_linearAnyClampSampler, sampleUv, 0.0).x;
			const F32 sigmaL = 4.0;
			const F32 lumaDiff = abs(luma - refLuma);
			const F32 wl = min(1.0, exp(-lumaDiff / (sigmaL * sqrt(varianceCenter + 0.001) + EPSILON)));

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

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

			// Include more weights
			w *= wl;
			// w *= KERNEL_WEIGHTS[abs(offsetx)] * KERNEL_WEIGHTS[abs(offsety)];

			// Sum
			ANKI_UNROLL for(U32 i = 0u; i < MAX_RT_SHADOW_LAYERS; ++i)
			{
				sumShadowLayers[i] += shadowLayers[i] * w;
			}

			sumVariance += w * w * variance;
			sumWeight += w;
		}
	}

	// Normalize
	sumWeight += EPSILON;

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

	sumVariance /= (sumWeight * sumWeight);

	// Store
	imageStore(u_shadowsImage, IVec2(gl_GlobalInvocationID), packRtShadows(sumShadowLayers));
#if !LAST_PASS
	imageStore(u_varianceImage, IVec2(gl_GlobalInvocationID), Vec4(sumVariance, 0.0, 0.0, 0.0));
#endif
}
#pragma anki end