UnrealEngineUWP/Engine/Shaders/Private/RayTracing/RayTracingDebug.usf

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

#include "../Common.ush"
#include "../ShadingCommon.ush"
#include "../ColorMap.ush"

#include "RayTracingCommon.ush"
#include "RayTracingHitGroupCommon.ush"

#include "/Engine/Shared/RayTracingDebugDefinitions.h"

RWTexture2D<float4> Output;
RaytracingAccelerationStructure TLAS;

uint VisualizationMode;
uint ShouldUsePreExposure;
uint OpaqueOnly;
float TimingScale;

float MaxTraceDistance;
float FarFieldMaxTraceDistance;
float3 FarFieldReferencePos;

// TODO: move MurmurMix and IntToColor into a common header, shared with Nanite/NaniteVisualize.usf

uint MurmurMix(uint Hash)
{
	Hash ^= Hash >> 16;
	Hash *= 0x85ebca6b;
	Hash ^= Hash >> 13;
	Hash *= 0xc2b2ae35;
	Hash ^= Hash >> 16;
	return Hash;
}

float3 IntToColor(uint Index)
{
	uint Hash = MurmurMix(Index);

	float3 Color = float3
	(
		(Hash >> 0) & 255,
		(Hash >> 8) & 255,
		(Hash >> 16) & 255
	);

	return Color * (1.0f / 255.0f);
}

