UnrealEngine/Engine/Shaders/Private/Lumen/LumenScreenSpaceBentNormal.usf

// Copyright Epic Games, Inc. All Rights Reserved.

#include "../Common.ush"
#include "LumenMaterial.ush"
#include "../DeferredShadingCommon.ush"
#include "../BRDF.ush"
#include "../MonteCarlo.ush"
#include "../BlueNoise.ush"
#include "../ShadingModelsSampling.ush"
#include "../SceneTextureParameters.ush"
#include "../HZB.ush"
#define IS_SSGI_SHADER 1
#include "LumenScreenTracing.ush"
#include "LumenHairTracing.ush"
#include "LumenScreenProbeCommon.ush"
#include "LumenScreenSpaceBentNormal.ush"

#ifndef THREADGROUP_SIZE
#define THREADGROUP_SIZE 1
#endif

#if OUTPUT_BENT_NORMAL
	RWTexture2DArray<uint> RWShortRangeAO;
#else
	RWTexture2DArray<UNORM float> RWShortRangeAO;
#endif

#if DOWNSAMPLE_FACTOR != 1
	RWTexture2D<float> RWDownsampledSceneDepth;
	RWTexture2D<UNORM float3> RWDownsampledSceneWorldNormal;
#endif

float SlopeCompareToleranceScale;
float MaxScreenTraceFraction;

uint2 ShortRangeAOViewMin;
uint2 ShortRangeAOViewSize;

