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

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

#pragma anki mutator GPU_WAVE_SIZE 16 32 64
#pragma anki mutator RADIANCE_OCTAHEDRON_MAP_SIZE 8 10 12 14 16 18 20
#pragma anki mutator IRRADIANCE_OCTAHEDRON_MAP_SIZE 4 5 6
#pragma anki mutator RT_MATERIAL_FETCH_CLIPMAP 0 1
#pragma anki mutator SPATIAL_RECONSTRUCT_TYPE 0 1
#pragma anki mutator IRRADIANCE_USE_SH_L2 0 1

#pragma anki technique RtMaterialFetch rgen mutators RT_MATERIAL_FETCH_CLIPMAP SPATIAL_RECONSTRUCT_TYPE
#pragma anki technique RtMaterialFetchInlineRt comp mutators
#pragma anki technique PopulateCaches comp mutators RADIANCE_OCTAHEDRON_MAP_SIZE
#pragma anki technique ComputeIrradiance comp mutators GPU_WAVE_SIZE RADIANCE_OCTAHEDRON_MAP_SIZE IRRADIANCE_OCTAHEDRON_MAP_SIZE IRRADIANCE_USE_SH_L2
#pragma anki technique Apply comp mutators SPATIAL_RECONSTRUCT_TYPE
#pragma anki technique ApplyInlineRt comp mutators SPATIAL_RECONSTRUCT_TYPE
#pragma anki technique SpatialReconstruct comp mutators SPATIAL_RECONSTRUCT_TYPE
#pragma anki technique TemporalDenoise comp mutators
#pragma anki technique BilateralDenoise comp mutators
#pragma anki technique VisualizeProbes vert pixel mutators

#include <AnKi/Shaders/Include/GpuSceneTypes.h>
#include <AnKi/Shaders/Functions.hlsl>
#include <AnKi/Shaders/Include/MiscRendererTypes.h>
#include <AnKi/Shaders/ImportanceSampling.hlsl>
#include <AnKi/Shaders/PackFunctions.hlsl>
#include <AnKi/Shaders/FastMathFunctions.hlsl>
#include <AnKi/Shaders/IndirectDiffuseClipmaps.hlsl>
#include <AnKi/Shaders/BilateralFilter.hlsl>
#include <AnKi/Shaders/TemporalAA.hlsl>
#include <AnKi/Shaders/VisibilityAndCollisionFunctions.hlsl>
#include <ThirdParty/SHforHLSL/SH.hlsli>

constexpr F32 kGaussianSigma = 0.55;
constexpr F32 kMaxBilateralSamplesPerDirection = 5.0;
constexpr Bool kLocalLightShadow = false;

struct ClipmapRegion
{
	UVec3 m_probesBegin;
	U32 m_partialUpdate;

	UVec3 m_probeCounts;
	U32 m_probeCount;
};

struct ProbeUpdateConsts
{
	U32 m_clipmapIdx;
	U32 m_radianceOctMapSize; // Have it here as well as well as a mutator. Can't use the mutator cause it will create may raygen variants
	U32 m_rayCountPerTexel; // Ray count per oct map texel
	U32 m_maxProbesToUpdate;

	ClipmapRegion m_clipmapRegion;
};

// ===========================================================================
// RtMaterialFetch and RtMaterialFetchInlineRt                               =
// ===========================================================================
#if(NOT_ZERO(ANKI_TECHNIQUE_RtMaterialFetch) && NOT_ZERO(RT_MATERIAL_FETCH_CLIPMAP)) || NOT_ZERO(ANKI_TECHNIQUE_RtMaterialFetchInlineRt)

#	define CLIPMAP_VOLUME
#	define INCLUDE_ALL
#	include <AnKi/Shaders/RtMaterialFetch.hlsl>

ANKI_FAST_CONSTANTS(ProbeUpdateConsts, g_consts)

