LumenScreenProbeGather:TraceScreen

[numthreads( 8 ,  8 , 1)]
void ScreenProbeTraceScreenTexturesCS(
	uint3 GroupId : SV_GroupID,
	uint3 DispatchThreadId : SV_DispatchThreadID,
	uint3 GroupThreadId : SV_GroupThreadID)
{
	uint ProbeTracingResolution =  0  ? ScreenProbeLightSampleResolutionXY : ScreenProbeTracingOctahedronResolution;
	uint2 ScreenProbeAtlasCoord = DispatchThreadId.xy / ProbeTracingResolution;
	uint2 TraceTexelCoord = DispatchThreadId.xy - ScreenProbeAtlasCoord * ProbeTracingResolution;

	uint ScreenProbeIndex = ScreenProbeAtlasCoord.y * ScreenProbeAtlasViewSize.x + ScreenProbeAtlasCoord.x;

	uint2 ScreenProbeScreenPosition = GetScreenProbeScreenPosition(ScreenProbeIndex);
	uint2 ScreenTileCoord = GetScreenTileCoord(ScreenProbeScreenPosition);

	if (ScreenProbeIndex < GetNumScreenProbes()
		&& ScreenProbeAtlasCoord.x < ScreenProbeAtlasViewSize.x
		&& all(TraceTexelCoord < ProbeTracingResolution))
	{
		float2 ScreenUV = GetScreenUVFromScreenProbePosition(ScreenProbeScreenPosition);
		float SceneDepth = GetScreenProbeDepth(ScreenProbeAtlasCoord);
		uint2 TraceBufferCoord = GetTraceBufferCoord(ScreenProbeAtlasCoord, TraceTexelCoord);

		if (ShouldTraceScreenProbeTexel(TraceBufferCoord, SceneDepth))
		{
			float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);

			float3 RayWorldDirection = 0;
			float TraceDistance = MaxTraceDistance;
			float ConeHalfAngle = 0;

			GetScreenProbeTexelRay(
				TraceBufferCoord,
				TraceTexelCoord,
				ScreenTileCoord,
				TranslatedWorldPosition,
				RayWorldDirection,
				TraceDistance,
				ConeHalfAngle);

			if (GetScreenProbeIsTwoSidedFoliage(ScreenProbeAtlasCoord))
			{
				bool bBackfaceRay = dot(GetScreenProbeNormal(ScreenProbeAtlasCoord), RayWorldDirection) < 0.0f;

				if (bBackfaceRay)
				{
					return;
				}
			}

			float DepthThresholdScale = HasDistanceFieldRepresentation(ScreenUV) ? 1.0f : ScreenTraceNoFallbackThicknessScale;

			{
				bool bCoveredByRadianceCache = false;

					FRadianceCacheCoverage Coverage = GetRadianceCacheCoverage(TranslatedWorldPosition -  LWCToFloat( GetPrimaryView() .PreViewTranslation) , RayWorldDirection, InterleavedGradientNoise(ScreenTileCoord,  (FixedJitterIndex < 0 ? View_StateFrameIndexMod8 : FixedJitterIndex) ));

					if (Coverage.bValid)
					{
						TraceDistance = min(TraceDistance, Coverage.MinTraceDistanceBeforeInterpolation);
					}

				bool bHit;
				bool bUncertain;
				float3 HitUVz;
				float3 LastVisibleHitUVz;
				float Unused;

				TraceScreen(
					TranslatedWorldPosition,
					RayWorldDirection,
					TraceDistance,
					HZBUvFactorAndInvFactor,
					MaxHierarchicalScreenTraceIterations,
					RelativeDepthThickness * DepthThresholdScale,
					NumThicknessStepsToDetermineCertainty,
					MinimumTracingThreadOccupancy,
					bHit,
					bUncertain,
					HitUVz,
					LastVisibleHitUVz,
					Unused);

				float Level = 1;
				bool bWriteDepthOnMiss = true;

				bHit = bHit && !bUncertain;

				bool bFastMoving = false;
				float HitDistance = min(sqrt(ComputeRayHitSqrDistance(TranslatedWorldPosition, HitUVz)), MaxTraceDistance);

				if (bHit)
				{
					float2 HitScreenUV = HitUVz.xy;
					float2 HitScreenPosition = (HitUVz.xy - View_ScreenPositionScaleBias.wz) / View_ScreenPositionScaleBias.xy;
					float HitDeviceZ = HitUVz.z;

					float3 HitHistoryScreenPosition = GetPrevScreenPosition(HitScreenPosition, HitScreenUV, HitDeviceZ);

					float Vignette = min(ComputeHitVignetteFromScreenPos(HitScreenPosition), ComputeHitVignetteFromScreenPos(HitHistoryScreenPosition.xy));
					float Noise = InterleavedGradientNoise(TraceBufferCoord + 0.5f, View_StateFrameIndexMod8);


					if (Vignette < Noise)
					{
						bHit = false;
					}

					if (bHit)
					{

						float PrevDeviceZ = HitHistoryScreenPosition.z;


						float2 HitHistoryScreenUVForDepth = HitHistoryScreenPosition.xy * PrevScreenPositionScaleBiasForDepth.xy + PrevScreenPositionScaleBiasForDepth.zw;
						float HistoryDeviceZ = Texture2DSampleLevel(HistorySceneDepth,  View_SharedPointClampedSampler , HitHistoryScreenUVForDepth, 0).x;

						bHit = abs(HistoryDeviceZ - PrevDeviceZ) < HistoryDepthTestRelativeThickness * lerp(.5f, 2.0f, Noise);
					}

					if (bHit && SkipFoliageHits > 0)
					{
						int ShadingModel = GetGBufferDataFromSceneTextures(HitScreenUV).ShadingModelID;

						if (ShadingModel ==  2  || ShadingModel ==  6 )
						{
							bHit = false;
						}
					}


					if (bHit)
					{
						float2 HitHistoryScreenUV = HitHistoryScreenPosition.xy * PrevScreenPositionScaleBias.xy + PrevScreenPositionScaleBias.zw;
						HitHistoryScreenUV = clamp(HitHistoryScreenUV, PrevSceneColorBilinearUVMin, PrevSceneColorBilinearUVMax);
						float3 Lighting = SampleScreenColor(PrevSceneColorTexture,  View_SharedPointClampedSampler , HitHistoryScreenUV).xyz * PrevSceneColorPreExposureCorrection;

						Lighting += GetSkylightLeaking(RayWorldDirection, HitDistance) * View_PreExposure;

							RWTraceRadiance[TraceBufferCoord] = Lighting;

						{
							float HitSceneDepth = ConvertFromDeviceZ(HitUVz.z);
							float3 HitTranslatedWorldPosition = mul(float4(HitScreenPosition * HitSceneDepth, HitSceneDepth, 1), View_ScreenToTranslatedWorld).xyz;

							float3 HitWorldVelocity = HitTranslatedWorldPosition - GetPrevTranslatedWorldPosition(HitHistoryScreenPosition);

							bFastMoving = IsTraceMoving(TranslatedWorldPosition, SceneDepth, ScreenProbeAtlasCoord, HitTranslatedWorldPosition, HitWorldVelocity);
						}
					}

				}

				if (bHit || bWriteDepthOnMiss)
				{
					WriteTraceHit(TraceBufferCoord, HitDistance, bHit, bFastMoving);
				}
			}
		}
	}
}