[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)]
void ScreenSpaceShortRangeAOCS(
	uint2 GroupId : SV_GroupID,
	uint2 DispatchThreadId : SV_DispatchThreadID,
	uint2 GroupThreadId : SV_GroupThreadID)
{
	const FLumenMaterialCoord Coord = GetLumenMaterialCoordDownsampled(DispatchThreadId, GroupId, GroupThreadId, DOWNSAMPLE_FACTOR, ShortRangeAOViewMin, ShortRangeAOViewSize);
	const uint3 ShortRangeAOCoord = Coord.DownsampledCoord;

	if (Coord.bIsValid)
	{
		const float2 ScreenUV = (Coord.SvPosition + .5f) * View.BufferSizeAndInvSize.zw;
		const float SceneDepth = CalcSceneDepth(ScreenUV);

		const float DummyMaxRoughnessToTrace = 0.5f;
		const FLumenMaterialData Material = ReadMaterialData(Coord, DummyMaxRoughnessToTrace);

		if (IsValid(Material))
		{
			float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);
			float3x3 TangentBasis = GetTangentBasis(Material.WorldNormal);
			float TraceDistance = MaxScreenTraceFraction * 2.0 * GetScreenRayLengthMultiplierForProjectionType(SceneDepth).x;
			float DepthThresholdScale = GetScreenTraceDepthThresholdScale(ScreenUV);
			uint NumPixelSamples = NUM_PIXEL_RAYS;

			float3 UnoccludedSum = 0;
			float3 BentNormal = 0;
			float AmbientOcclusion = 0.0f;

			UNROLL
			for (uint PixelRayIndex = 0; PixelRayIndex < NumPixelSamples; PixelRayIndex++)
			{
				float2 UniformRandom = BlueNoiseVec2(Coord.SvPosition, (ScreenProbeGatherStateFrameIndex * NumPixelSamples + PixelRayIndex));

				//@todo - other shading models
				float4 HemisphereSample = CosineSampleHemisphere(UniformRandom);
				float3 RayDirection = mul(HemisphereSample.xyz, TangentBasis);
				float DirectionVisible = 1;

				#define TRACE_SCREEN 1
				#if TRACE_SCREEN
				{
					uint NumSteps = 4;
					float StartMipLevel = 0;

					uint2 NoiseCoord = Coord.SvPosition * uint2(NumPixelSamples, 1) + uint2(PixelRayIndex, 0);
					float StepOffset = InterleavedGradientNoise(NoiseCoord + 0.5f, 0.0f);

					float RayRoughness = .2f;
		
					FSSRTCastingSettings CastSettings = CreateDefaultCastSettings();
					CastSettings.bStopWhenUncertain = true;

					bool bHit = false;
					float Level;
					float3 HitUVz;
					bool bRayWasClipped;

					FSSRTRay Ray = InitScreenSpaceRayFromWorldSpace(
						TranslatedWorldPosition, RayDirection,
						/* WorldTMax = */ TraceDistance,
						/* SceneDepth = */ SceneDepth,
						/* SlopeCompareToleranceScale */ SlopeCompareToleranceScale * DepthThresholdScale * (float)NumSteps,
						/* bExtendRayToScreenBorder = */ false,
						/* out */ bRayWasClipped);

					bool bUncertain;
					float3 DebugOutput;

					CastScreenSpaceRay(
						FurthestHZBTexture, FurthestHZBTextureSampler,
						StartMipLevel,
						CastSettings,
						Ray, RayRoughness, NumSteps, StepOffset - .9f,
						HZBUvFactorAndInvFactor, false,
						/* out */ DebugOutput,
						/* out */ HitUVz,
						/* out */ Level,
						/* out */ bHit,
						/* out */ bUncertain);

	#if USE_HAIRSTRANDS_SCREEN
					if (!bHit)
					{
						float3 Hair_DebugOutput;
						float3 Hair_HitUVz;
						float Hair_Level;
						bool Hair_bHit = false;
						bool Hair_bUncertain = bUncertain;

						CastScreenSpaceRay(
							HairStrands.HairOnlyDepthFurthestHZBTexture, FurthestHZBTextureSampler,
							StartMipLevel,
							CastSettings,
							Ray, RayRoughness, NumSteps, StepOffset,
							HZBUvFactorAndInvFactor, false,
							/* out */ Hair_DebugOutput,
							/* out */ Hair_HitUVz,
							/* out */ Hair_Level,
							/* out */ Hair_bHit,
							/* out */ Hair_bUncertain);

						if (Hair_bHit && !Hair_bUncertain)
						{
							DebugOutput = Hair_DebugOutput;
							HitUVz = Hair_HitUVz;
							bHit = Hair_bHit;
							bUncertain = Hair_bUncertain;
						}
					}
	#endif

					bHit = bHit && !bUncertain;

	#if USE_HAIRSTRANDS_VOXEL
					if (!bHit)
					{
						bool bHairHit;
						float HairTransparency;
						float HairHitT;

						TraceHairVoxels(
							NoiseCoord,
							SceneDepth,
							TranslatedWorldPosition,
							RayDirection,
							TraceDistance,
							false,
							bHairHit,
							HairTransparency,
							HairHitT);

						bHit = bHairHit && HairHitT < TraceDistance;
					}
	#endif

					DirectionVisible = bHit ? 0.0f : 1.0f;
				}
				#endif

				UnoccludedSum += RayDirection;
				BentNormal += RayDirection * DirectionVisible;
				AmbientOcclusion += DirectionVisible;
			}

			float NormalizeFactor = length(UnoccludedSum);

			if (NormalizeFactor > 0)
			{
				BentNormal /= NormalizeFactor;
			}

			AmbientOcclusion /= NumPixelSamples;

			// Debug passthrough
			//BentNormal = Material.WorldNormal;

			#if OUTPUT_BENT_NORMAL
				RWShortRangeAO[ShortRangeAOCoord] = PackScreenBentNormal(BentNormal);
			#else
				RWShortRangeAO[ShortRangeAOCoord] = AmbientOcclusion;
			#endif

			#if DOWNSAMPLE_FACTOR != 1
			if (ShortRangeAOCoord.z == 0)
			{
				RWDownsampledSceneDepth[ShortRangeAOCoord.xy] = Material.SceneDepth;
				RWDownsampledSceneWorldNormal[ShortRangeAOCoord.xy] = EncodeNormal(Material.WorldNormal);
			}
			#endif
		}
	}
}