#	if ANKI_COMPUTE_SHADER
[numthreads(64, 1, 1)] void main(COMPUTE_ARGS)
#	else
[Shader("raygeneration")] void main()
#	endif
{
#	if ANKI_COMPUTE_SHADER
	const U32 tid = svDispatchThreadId.x;
#	else
	const U32 tid = DispatchRaysIndex().x;
#	endif

	const IndirectDiffuseClipmapConstants idConsts = g_globalRendererConstants.m_indirectDiffuseClipmaps;
	const U32 octMapTexelCount = square(g_consts.m_radianceOctMapSize);

	// Compute probe info. Make sure you shoot coherent rays as much as possible by using the same direction on a specific wave
	U32 probeIdx;
	U32 subRayIdx;
	U32 octMapTexelIdx;
	unflatten3dArrayIndex(octMapTexelCount, g_consts.m_maxProbesToUpdate, g_consts.m_rayCountPerTexel, tid, octMapTexelIdx, probeIdx, subRayIdx);
#	if ANKI_COMPUTE_SHADER
	if(octMapTexelIdx >= octMapTexelCount || probeIdx >= g_consts.m_maxProbesToUpdate || subRayIdx >= g_consts.m_rayCountPerTexel)
	{
		return;
	}
#	endif

	if(g_consts.m_clipmapRegion.m_partialUpdate)
	{
		// Choose every other probe depending on the budget
		const U32 div = g_consts.m_clipmapRegion.m_probeCount / g_consts.m_maxProbesToUpdate;

		probeIdx = g_globalRendererConstants.m_frame + div * probeIdx;
		probeIdx = probeIdx % g_consts.m_clipmapRegion.m_probeCount;
	}

	UVec3 probeId;
	unflatten3dArrayIndex(g_consts.m_clipmapRegion.m_probeCounts.z, g_consts.m_clipmapRegion.m_probeCounts.y,
						  g_consts.m_clipmapRegion.m_probeCounts.x, probeIdx, probeId.z, probeId.y, probeId.x);
	probeId += g_consts.m_clipmapRegion.m_probesBegin;
	probeIdx = probeId.z * idConsts.m_probeCounts.x * idConsts.m_probeCounts.y + probeId.y * idConsts.m_probeCounts.x + probeId.x;

	// Check
	{
		const UVec3 probeIdBegin = g_consts.m_clipmapRegion.m_probesBegin;
		[unroll] for(U32 i = 1; i < 2; ++i)
		{
			ANKI_ASSERT(probeId[i] >= probeIdBegin[i] && probeId[i] < probeIdBegin[i] + g_consts.m_clipmapRegion.m_probeCounts[i]);
		}
	}

	const Vec3 probeSize = idConsts.m_sizes[g_consts.m_clipmapIdx] / idConsts.m_probeCounts;
	const Vec3 probeWorldPos = probeId * probeSize + probeSize * 0.5 + idConsts.m_aabbMins[g_consts.m_clipmapIdx];

	// Generate direction
	const UVec2 radianceOctCoord = UVec2(octMapTexelIdx % g_consts.m_radianceOctMapSize, octMapTexelIdx / g_consts.m_radianceOctMapSize);
	ANKI_ASSERT(all(radianceOctCoord < g_consts.m_radianceOctMapSize));
	const U32 sampleIdx = (g_globalRendererConstants.m_frame * g_consts.m_rayCountPerTexel + subRayIdx) % 16;
	const Vec2 sampleCoord = radianceOctCoord + 0.5 + generateMsaa16x(sampleIdx) / (16.0 * 2.0);
	const HVec3 dir = octahedronDecode(sampleCoord / g_consts.m_radianceOctMapSize);

	// Trace
	const F32 tMax = 1000.0; // TODO
	constexpr U32 traceFlags = RAY_FLAG_FORCE_OPAQUE | RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES;
	GBufferLight<F16> gbuffer = (GBufferLight<F16>)0;
	F32 rayT = 0.0;
	Bool backfacing = false;
#	if ANKI_COMPUTE_SHADER
	const Bool hit = materialRayTraceInlineRt<F16>(probeWorldPos, dir, 0.0, tMax, 1000.0, gbuffer, rayT, backfacing);
#	else
	const Bool hit = materialRayTrace<F16>(probeWorldPos, dir, 0.0, tMax, 1000.0, gbuffer, rayT, backfacing, traceFlags);
#	endif

	HVec3 radiance;
	if(backfacing)
	{
		radiance = HVec3(1.0, 0.0, 1.0);
	}
	else
	{
		const Vec3 hitPos = probeWorldPos + dir * (rayT - 0.01);
		radiance = directLighting<F16>(gbuffer, hitPos, !hit, true, tMax, kLocalLightShadow, traceFlags | RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH);

		// Apply indirect
		constexpr SampleClipmapFlag flags =
			kSampleClipmapFlagBackfacingProbeRejection | kSampleClipmapFlagBiasSamplePointSurfaceNormal | kSampleClipmapFlagUsePreviousFrame;
		const Vec3 irradiance = sampleClipmapIrradiance(hitPos, gbuffer.m_worldNormal, g_globalRendererConstants.m_cameraPosition, idConsts,
														g_linearAnyRepeatSampler, flags);

		radiance += irradiance * gbuffer.m_diffuse / kPi;
	}

	// Store result
	const F32 kMaxDist = 1000.0; // Chose something small and make sure its square doesn't overflow F16
	TEX(g_lightResultTex, UVec2(probeIdx, octMapTexelIdx * g_consts.m_rayCountPerTexel + subRayIdx)) = HVec4(radiance, min(rayT, kMaxDist));
}
#endif

// ===========================================================================
// PopulateCaches                                                            =
// ===========================================================================
#if NOT_ZERO(ANKI_TECHNIQUE_PopulateCaches)
Texture2D<Vec4> g_rtResultTex : register(t0);

RWTexture3D<Vec4> g_radianceVolume : register(u0);
RWTexture3D<Vec4> g_distanceMomentsVolume : register(u1);
RWTexture3D<Vec4> g_probeValidiryVolume : register(u2);

ConstantBuffer<GlobalRendererConstants> g_globalRendererConstants : register(b0);

ANKI_FAST_CONSTANTS(ProbeUpdateConsts, g_consts)