bool ShouldTraceScreenProbeTexel(uint2 TraceBufferCoord, float ProbeSceneDepth)
{
	bool bShouldTrace = ProbeSceneDepth > 0.0f;

#if TRACE_LIGHT_SAMPLES
	bool bValidLightSample;
	bool bCastShadow;
	DecodeLightSampleFlags(ScreenProbeLightSampleFlags[TraceBufferCoord], bValidLightSample, bCastShadow);
	bShouldTrace = bShouldTrace && bValidLightSample && bCastShadow;
#endif

	return bShouldTrace;
}

void GetScreenProbeTexelRay(
	uint2 TraceBufferCoord,
	uint2 TraceTexelCoord,
	uint2 ScreenTileCoord,
	float3 TranslatedWorldPosition,
	inout float3 RayDirection,
	inout float TraceDistance,
	inout float ConeHalfAngle)
{
#if TRACE_LIGHT_SAMPLES
	{
		float4 LightSampleDirectionAndTraceDistance = ScreenProbeLightSampleDirection[TraceBufferCoord];
		RayDirection = LightSampleDirectionAndTraceDistance.xyz;
		TraceDistance = LightSampleDirectionAndTraceDistance.w;
	}
#else
	{
		float2 ProbeUV;
		GetProbeTracingUV(TraceBufferCoord, TraceTexelCoord, GetProbeTexelCenter(ScreenTileCoord), 1, ProbeUV, ConeHalfAngle);

		RayDirection = OctahedronToUnitVector(ProbeUV * 2.0 - 1.0);
	}
#endif
}