float4 DebugRayTracingMaterial(RayDesc Ray, bool bIsFarField)
{
	uint2 PixelCoord = DispatchRaysIndex().xy + View.ViewRectMin.xy;

	FRayCone RayCone = (FRayCone)0;
	RayCone.SpreadAngle = View.EyeToPixelSpreadAngle;

	const uint RayFlags = RAY_FLAG_CULL_BACK_FACING_TRIANGLES;
	const uint InstanceInclusionMask = (bIsFarField) ? RAY_TRACING_MASK_FAR_FIELD : (OpaqueOnly ? RAY_TRACING_MASK_OPAQUE : RAY_TRACING_MASK_ALL);
	const bool bEnableSkyLightContribution = true;
	const bool bIgnoreTranslucentMaterials = false;

#if PLATFORM_SUPPORTS_SHADER_TIMESTAMP
	uint64_t TimeBegin = GetShaderTimestamp();
#endif

	if (bIsFarField)
	{
		Ray.TMin = MaxTraceDistance;
		Ray.TMax = FarFieldMaxTraceDistance;
		Ray.Origin += FarFieldReferencePos;
	}
	else
	{
		Ray.TMax = (MaxTraceDistance > 0.0f) ? MaxTraceDistance : Ray.TMax;
	}
	
	FMaterialClosestHitPayload Payload = TraceMaterialRay(
		TLAS,
		RayFlags,
		InstanceInclusionMask,
		Ray,
		RayCone,
		PixelCoord,
		bEnableSkyLightContribution,
		bIgnoreTranslucentMaterials);

#if PLATFORM_SUPPORTS_SHADER_TIMESTAMP
	uint64_t TimeEnd = GetShaderTimestamp();
#endif

	// HitT is packed into the alpha for far field compositing
	float4 Result = (RAY_TRACING_DEBUG_VIZ_FAR_FIELD) ? float4(0, 0, 0, -1) : float4(0, 0, 0, 1);

	if (Payload.IsHit())
	{
		switch (VisualizationMode)
		{
		default:
		case RAY_TRACING_DEBUG_VIZ_RADIANCE:
			Result = float4(Payload.Radiance, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_WORLD_NORMAL:
			Result = float4(Payload.WorldNormal * 0.5 + 0.5, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_OPACITY:
			Result = float4(1.0f - Payload.Opacity, 1.0f - Payload.Opacity, 1.0f - Payload.Opacity, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_BLENDING_MODE:
			Result = float4(Payload.BlendingMode, Payload.BlendingMode, Payload.BlendingMode, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_LIGHTING_CHANNEL_MASK:
			Result = float4(Payload.PrimitiveLightingChannelMask & 0x1, Payload.PrimitiveLightingChannelMask & 0x2, Payload.PrimitiveLightingChannelMask & 0x4, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_INDIRECT_IRRADIANCE:
			Result = float4(Payload.IndirectIrradiance, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_WORLD_POSITION:
			Result = float4(Payload.TranslatedWorldPos, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_HITKIND:
			Result = Payload.IsFrontFace() ? float4(0.0, 1.0, 0.0, 1.0f) : float4(1.0, 0.0, 0.0, 1.0f);
			break;
		#if PLATFORM_SUPPORTS_SHADER_TIMESTAMP
		case RAY_TRACING_DEBUG_VIZ_PERFORMANCE:
			Result.rgb = ColorMapInferno(float(TimeEnd - TimeBegin) * TimingScale);
			break;
		#endif // PLATFORM_SUPPORTS_SHADER_TIMESTAMP

		// STRATA_TODO
	#if !STRATA_ENABLED
		case RAY_TRACING_DEBUG_VIZ_FAR_FIELD:
			Result = float4(Payload.BaseColor.rgb, Payload.HitT);
			break;
		case RAY_TRACING_DEBUG_VIZ_GBUFFER_AO:
			Result = float4(Payload.GBufferAO, Payload.GBufferAO, Payload.GBufferAO, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_SHADING_MODEL:
			Result = float4(GetShadingModelColor(Payload.ShadingModelID), 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_BASE_COLOR:
			Result = float4(Payload.BaseColor.rgb, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_DIFFUSE_COLOR:
			Result = float4(Payload.DiffuseColor.rgb, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_SPECULAR_COLOR:
			Result = float4(Payload.SpecularColor.rgb, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_METALLIC:
			Result = float4(Payload.Metallic, Payload.Metallic, Payload.Metallic, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_SPECULAR:
			Result = float4(Payload.Specular, Payload.Specular, Payload.Specular, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_ROUGHNESS:
			Result = float4(pow(Payload.Roughness, 2.2f).xxx, 1.0f); // NOTE: pow 2.2 is used for consistency with BufferVisualization/Roughness shader
			break;
		case RAY_TRACING_DEBUG_VIZ_IOR:
			Result = float4(Payload.Ior, Payload.Ior, Payload.Ior, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_CUSTOM_DATA:
			Result = float4(Payload.CustomData);
			break;
		case RAY_TRACING_DEBUG_VIZ_WORLD_TANGENT:
			Result = float4(Payload.WorldTangent * 0.5 + 0.5, 1.0f);
			break;
		case RAY_TRACING_DEBUG_VIZ_ANISOTROPY:
			{	
			// Follow same mapping than raster buffer visualization (BP located in Engine/Contents/Anisotropy)
			float AniG = saturate(Payload.Anisotropy);
			float AniB = saturate(-1.0 * Payload.Anisotropy);
			Result = float4(0.0, pow(AniG, 2.2), pow(AniB, 2.2), 1.0f);
			break;
			}
	#endif // STRATA_ENABLED
		}
	}
	return Result;
}

struct FRayTracingDebugPayload : FMinimalPayload
{
	uint InstanceHash;
	uint TriangleHash;
};

float4 DebugRayTracingInstance(RayDesc Ray)
{
	float4 Result = float4(0, 0, 0, 1);

	FRayTracingDebugPayload Payload = (FRayTracingDebugPayload)0;
	Payload.SetMiss();

	TraceRay(
		TLAS,
		RAY_FLAG_CULL_BACK_FACING_TRIANGLES  /*RayFlags*/,
		0xff /*InstanceInclusionMask*/,
		0 /*RayContributionToHitGroupIndex*/,
		RAY_TRACING_NUM_SHADER_SLOTS /*MultiplierForGeometryContributionToShaderIndex*/,
		0 /*MissShaderIndex*/,
		Ray,
		Payload);

	if (Payload.IsHit())
	{
		switch (VisualizationMode)
		{
		default:
		case RAY_TRACING_DEBUG_VIZ_INSTANCES:
			Result.rgb = IntToColor(Payload.InstanceHash);
			break;
		case RAY_TRACING_DEBUG_VIZ_TRIANGLES:
			Result.rgb = IntToColor(Payload.TriangleHash);
			break;
		}
	}

	return Result;
}

RAY_TRACING_ENTRY_RAYGEN(RayTracingDebugMainRGS)
{
	uint2 PixelCoord = DispatchRaysIndex().xy + View.ViewRectMin;

	float2 RenderTargetUV = (float2(PixelCoord) + .5f) * View.BufferSizeAndInvSize.zw;

	RayDesc Ray = CreatePrimaryRay(RenderTargetUV);

	float4 Result = (float4)0;

	if (VisualizationMode == RAY_TRACING_DEBUG_VIZ_INSTANCES || VisualizationMode == RAY_TRACING_DEBUG_VIZ_TRIANGLES)
	{
		Result = DebugRayTracingInstance(Ray);
	}
	else
	{
		bool bIsFarField = false;
		Result = DebugRayTracingMaterial(Ray, bIsFarField);
		
		const bool bNearFieldMiss = Result.w <= 0.0;
		if (VisualizationMode == RAY_TRACING_DEBUG_VIZ_FAR_FIELD && bNearFieldMiss)
		{
			bIsFarField = true;
			const float4 FarFieldTint = float4(1.0f, 0.5f, 0.5f, 1.0f);
			Result = DebugRayTracingMaterial(Ray, bIsFarField) * FarFieldTint;
		}
	}

	if (ShouldUsePreExposure)
	{
		// Only  when the output is tonemapped
		Result *= View.PreExposure;
	}
	
	Output[PixelCoord] = Result;
}

RAY_TRACING_ENTRY_CLOSEST_HIT(RayTracingDebugMainCHS,
	FRayTracingDebugPayload, Payload,
	FRayTracingIntersectionAttributes, Attributes)
{
	Payload.HitT = RayTCurrent();

	// Can't use InstanceIndex() because it's not stable due to culling.
	// Instad we can use InstanceID() (user data on the instance) combined with
	// a hash of the instance transform. This gives us a reasonable overall instance hash
	// that can be used for debug visualizations. 
	// Dynamic objects can't be visualized in a stable way, as their transform is constantly updated.

	Payload.InstanceHash = InstanceID();
	float4x3 Mat = ObjectToWorld4x3();
	for (int i = 0; i < 4; ++i)
	{
		for (int j = 0; j < 3; ++j)
		{
			Payload.InstanceHash += MurmurMix(asuint(Mat[i][j]));
		}
	}
	Payload.TriangleHash = MurmurMix(Payload.InstanceHash + PrimitiveIndex());
}