[numthreads(64, 1, 1)] void main(COMPUTE_ARGS)
{
	const IndirectDiffuseClipmapConstants idConsts = g_globalRendererConstants.m_indirectDiffuseClipmaps;
	const U32 clipmapIdx = g_consts.m_clipmapIdx;
	const Vec3 clipmapSize = idConsts.m_sizes[clipmapIdx].xyz;
	const Vec3 prevClipmapAabbMin = idConsts.m_previousFrameAabbMins[clipmapIdx].xyz;
	const U32 octMapTexelCount = square(RADIANCE_OCTAHEDRON_MAP_SIZE);

	U32 probeIdx = svDispatchThreadId.x / octMapTexelCount;
	const U32 octMapTexelIdx = svDispatchThreadId.x % octMapTexelCount;

	if(octMapTexelIdx >= octMapTexelCount || probeIdx >= g_consts.m_maxProbesToUpdate)
	{
		return;
	}

	if(g_consts.m_clipmapRegion.m_partialUpdate)
	{
		// Choose every other probe depending on the budget
		const U32 div = g_consts.m_clipmapRegion.m_probeCount / g_consts.m_maxProbesToUpdate;

		probeIdx = g_globalRendererConstants.m_frame + div * probeIdx;
		probeIdx = probeIdx % g_consts.m_clipmapRegion.m_probeCount;
	}

	UVec3 probeId;
	unflatten3dArrayIndex(g_consts.m_clipmapRegion.m_probeCounts.z, g_consts.m_clipmapRegion.m_probeCounts.y,
						  g_consts.m_clipmapRegion.m_probeCounts.x, probeIdx, probeId.z, probeId.y, probeId.x);
	probeId += g_consts.m_clipmapRegion.m_probesBegin;
	probeIdx = probeId.z * idConsts.m_probeCounts.x * idConsts.m_probeCounts.y + probeId.y * idConsts.m_probeCounts.x + probeId.x;

	// Check
	{
		const UVec3 probeIdBegin = g_consts.m_clipmapRegion.m_probesBegin;
		[unroll] for(U32 i = 1; i < 2; ++i)
		{
			ANKI_ASSERT(probeId[i] >= probeIdBegin[i] && probeId[i] < probeIdBegin[i] + g_consts.m_clipmapRegion.m_probeCounts[i]);
		}
	}

	// Read the result of RT
	HVec3 radiance = 0.0;
	Vec2 moments = 0.0;
	F32 weightSum = 0.0;
	for(U32 subray = 0; subray < g_consts.m_rayCountPerTexel; ++subray)
	{
		HVec4 comp = TEX(g_rtResultTex, UVec2(probeIdx, octMapTexelIdx * g_consts.m_rayCountPerTexel + subray));
		const F32 weight = 1.0 / g_consts.m_rayCountPerTexel;

		if(any(comp.xyz != HVec3(1.0, 0.0, 1.0)))
		{
			radiance += comp.xyz * weight;
			moments += Vec2(comp.w, square(comp.w)) * weight;

			weightSum += weight;
		}
	}

	if(weightSum > 0.0)
	{
		radiance /= weightSum;
		moments /= weightSum;
	}

	// Compute probe info
	const Vec3 probeSize = clipmapSize / idConsts.m_probeCounts;
	const Vec3 probeWorldPos = probeId * probeSize + probeSize * 0.5 + idConsts.m_aabbMins[clipmapIdx].xyz;

	const Bool blendWithHistory =
		g_consts.m_clipmapRegion.m_partialUpdate && all(probeWorldPos > prevClipmapAabbMin) && all(probeWorldPos < prevClipmapAabbMin + clipmapSize);

	UVec3 noOctTexCoord = frac(probeWorldPos / clipmapSize) * idConsts.m_probeCounts;
	noOctTexCoord = min(noOctTexCoord, idConsts.m_probeCounts - 1u);
	noOctTexCoord = noOctTexCoord.xzy;

	// Update the radiance and distance moments volumes
	{
		// Compute oct coord
		const UVec2 radianceOctCoord = UVec2(octMapTexelIdx % RADIANCE_OCTAHEDRON_MAP_SIZE, octMapTexelIdx / RADIANCE_OCTAHEDRON_MAP_SIZE);
		ANKI_ASSERT(all(radianceOctCoord < RADIANCE_OCTAHEDRON_MAP_SIZE));

		UVec3 actualVolumeTexCoord;
		actualVolumeTexCoord.xy = radianceOctCoord + noOctTexCoord * (RADIANCE_OCTAHEDRON_MAP_SIZE + 2) + 1;
		actualVolumeTexCoord.z = noOctTexCoord.z;

		HVec3 avgRadiance = 0.0;
		Vec2 avgMoments = 0.0;
		if(blendWithHistory)
		{
			const F16 blendFactor = 0.1 / 4.0;

			const HVec3 prevValue = TEX(g_radianceVolume, actualVolumeTexCoord).xyz;
			avgRadiance = lerp(prevValue, radiance, blendFactor);

			const Vec2 prevValue2 = TEX(g_distanceMomentsVolume, actualVolumeTexCoord).xy;
			avgMoments = lerp(prevValue2, moments, blendFactor);
		}
		else
		{
			avgRadiance = radiance;
			avgMoments = moments;
		}

		TEX(g_radianceVolume, actualVolumeTexCoord).xyz = avgRadiance;
		TEX(g_distanceMomentsVolume, actualVolumeTexCoord).xy = avgMoments;

		// Set oct borders
		IVec2 borders[3];
		const U32 borderCount = octahedronBorder(RADIANCE_OCTAHEDRON_MAP_SIZE, radianceOctCoord, borders);
		for(U32 i = 0; i < borderCount; ++i)
		{
			IVec3 actualVolumeTexCoord;
			actualVolumeTexCoord.xy = radianceOctCoord + noOctTexCoord * (RADIANCE_OCTAHEDRON_MAP_SIZE + 2) + 1;
			actualVolumeTexCoord.xy += borders[i];
			actualVolumeTexCoord.z = noOctTexCoord.z;

			TEX(g_radianceVolume, actualVolumeTexCoord).xyz = avgRadiance;
			TEX(g_distanceMomentsVolume, actualVolumeTexCoord).xy = avgMoments;
		}
	}

	// Update probe validity
	if(octMapTexelIdx == 0)
	{
		// Loop the directions again
		F32 invalidRayCount = 0.0;
		for(U32 i = 0; i < octMapTexelCount; ++i)
		{
			const U32 subray = 0;
			const HVec3 radiance = TEX(g_rtResultTex, UVec2(probeIdx, i * g_consts.m_rayCountPerTexel + subray));
			if(all(radiance == HVec3(1.0, 0.0, 1.0)))
			{
				invalidRayCount += 1.0;
			}
		}

		F32 valid = 1.0 - min(1.0, invalidRayCount / F32(octMapTexelCount / 4));
		if(blendWithHistory)
		{
			const F32 prev = TEX(g_probeValidiryVolume, noOctTexCoord).x;
			valid = lerp(prev, valid, 0.05);
		}
		TEX(g_probeValidiryVolume, noOctTexCoord).x = valid;
	}
}
#endif

// ===========================================================================
// ComputeIrradiance                                                         =
// ===========================================================================
#if NOT_ZERO(ANKI_TECHNIQUE_ComputeIrradiance)

// Each thread is touching a radiance texel
constexpr U32 kThreadCount = (RADIANCE_OCTAHEDRON_MAP_SIZE * RADIANCE_OCTAHEDRON_MAP_SIZE + 32 - 1) / 32 * 32;

Texture3D<Vec4> g_radianceVolumes[kIndirectDiffuseClipmapCount] : register(t0);

RWTexture3D<Vec4> g_irradianceVolumes[kIndirectDiffuseClipmapCount] : register(u0);
RWTexture3D<Vec4> g_avgIrradianceVolumes[kIndirectDiffuseClipmapCount] : register(u3);

ConstantBuffer<GlobalRendererConstants> g_globalRendererConstants : register(b0);

#	if IRRADIANCE_USE_SH_L2
#		define SH_TYPE SH::L2_F16_RGB
#		define SH_PROJECT_ONTO_LX SH::ProjectOntoL2
#		define SH_TO_L1(x) SH::L2toL1(x)
#	else
#		define SH_TYPE SH::L1_F16_RGB
#		define SH_PROJECT_ONTO_LX SH::ProjectOntoL1
#		define SH_TO_L1(x) (x)
#	endif

groupshared SH_TYPE g_sh[kThreadCount];

struct StoreBorderFunc
{
	IVec3 m_startOfOctCoord;
	Vec3 m_value;
	U32 m_clipmapIdx;