void TraceScreen(
	float3 RayTranslatedWorldOrigin,
	float3 RayWorldDirection,
	float MaxWorldTraceDistance,
	float4 HZBUvFactorAndInvFactor,
	float MaxIterations,
	float RelativeDepthThickness,
	float NumThicknessStepsToDetermineCertainty,
	uint MinimumTracingThreadOccupancy,
	inout bool bHit,
	inout bool bUncertain,
	inout float3 OutScreenUV,
	inout float3 OutLastVisibleScreenUV,
	inout float OutHitTileZ)
{
	InternalTraceScreen(
		SceneDepthTexture,
		ClosestHZBTexture,
		HZBBaseTexelSize,
		HZBUVToScreenUVScaleBias,
		RayTranslatedWorldOrigin,
		RayWorldDirection,
		MaxWorldTraceDistance,
		HZBUvFactorAndInvFactor,
		MaxIterations,
		RelativeDepthThickness,
		NumThicknessStepsToDetermineCertainty,
		MinimumTracingThreadOccupancy,
		bHit,
		bUncertain,
		OutScreenUV,
		OutLastVisibleScreenUV,
		OutHitTileZ);
} 

void InternalTraceScreen(
	Texture2D InSceneDepthTexture,
	Texture2D InClosestHZBTexture,
	float2 InHZBBaseTexelSize,
	float4 InHZBUVToScreenUVScaleBias,
	float3 RayTranslatedWorldOrigin, 
	float3 RayWorldDirection,
	float MaxWorldTraceDistance,
	float4 HZBUvFactorAndInvFactor,
	float MaxIterations,
	float RelativeDepthThickness,
	float NumThicknessStepsToDetermineCertainty,
	uint MinimumTracingThreadOccupancy,
	inout bool bHit,
	inout bool bUncertain,
	inout float3 OutScreenUV,
	inout float3 OutLastVisibleScreenUV,
	inout float OutHitTileZ)
{
	float3 RayStartScreenUV;
	{
		float4 RayStartClip = mul(float4(RayTranslatedWorldOrigin, 1.0f), View.TranslatedWorldToClip);
		float3 RayStartScreenPosition = RayStartClip.xyz / max(RayStartClip.w, 1.0f);
		RayStartScreenUV = float3((RayStartScreenPosition.xy * float2(0.5f, -0.5f) + 0.5f) * HZBUvFactorAndInvFactor.xy, RayStartScreenPosition.z);
	}
	
	float3 RayEndScreenUV;
	{
		float3 ViewRayDirection = mul(float4(RayWorldDirection, 0.0), View.TranslatedWorldToView).xyz;
		float SceneDepth = mul(float4(RayTranslatedWorldOrigin, 1.0f), View.TranslatedWorldToView).z;
		// Clamps the ray to end at the Z == 0 plane so the end point will be valid in NDC space for clipping
		float RayEndWorldDistance = ViewRayDirection.z < 0.0 ? min(-0.99f * SceneDepth / ViewRayDirection.z, MaxWorldTraceDistance) : MaxWorldTraceDistance;

		float3 RayWorldEnd = RayTranslatedWorldOrigin + RayWorldDirection * RayEndWorldDistance;
		float4 RayEndClip = mul(float4(RayWorldEnd, 1.0f), View.TranslatedWorldToClip);
		float3 RayEndScreenPosition = RayEndClip.xyz / RayEndClip.w;
		RayEndScreenUV = float3((RayEndScreenPosition.xy * float2(0.5f, -0.5f) + 0.5f) * HZBUvFactorAndInvFactor.xy, RayEndScreenPosition.z);

		float2 ScreenEdgeIntersections = LineBoxIntersect(RayStartScreenUV, RayEndScreenUV, float3(0, 0, 0), float3(HZBUvFactorAndInvFactor.xy, 1));

		// Recalculate end point where it leaves the screen
		RayEndScreenUV = RayStartScreenUV + (RayEndScreenUV - RayStartScreenUV) * ScreenEdgeIntersections.y;
	}

	float BaseMipLevel = HZB_TRACE_INCLUDE_FULL_RES_DEPTH ? -1 : 0;
	float MipLevel = BaseMipLevel;

	float3 RayDirectionScreenUV = RayEndScreenUV - RayStartScreenUV;
	float CurrentIntersectionTime = 0;

    // Offset to pick which XY boundary planes to intersect
    float2 FloorOffset = select(RayDirectionScreenUV.xy < 0, 0.0, 1.0);

	float3 RayScreenUV = RayStartScreenUV;

	// Step out of current tile without hit test to avoid self-intersection
	bool bStepOutOfCurrentTile = true;

	if (bStepOutOfCurrentTile)
	{
		float MipLevelForStepOut = MipLevel;
		float2 CurrentMipTexelSize = exp2(MipLevelForStepOut) * InHZBBaseTexelSize;
		float2 CurrentMipResolution = 1.0f / CurrentMipTexelSize;

		float2 UVOffset = .005f * CurrentMipTexelSize;
		UVOffset = select(RayDirectionScreenUV.xy < 0, -UVOffset, UVOffset);

		float2 XYPlane = floor(RayScreenUV.xy * CurrentMipResolution) + FloorOffset;
		XYPlane = XYPlane * CurrentMipTexelSize + UVOffset;
		
		float2 PlaneIntersectionTimes = (XYPlane - RayStartScreenUV.xy) / RayDirectionScreenUV.xy;
		float IntersectionTime = min(PlaneIntersectionTimes.x, PlaneIntersectionTimes.y);
		CurrentIntersectionTime = IntersectionTime;
		RayScreenUV = RayStartScreenUV + CurrentIntersectionTime * RayDirectionScreenUV;
	}

	float NumIterations = 0;
	bHit = false;
	bUncertain = false;
	OutHitTileZ = 0;
	float LastAboveSurfaceTime = CurrentIntersectionTime;

	// Stackless HZB traversal
	while (MipLevel >= BaseMipLevel 
		&& NumIterations < MaxIterations 
		&& CurrentIntersectionTime < 1.0f
#if TERMINATE_ON_LOW_OCCUPANCY
		&& WaveActiveCountBits(true) > MinimumTracingThreadOccupancy
#endif
		)
	{
		float2 CurrentMipTexelSize = exp2(MipLevel) * InHZBBaseTexelSize;
		float2 CurrentMipResolution = 1.0f / CurrentMipTexelSize;

		float2 UVOffset = .005f * CurrentMipTexelSize;
		UVOffset = select(RayDirectionScreenUV.xy < 0, -UVOffset, UVOffset);

		float2 XYPlane = floor(RayScreenUV.xy * CurrentMipResolution) + FloorOffset;
		XYPlane = XYPlane * CurrentMipTexelSize + UVOffset;

		float TileZ;

#if HZB_TRACE_INCLUDE_FULL_RES_DEPTH
		if (MipLevel < 0)
		{
			TileZ = InSceneDepthTexture.SampleLevel(GlobalPointClampedSampler, RayScreenUV.xy * InHZBUVToScreenUVScaleBias.xy + InHZBUVToScreenUVScaleBias.zw, 0).r;
		}
		else
#endif
		{
			TileZ = InClosestHZBTexture.SampleLevel(GlobalPointClampedSampler, RayScreenUV.xy, MipLevel).x;
#if !HZB_TRACE_INCLUDE_FULL_RES_DEPTH
			TileZ *= lerp(.99f, 1.0f, saturate(CurrentIntersectionTime * 10.0f));
#endif
		}

		float3 BoundaryPlanes = float3(XYPlane, TileZ);

		float3 PlaneIntersectionTimes = (BoundaryPlanes - RayStartScreenUV) / RayDirectionScreenUV;
		PlaneIntersectionTimes.z = RayDirectionScreenUV.z < 0 ? PlaneIntersectionTimes.z : 1.0f;
		float IntersectionTime = min3(PlaneIntersectionTimes.x, PlaneIntersectionTimes.y, PlaneIntersectionTimes.z);

		bool bAboveSurface = RayScreenUV.z > TileZ;
		bool bSkippedTile = IntersectionTime != PlaneIntersectionTimes.z && bAboveSurface;

		if (bSkippedTile)
		{
			LastAboveSurfaceTime = IntersectionTime;
		}
		
		CurrentIntersectionTime = bAboveSurface ? IntersectionTime : CurrentIntersectionTime;
		RayScreenUV = RayStartScreenUV + min(CurrentIntersectionTime, 1.0f) * RayDirectionScreenUV;
		MipLevel += bSkippedTile ? 1 : -1;

		NumIterations++;
	}

	if (MipLevel < BaseMipLevel)
	{
		float TileZ;

#if HZB_TRACE_INCLUDE_FULL_RES_DEPTH
		TileZ = InSceneDepthTexture.SampleLevel(GlobalPointClampedSampler, RayScreenUV.xy * InHZBUVToScreenUVScaleBias.xy + InHZBUVToScreenUVScaleBias.zw, 0).r;
#else
		TileZ = InClosestHZBTexture.SampleLevel(GlobalPointClampedSampler, RayScreenUV.xy, 0).x;
#endif

		OutHitTileZ = TileZ;

		float HitSceneDepth = ConvertFromDeviceZ(TileZ);
		float RaySceneDepth = ConvertFromDeviceZ(RayScreenUV.z);

		bHit = (RaySceneDepth - HitSceneDepth) < RelativeDepthThickness * max(HitSceneDepth, .00001f);

		if (!bHit)
		{
			// We went below the surface and couldn't count it as a hit, rewind to the last time we were above
			RayScreenUV = RayStartScreenUV + LastAboveSurfaceTime * RayDirectionScreenUV;
		}
	}

	// Linear steps to determine feature thickness along the ray, to reject hits behind very thin surfaces (grass / hair / foliage)
	if (bHit && !bUncertain && NumThicknessStepsToDetermineCertainty > 0)
	{
		float ThicknessSearchMipLevel = 0.0f;
		float MipNumTexels = exp2(ThicknessSearchMipLevel);
		float2 HZBTileSize = MipNumTexels * InHZBBaseTexelSize;
		float NumSteps = NumThicknessStepsToDetermineCertainty / MipNumTexels;
		float ThicknessSearchEndTime = min(length(RayDirectionScreenUV.xy * HZBTileSize * NumSteps) / length(RayEndScreenUV.xy - RayScreenUV.xy), 1.0f);

		for (float I = 0; I < NumSteps; I++)
		{
			float3 SampleUV = RayScreenUV + (I / NumSteps) * ThicknessSearchEndTime * (RayEndScreenUV - RayScreenUV);

			if (all(and(SampleUV.xy > 0, SampleUV.xy < HZBUvFactorAndInvFactor.xy)))
			{
				float SampleTileZ = InClosestHZBTexture.SampleLevel(GlobalPointClampedSampler, SampleUV.xy, ThicknessSearchMipLevel).x;

				if (SampleUV.z > SampleTileZ)
				{
					bUncertain = true;
				}
			}
		}
	}

	// Debug - visualize clipped endpoints
	//RayScreenUV = RayEndScreenUV;
	//bHit = bUncertain;

	// Output in GBuffer SceneUV space for xy
	OutScreenUV.xy = RayScreenUV.xy * InHZBUVToScreenUVScaleBias.xy + InHZBUVToScreenUVScaleBias.zw;
	OutScreenUV.z = RayScreenUV.z;

	float3 LastVisibleScreenUV = RayStartScreenUV + LastAboveSurfaceTime * RayDirectionScreenUV;
	OutLastVisibleScreenUV.xy = LastVisibleScreenUV.xy * InHZBUVToScreenUVScaleBias.xy + InHZBUVToScreenUVScaleBias.zw;
	OutLastVisibleScreenUV.z = LastVisibleScreenUV.z;
}

float4 SampleScreenColor(Texture2D Texture, SamplerState Sampler, float2 UV)
{
	float4 OutColor;

	OutColor.rgb = Texture.SampleLevel( Sampler, UV, 0 ).rgb;
	// Transform NaNs to black, transform negative colors to black.
	OutColor.rgb = -min(-OutColor.rgb, 0.0);
	OutColor.a = 1;
	
	return OutColor;
}