	void operator()(IVec2 offset)
	{
		const IVec3 coord = m_startOfOctCoord + IVec3(offset, 0);
		TEX(g_irradianceVolumes[m_clipmapIdx], coord) = Vec4(m_value, 0.0);
	}
};

// The group services a single probe.
// - Every thread reads a radiance value, converts it to SH and stores is in groupshared
// - Then we do a reduction of all SH
// - Then we use the SH to populate the irradiance
[numthreads(kThreadCount, 1, 1)] void main(COMPUTE_ARGS)
{
	const IndirectDiffuseClipmapConstants idConsts = g_globalRendererConstants.m_indirectDiffuseClipmaps;
	const U32 clipmapIdx = svGroupId.x / idConsts.m_probeCounts.x;
	const UVec3 probeId = UVec3(svGroupId.x % idConsts.m_probeCounts.x, svGroupId.y, svGroupId.z);

	// Read the radiance and put it in SH
	if(svGroupIndex < square(RADIANCE_OCTAHEDRON_MAP_SIZE))
	{
		UVec2 radianceOctCoordLocal;
		radianceOctCoordLocal.y = svGroupIndex / RADIANCE_OCTAHEDRON_MAP_SIZE;
		radianceOctCoordLocal.x = svGroupIndex % RADIANCE_OCTAHEDRON_MAP_SIZE;

		UVec3 radianceTexelCoordStart = probeId.xzy;
		radianceTexelCoordStart.xy *= RADIANCE_OCTAHEDRON_MAP_SIZE + 2;
		radianceTexelCoordStart.xy += 1;

		const Vec2 uv = (radianceOctCoordLocal + 0.5) / RADIANCE_OCTAHEDRON_MAP_SIZE;
		const Vec3 sampleDir = octahedronDecode(uv);

		const UVec3 coord = radianceTexelCoordStart + UVec3(radianceOctCoordLocal, 0);
		const Vec3 radiance = TEX(g_radianceVolumes[clipmapIdx], coord);

		const F16 sampleCountf = square(RADIANCE_OCTAHEDRON_MAP_SIZE);
		const F16 normalization = 1.0 / (sampleCountf * sampleDirectionSpherePdf());

		g_sh[svGroupIndex] = SH_PROJECT_ONTO_LX(HVec3(sampleDir), HVec3(radiance)) * normalization;
	}
	else
	{
		g_sh[svGroupIndex] = SH_TYPE::Zero();
	}

	// Integrate, like parallel prefix sum
	GroupMemoryBarrierWithGroupSync();

	[loop] for(U32 s = kThreadCount / 2u; s > 0u; s >>= 1u)
	{
		if(svGroupIndex < s)
		{
			g_sh[svGroupIndex] = g_sh[svGroupIndex] + g_sh[svGroupIndex + s];
		}

#	if ANKI_PLATFORM_MOBILE
		if(s > WaveGetLaneCount())
		{
			GroupMemoryBarrierWithGroupSync();
		}
#	else
		GroupMemoryBarrierWithGroupSync();
#	endif
	}

	const SH_TYPE sh = g_sh[0];

	// Store the irradiance
	if(svGroupIndex < square(IRRADIANCE_OCTAHEDRON_MAP_SIZE))
	{
		UVec2 octCoordLocal;
		octCoordLocal.y = svGroupIndex / IRRADIANCE_OCTAHEDRON_MAP_SIZE;
		octCoordLocal.x = svGroupIndex % IRRADIANCE_OCTAHEDRON_MAP_SIZE;

		UVec3 texelCoordStart = probeId.xzy;
		texelCoordStart.xy *= IRRADIANCE_OCTAHEDRON_MAP_SIZE + 2;
		texelCoordStart.xy += 1;

		const Vec2 uv = (octCoordLocal + 0.5) / IRRADIANCE_OCTAHEDRON_MAP_SIZE;
		const Vec3 sampleDir = octahedronDecode(uv);

		const Vec3 irradiance = SH::CalculateIrradiance<F16>(sh, sampleDir);

		const UVec3 coord = texelCoordStart + UVec3(octCoordLocal, 0);
		TEX(g_irradianceVolumes[clipmapIdx], coord) = Vec4(irradiance, 0.0);

		// Write the borders
		StoreBorderFunc func;
		func.m_clipmapIdx = clipmapIdx;
		func.m_startOfOctCoord = texelCoordStart;
		func.m_value = irradiance;
		storeOctahedronBorder(IRRADIANCE_OCTAHEDRON_MAP_SIZE, octCoordLocal, func);
	}

	// Store the average irradiance
	HVec3 dir;
	HVec3 color;
	SH::ApproximateDirectionalLight(SH_TO_L1(sh), dir, color);
	if(isInfOrNan(Vec3(color)))
	{
		color = 0.0;
	}

	TEX(g_avgIrradianceVolumes[clipmapIdx], probeId.xzy) = Vec4(color, 0.0);
}
#endif

// ===========================================================================
// Apply                                                                     =
// ===========================================================================
#if NOT_ZERO(ANKI_TECHNIQUE_Apply)
Texture2D<Vec4> g_depthTex : register(t0);
Texture2D<Vec4> g_gbufferRt2 : register(t1);

Texture2D<Vec4> g_blueNoiseTex : register(t2);

RWTexture2D<Vec4> g_outTex : register(u0);

ConstantBuffer<GlobalRendererConstants> g_globalRendererConstants : register(b0);

SamplerState g_linearAnyRepeatSampler : register(s0);

[numthreads(64, 1, 1)] void main(COMPUTE_ARGS)
{
	Vec2 lowTextureSize;
	g_outTex.GetDimensions(lowTextureSize.x, lowTextureSize.y);

	const UVec2 realSvDispatchThreadId = getOptimalDispatchThreadId8x8Amd(svGroupIndex, svGroupId.xy);
#	if SPATIAL_RECONSTRUCT_TYPE == 0
	const Vec2 fullViewportSize = lowTextureSize * Vec2(2.0, 1.0);
	const Vec2 coord = Vec2(realSvDispatchThreadId.x * 2u + (realSvDispatchThreadId.y & 1u), realSvDispatchThreadId.y);
#	else
	const Vec2 fullViewportSize = lowTextureSize * 2.0;
	const Vec2 coord = Vec2(realSvDispatchThreadId * 2u);
#	endif

	if(any(coord >= fullViewportSize))
	{
		return;
	}

	const F32 depth = TEX(g_depthTex, coord).r;
	const Vec2 uv = (coord + 0.5) / fullViewportSize;
	const Vec2 ndc = uvToNdc(uv);
	const Vec4 worldPos4 = mul(g_globalRendererConstants.m_matrices.m_invertedViewProjectionJitter, Vec4(ndc, depth, 1.0));
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;

	const Vec3 normal = unpackNormalFromGBuffer(TEX(g_gbufferRt2, coord));

	// Rand
	UVec2 noiseTexSize;
	g_blueNoiseTex.GetDimensions(noiseTexSize.x, noiseTexSize.y);
	Vec3 noise3 = TEX(g_blueNoiseTex, realSvDispatchThreadId % noiseTexSize);
	noise3 = animateBlueNoise(noise3, g_globalRendererConstants.m_frame);
	const F32 noise = noise3.x;

	const U32 method = 0;
	Vec3 irradiance;
	if(method == 0)
	{
		const SampleClipmapFlag flags = kSampleClipmapFlagFullQuality | kSampleClipmapFlagBiasSamplePointTowardsCamera;
		irradiance = sampleClipmapIrradiance(worldPos, normal, g_globalRendererConstants.m_cameraPosition,
											 g_globalRendererConstants.m_indirectDiffuseClipmaps, g_linearAnyRepeatSampler, flags, noise);
	}
	else
	{
		const SampleClipmapFlag flags = kSampleClipmapFlagFullQuality | kSampleClipmapFlagBiasSamplePointTowardsCamera;
		irradiance = sampleClipmapAvgIrradiance(worldPos, normal, g_globalRendererConstants.m_cameraPosition,
												g_globalRendererConstants.m_indirectDiffuseClipmaps, g_linearAnyRepeatSampler, flags, noise);
	}

	TEX(g_outTex, realSvDispatchThreadId) = Vec4(irradiance, 0.0);
}
#endif

// ===========================================================================
// RtMaterialFetch (Apply)                                                   =
// ===========================================================================
#if(ANKI_TECHNIQUE_RtMaterialFetch && !RT_MATERIAL_FETCH_CLIPMAP) || ANKI_TECHNIQUE_ApplyInlineRt

#	define INCLUDE_ALL
#	include <AnKi/Shaders/RtMaterialFetch.hlsl>

#	define INLINE_RT ANKI_TECHNIQUE_ApplyInlineRt

struct Consts
{
	F32 m_rayMax;
	F32 m_padding1;
	F32 m_padding2;
	F32 m_padding3;

	Vec4 m_padding[2];
};
ANKI_FAST_CONSTANTS(Consts, g_consts)

#	if INLINE_RT
[numthreads(8, 8, 1)] void main(COMPUTE_ARGS)
#	else
[Shader("raygeneration")] void main()
#	endif
{
#	if INLINE_RT
	UVec2 texSize;
	g_colorAndPdfTex.GetDimensions(texSize.x, texSize.y);
	if(any(svDispatchThreadId.xy >= texSize))
	{
		return;
	}
#	else
	const UVec2 svDispatchThreadId = DispatchRaysIndex();
	const UVec2 texSize = DispatchRaysDimensions().xy;
#	endif

#	if SPATIAL_RECONSTRUCT_TYPE == 0
	const UVec2 fullCoord = UVec2(svDispatchThreadId.x * 2u + (svDispatchThreadId.y & 1u), svDispatchThreadId.y);
	const Vec2 uv = (fullCoord + 0.5) / (texSize * UVec2(2, 1));
#	else
	const UVec2 fullCoord = svDispatchThreadId.xy * 2u;
	const Vec2 uv = (fullCoord + 0.5) / (texSize * 2);
#	endif

	const F32 depth = TEX(g_depthTex, fullCoord).x;
	const Vec4 rt2 = TEX(g_gbufferRt2, fullCoord);
	const Vec3 worldNormal = unpackNormalFromGBuffer(rt2);

	const Vec4 v4 = mul(g_globalRendererConstants.m_matrices.m_invertedViewProjectionJitter, Vec4(uvToNdc(uv), depth, 1.0));
	const Vec3 worldPos = v4.xyz / v4.w;

	const Vec3 biasDir = normalize(g_globalRendererConstants.m_cameraPosition - worldPos);
	const Vec3 biasedWorldPos = worldPos + biasDir * 0.1;

	// Rand
	const UVec3 seed = rand3DPCG16(UVec3(fullCoord, g_globalRendererConstants.m_frame % 8u));
	const Vec2 randFactors = hammersleyRandom16(g_globalRendererConstants.m_frame % 64u, 64u, seed);

	const Mat3 tbn = rotationFromDirection(worldNormal);
	const Vec3 rayDir = normalize(mul(tbn, hemisphereSampleCos(randFactors)));

	// Trace
	const F32 tMax = g_consts.m_rayMax;
	constexpr U32 traceFlags = RAY_FLAG_FORCE_OPAQUE | RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES;
	GBufferLight<F16> gbuffer = (GBufferLight<F16>)0;
	F32 rayT = 0.0;
	Bool backfacing = false;
#	if INLINE_RT
	const Bool hit = materialRayTraceInlineRt<F16>(biasedWorldPos, rayDir, 0.01, tMax, 1000.0, gbuffer, rayT, backfacing);
#	else
	const Bool hit = materialRayTrace<F16>(biasedWorldPos, rayDir, 0.01, tMax, 1000.0, gbuffer, rayT, backfacing, traceFlags);
#	endif

	HVec3 radiance = 0.0;
	Vec3 hitPos = 0.0;
	if(hit)
	{
		hitPos = biasedWorldPos + rayDir * (rayT - 0.01);
		radiance = directLighting<F16>(gbuffer, hitPos, !hit, true, 1000.0, kLocalLightShadow, traceFlags | RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH);
	}

	Vec3 rayOrigin;
	Vec3 rayDir2;
	if(hit)
	{
		// 2nd bounce
		rayOrigin = hitPos;
		rayDir2 = gbuffer.m_worldNormal;
	}
	else
	{
		rayOrigin = biasedWorldPos;
		rayDir2 = worldNormal;
	}
	const SampleClipmapFlag flags = kSampleClipmapFlagFullQuality | kSampleClipmapFlagBiasSamplePointSurfaceNormal;
	const Vec3 irradiance =
		sampleClipmapIrradiance(rayOrigin, rayDir2, g_globalRendererConstants.m_cameraPosition, g_globalRendererConstants.m_indirectDiffuseClipmaps,
								g_linearAnyRepeatSampler, flags, randFactors.x);

	Vec3 final;
	if(hit)
	{
		final = radiance + irradiance * gbuffer.m_diffuse;
	}
	else
	{
		final = irradiance;
	}

	TEX(g_colorAndPdfTex, svDispatchThreadId.xy) = Vec4(final, 0.0);
	// TEX(g_colorAndPdfTex, DispatchRaysIndex().xy) = lerp(TEX(g_colorAndPdfTex, DispatchRaysIndex().xy), Vec4(final, 0.0), 0.05);
}
#endif

// ===========================================================================
// SpatialReconstruct                                                        =
// ===========================================================================
#if NOT_ZERO(ANKI_TECHNIQUE_SpatialReconstruct)
Texture2D<Vec3> g_inTex : register(t0);
Texture2D<F32> g_depthTex : register(t1);

RWTexture2D<Vec4> g_outTex : register(u0);

void appendSample(F32 refDepth, F32 sampleDepth, Vec3 sample, inout Vec3 sampleSum, inout F32 weightSum, F32 extraWeight = 1.0)
{
	const F32 weight = calculateBilateralWeightDepth<F32>(refDepth, sampleDepth, 1.0) * extraWeight;
	sampleSum += sample * weight;
	weightSum += weight;
}

void normalizeSum(F32 weightSum, inout Vec3 sampleSum)
{
	if(weightSum > kEpsilonF32 * 10.0)
	{
		sampleSum /= weightSum;
	}
	else
	{
		sampleSum = 0.0;
	}
}

void oneIn4Reconstruct(IVec2 svDispatchThreadId)
{
	IVec2 viewportSize;
	g_outTex.GetDimensions(viewportSize.x, viewportSize.y);
	const IVec2 quarterViewportSize = viewportSize / 2;
	const IVec2 quarterCoord = svDispatchThreadId; // Coord in quarter rez

	// This is the pattern we are trying to fill
	// +---+---+
	// | 0 | 1 |
	// +---+---+
	// | 3 | 2 |
	// +---+---+

	// Gather the color of the neighbours and their depth
	Vec3 samples[2][2];
	F32 sampleDepths[2][2];
	F32 maxLuma = 0.0;
	IVec2 maxLumaPixel = 0;
	[unroll] for(U32 x = 0; x < 2; ++x)
	{
		[unroll] for(U32 y = 0; y < 2; ++y)
		{
			IVec2 coord = quarterCoord + IVec2(x, y);
			coord = min(coord, quarterViewportSize - 1);
			samples[x][y] = TEX(g_inTex, coord);

			const F32 luma = computeLuminance(samples[x][y]);
			if(luma > maxLuma)
			{
				maxLuma = luma;
				maxLumaPixel = IVec2(x, y);
			}

			coord = quarterCoord * 2 + IVec2(x, y);
			sampleDepths[x][y] = TEX(g_depthTex, coord);
		}
	}

	// Remove fireflies
	F32 avgLumaOf3 = 0.0;
	[unroll] for(U32 x = 0; x < 2; ++x)
	{
		[unroll] for(U32 y = 0; y < 2; ++y)
		{
			if(any(maxLumaPixel != IVec2(x, y)))
			{
				const F32 luma = computeLuminance(samples[x][y]);
				avgLumaOf3 += luma / 3.0;
			}
		}
	}

	if(maxLuma > avgLumaOf3 * 5.0)
	{
		// Firefly, tone it down
		samples[maxLumaPixel.x][maxLumaPixel.y] *= avgLumaOf3 / maxLuma;
	}

	// 0 is already filled, just write it
	IVec2 coord = quarterCoord * 2;
	TEX(g_outTex, coord) = Vec4(samples[0][0], 0.0);

	// For 2 use 4 samples
	coord = quarterCoord * 2 + 1;
	F32 refDepth = TEX(g_depthTex, coord);
	Vec3 sampleSum = 0.0;
	F32 weightSum = 0.0;
	appendSample(refDepth, sampleDepths[0][0], samples[0][0], sampleSum, weightSum);
	appendSample(refDepth, sampleDepths[1][0], samples[1][0], sampleSum, weightSum);
	appendSample(refDepth, sampleDepths[1][1], samples[1][1], sampleSum, weightSum);
	appendSample(refDepth, sampleDepths[0][1], samples[0][1], sampleSum, weightSum);
	normalizeSum(weightSum, sampleSum);
	TEX(g_outTex, coord) = Vec4(sampleSum, 0.0);
	const Vec3 sample2 = sampleSum;
	const F32 depth2 = refDepth;

	// For 1 use 3 samples
	coord = quarterCoord * 2 + IVec2(1, 0);
	refDepth = TEX(g_depthTex, coord);
	sampleSum = 0.0;
	weightSum = 0.0;
	appendSample(refDepth, sampleDepths[0][0], samples[0][0], sampleSum, weightSum, 1.0);
	appendSample(refDepth, sampleDepths[1][0], samples[1][0], sampleSum, weightSum, 1.0);
	appendSample(refDepth, depth2, sample2, sampleSum, weightSum, 0.5); // Less weight on that since it's reconstructed
	normalizeSum(weightSum, sampleSum);
	TEX(g_outTex, coord) = Vec4(sampleSum, 0.0);

	// For 4 use 3 samples
	coord = quarterCoord * 2 + IVec2(0, 1);
	refDepth = TEX(g_depthTex, coord);
	sampleSum = 0.0;
	weightSum = 0.0;
	appendSample(refDepth, sampleDepths[0][0], samples[0][0], sampleSum, weightSum, 1.0);
	appendSample(refDepth, sampleDepths[0][1], samples[0][1], sampleSum, weightSum, 1.0);
	appendSample(refDepth, depth2, sample2, sampleSum, weightSum, 0.5); // Less weight on that since it's reconstructed
	normalizeSum(weightSum, sampleSum);
	TEX(g_outTex, coord) = Vec4(sampleSum, 0.0);
}

void checkerboardReconstruct(IVec2 svDispatchThreadId)
{
	IVec2 viewportSize;
	g_outTex.GetDimensions(viewportSize.x, viewportSize.y);

	const IVec2 filledCoord = IVec2(svDispatchThreadId.x * 2 + (svDispatchThreadId.y & 1), svDispatchThreadId.y);
	const IVec2 toBeFilledCoord = IVec2(svDispatchThreadId.x * 2 + ((svDispatchThreadId.y + 1) & 1), svDispatchThreadId.y);

	const F32 refDepth = TEX(g_depthTex, toBeFilledCoord);

	Vec3 toBeFilledColor = 0.0;
	F32 weightSum = 0.0;
	const IVec2 offsets[4] = {IVec2(-1, 0), IVec2(1, 0), IVec2(0, -1), IVec2(0, 1)};
	[unroll] for(U32 i = 0; i < 4; ++i)
	{
		const IVec2 sampleCoord = toBeFilledCoord + offsets[i];
		if(all(sampleCoord >= 0) && all(sampleCoord < viewportSize))
		{
			const F32 sampleDepth = TEX(g_depthTex, sampleCoord);
			const Vec3 sample = TEX(g_inTex, IVec2(sampleCoord.x / 2, sampleCoord.y));

			appendSample(refDepth, sampleDepth, sample, toBeFilledColor, weightSum);

			if(all(sampleCoord == filledCoord))
			{
				TEX(g_outTex, filledCoord) = Vec4(sample, 0.0);
			}
		}
	}

	normalizeSum(weightSum, toBeFilledColor);
	TEX(g_outTex, toBeFilledCoord) = Vec4(toBeFilledColor, 0.0);
}

[numthreads(64, 1, 1)] void main(COMPUTE_ARGS)
{
	const IVec2 realSvDispatchThreadId = getOptimalDispatchThreadId8x8Amd(svGroupIndex, svGroupId.xy);
#	if SPATIAL_RECONSTRUCT_TYPE == 0
	checkerboardReconstruct(realSvDispatchThreadId);
#	else
	oneIn4Reconstruct(realSvDispatchThreadId);
#	endif
}
#endif

// ===========================================================================
// TemporalDenoise                                                           =
// ===========================================================================
#if NOT_ZERO(ANKI_TECHNIQUE_TemporalDenoise)
Texture2D<F32> g_historyLengthTex : register(t0);
Texture2D<Vec2> g_motionVectorsTex : register(t1);
Texture2D<Vec3> g_historyTex : register(t2);
Texture2D<Vec4> g_currentTex : register(t3);

RWTexture2D<Vec4> g_outTex : register(u0);

SamplerState g_linearAnyClampSampler : register(s0);

ConstantBuffer<GlobalRendererConstants> g_globalRendererConsts : register(b0);

[numthreads(64, 1, 1)] void main(COMPUTE_ARGS)
{
	const Vec2 coord = getOptimalDispatchThreadId8x8Amd(svGroupIndex, svGroupId.xy);

	const F32 minBlendFactor = 0.025;
	const F32 maxBlendFactor = 0.9;

	const F32 historyLen = TEX(g_historyLengthTex, coord) * kMaxHistoryLength;

	F32 blendFactor = min(1.0, historyLen / 1.0);
	blendFactor = lerp(maxBlendFactor, minBlendFactor, blendFactor);

	Vec3 outColor = TEX(g_currentTex, coord);
	if(blendFactor > maxBlendFactor * 0.9)
	{
		// Don't accumulate
	}
	else
	{
		Vec2 viewport;
		g_historyLengthTex.GetDimensions(viewport.x, viewport.y);

		const Vec2 uv = (coord + 0.5) / viewport;
		const Vec2 historyUv =
			uv + TEX(g_motionVectorsTex, coord)
			+ (g_globalRendererConsts.m_previousMatrices.m_jitterOffsetNdc - g_globalRendererConsts.m_matrices.m_jitterOffsetNdc) / Vec2(2.0, -2.0);

		const Vec3 history = g_historyTex.SampleLevel(g_linearAnyClampSampler, historyUv, 0.0);

		outColor = lerp(history, outColor, blendFactor);
	}

	TEX(g_outTex, coord) = Vec4(outColor, historyLen);
}
#endif

// ===========================================================================
// BilateralDenoise                                                          =
// ===========================================================================
#if NOT_ZERO(ANKI_TECHNIQUE_BilateralDenoise)
Texture2D<Vec4> g_inTex : register(t0);
Texture2D<F32> g_depthTex : register(t1);

RWTexture2D<Vec4> g_outTex : register(u0);

[numthreads(64, 1, 1)] void main(COMPUTE_ARGS)
{
	const IVec2 coord = getOptimalDispatchThreadId8x8Amd(svGroupIndex, svGroupId.xy);
	IVec2 viewport;
	g_outTex.GetDimensions(viewport.x, viewport.y);

	if(any(coord >= viewport))
	{
		return;
	}

	const F32 refDepth = TEX(g_depthTex, coord);

	F32 weightSum = calculateBilateralWeightDepth<F32>(0.0, 0.0, 1.0); // Highest weight that this function can give

	const Vec4 rgba = TEX(g_inTex, coord);

	const F32 historyLen = rgba.w;

	Vec3 colorSum = rgba.xyz * weightSum;

	const F32 blurFactor = 1.0 - min(1.0, historyLen / 12.0);
	const I32 sampleCount = max(1.0, kMaxBilateralSamplesPerDirection * blurFactor);

	for(I32 x = -sampleCount; x <= sampleCount; ++x)
	{
		for(I32 y = -sampleCount; y <= sampleCount; ++y)
		{
			if(x == 0.0 && y == 0.0)
			{
				continue;
			}

			IVec2 newCoord = coord + IVec2(x, y);
			newCoord = clamp(newCoord, 0, viewport - 1);

			const Vec3 sampleColor = TEX(g_inTex, newCoord);
			const F32 sampleDepth = TEX(g_depthTex, newCoord);

			const F32 depthWeight = calculateBilateralWeightDepth<F32>(refDepth, sampleDepth, 1.0);
			const F32 weight = depthWeight;

			colorSum += sampleColor * weight;
			weightSum += weight;
		}
	}

	colorSum /= weightSum;
	TEX(g_outTex, coord) = Vec4(colorSum, 0.0);
}
#endif

// ===========================================================================
// VisualizeProbes                                                           =
// ===========================================================================
#if NOT_ZERO(ANKI_TECHNIQUE_VisualizeProbes)

struct VertIn
{
	U32 m_svVertexId : SV_VERTEXID;
	U32 m_svInstanceId : SV_INSTANCEID;
};

struct VertOut
{
	Vec4 m_svPosition : SV_POSITION;

	Vec3 m_probeCenter : PROBE_CENTER;
};

struct FragOut
{
	Vec4 m_color : SV_TARGET0;
	F32 m_svDepth : SV_Depth;
};

struct Consts
{
	U32 m_clipmapIdx;
	U32 m_padding1;
	U32 m_padding2;
	U32 m_padding3;
};
ANKI_FAST_CONSTANTS(Consts, g_consts)

ConstantBuffer<GlobalRendererConstants> g_globalRendererConstants : register(b0);

Texture3D<Vec4> g_volume : register(t0);
Texture3D<Vec4> g_probeValidityVolume : register(t1);

SamplerState g_linearAnyRepeatSampler : register(s0);

constexpr F32 kSphereRadius = 0.05;

#	if ANKI_VERTEX_SHADER
// Cube vertex positions indexed via SV_VERTEXID
constexpr Vec3 cubeVertices[8] = {Vec3(-1, -1, -1), Vec3(1, -1, -1), Vec3(1, 1, -1), Vec3(-1, 1, -1),
								  Vec3(-1, -1, 1),  Vec3(1, -1, 1),  Vec3(1, 1, 1),  Vec3(-1, 1, 1)};

// Index order for drawing the cube as a triangle list (36 indices, 12 triangles)
constexpr U32 cubeIndices[36] = {0, 1, 2, 2, 3, 0, 1, 5, 6, 6, 2, 1, 5, 4, 7, 7, 6, 5, 4, 0, 3, 3, 7, 4, 3, 2, 6, 6, 7, 3, 4, 5, 1, 1, 0, 4};

VertOut main(VertIn input)
{
	const Vec3 camPos = g_globalRendererConstants.m_cameraPosition;
	const IndirectDiffuseClipmapConstants idConsts = g_globalRendererConstants.m_indirectDiffuseClipmaps;
	const UVec3 probeCounts = idConsts.m_probeCounts;
	const U32 clipmapIdx = g_consts.m_clipmapIdx;

	UVec3 probeCoord;
	unflatten3dArrayIndex(probeCounts.z, probeCounts.y, probeCounts.x, input.m_svInstanceId, probeCoord.z, probeCoord.y, probeCoord.x);

	const Vec3 probeSize = idConsts.m_sizes[clipmapIdx] / probeCounts;
	const Vec3 probeWorldPos = probeCoord * probeSize + probeSize * 0.5 + idConsts.m_aabbMins[clipmapIdx];

	// Vert pos
	const U32 index = cubeIndices[input.m_svVertexId];
	Vec3 vertPos = cubeVertices[index];
	vertPos *= kSphereRadius;
	vertPos += probeWorldPos;

	VertOut output;
	output.m_svPosition = mul(g_globalRendererConstants.m_matrices.m_viewProjectionJitter, Vec4(vertPos, 1.0));
	output.m_probeCenter = probeWorldPos;

	return output;
}
#	endif // ANKI_VERTEX_SHADER

#	if ANKI_PIXEL_SHADER
FragOut main(VertOut input)
{
	const IndirectDiffuseClipmapConstants idConsts = g_globalRendererConstants.m_indirectDiffuseClipmaps;
	const UVec3 probeCounts = idConsts.m_probeCounts;
	const U32 clipmapIdx = g_consts.m_clipmapIdx;

	FragOut output;

	// Compute the far point
	const Vec2 ndc = uvToNdc(input.m_svPosition.xy / g_globalRendererConstants.m_renderingSize);
	const Vec4 v4 = mul(g_globalRendererConstants.m_matrices.m_invertedViewProjectionJitter, Vec4(ndc, 1.0, 1.0));
	const Vec3 farPoint = v4.xyz / v4.w;

	// Do sphere to view vec collision
	const Vec3 rayDir = normalize(farPoint - g_globalRendererConstants.m_cameraPosition);
	F32 t0, t1;
	const Bool collides = testRaySphere(g_globalRendererConstants.m_cameraPosition, rayDir, input.m_probeCenter, kSphereRadius, t0, t1);

	if(!collides)
	{
		discard;
	}

	const F32 t = min(t0, t1);
	const Vec3 collisionPoint = g_globalRendererConstants.m_cameraPosition + rayDir * t;

	const Vec4 p = mul(g_globalRendererConstants.m_matrices.m_viewProjectionJitter, Vec4(collisionPoint, 1.0));
	output.m_svDepth = p.z / p.w;

	UVec3 texSize;
	g_volume.GetDimensions(texSize.x, texSize.y, texSize.z);

	const Bool hasOctMap = texSize.x != probeCounts.x;

	Vec3 uvw = frac(input.m_probeCenter.xzy / idConsts.m_sizes[clipmapIdx].xzy);
	const UVec3 texelCoord = uvw * probeCounts.xzy;

	if(hasOctMap)
	{
		const U32 octProbeSize = texSize.x / probeCounts.x - 2;

		const Vec3 normal = normalize(collisionPoint - input.m_probeCenter);

		uvw.xy = texelCoord.xy * (octProbeSize + 2);
		uvw.xy += octahedronEncode(normal) * octProbeSize + 1.0;
		uvw.xy /= probeCounts.xz * (octProbeSize + 2);

		uvw.z = (texelCoord.z + 0.5) / probeCounts.y;
	}
	else
	{
		uvw = texelCoord + 0.5;
		uvw /= texSize;
	}

	const Bool valid = TEX(g_probeValidityVolume, texelCoord).x > 0.9;
	Vec3 radiance;
	if(valid)
	{
		radiance = g_volume.SampleLevel(g_linearAnyRepeatSampler, uvw, 0.0).xyz;
	}
	else
	{
		radiance = Vec3(1.0, 0.0, 1.0);
	}

	output.m_color = Vec4(radiance, 0.0);
	return output;
}
#	endif // ANKI_PIXEL_SHADER

#endif