UnrealEngineUWP/Engine/Shaders/Private/Strata/StrataMaterialClassification.usf

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

#include "/Engine/Private/Common.ush"

#define STRATA_INLINE_SHADING 0
#define STRATA_SSS_MATERIAL_OVERRIDE 0
#include "/Engine/Private/Strata/Strata.ush"

#define TILE_SIZE 8
#define GROUP_THREAD_COUNT (TILE_SIZE*TILE_SIZE)

#if SHADER_CATEGORIZATION

uint MaxBytesPerPixel;
ByteAddressBuffer MaterialLobesBuffer;

// Note: 
// This code is partial copy of PackStrataOut in Strata.ush. 
// The command part will progressively be removed as we remove data from the gbuffer.
void ComputePostProcessBSDFData(
	const FStrataPixelHeader Header,
	const FStrataBSDF BSDF,
	const FStrataAddressing StrataAddressing,
	float3 V,
	inout FStrataClassification Classification,
	inout FStrataSubsurfaceData SSS,
	inout FStrataTopLayerData TopLayer,
	inout float TopLayerRoughness, 
	inout float TopLayerRoughnessAverageCount)
{
	const float OpaqueBSDFThroughput = 0.0f;
	const float FullThroughput = 1.0f;
	const float FullyRough = 1.0f;

	// FinalWeight does not always contains coverage (see Strata.ush)
	const float  BSDFWeightAvg = dot(BSDF.FinalWeight, 1.0f/3.0f);

	// STRATA_TODO apply the same normal transform as for the material
	float3x3 TangentBasis = StrataGetBSDFSharedBasis(Header, BSDF, StrataAddressing);
	float3 N = TangentBasis[2];
	//N = normalize(N);


	BxDFContext ContextNoLight = (BxDFContext)0;
	const float3 FakeL = float3(0, 0, 1);
	Init(ContextNoLight, N, V, FakeL);

	const bool bTopLayer = BSDF_GETISTOPLAYER(BSDF);
	if (bTopLayer)
	{
		// We must take into account the weight in case horizontal blending is used,
		// together with a threshold to avoid 0-length normal after normalization in case influence is fadeout using the Weight node.
		TopLayer.WorldNormal += N * max(1e-5, BSDFWeightAvg); 
	}

	Classification.ShadingModels |= StrataShadingModelBit(BSDF_GETTYPE(BSDF));

	switch (BSDF_GETTYPE(BSDF))
	{
	case STRATA_BSDF_TYPE_SLAB:
	{
		Classification.bIsSimple = Classification.bIsSimple && IsStrataSlabFastPathCompatible(BSDF);

		if (bTopLayer)
		{
			// Setup SSR for the top layer
			TopLayerRoughness += BSDFWeightAvg * SLAB_ROUGHNESSX(BSDF);
			TopLayerRoughnessAverageCount += BSDFWeightAvg;
		}

		// Select the latest encoutered SSS BSDF
		if (!BSDF_GETISTHIN(BSDF) &&					// STRATA_TODO A material cannot have subsurface if using Thin mode. This is because DMFP does not match MFP at all (right now, SSS would just smears thin lighting model details out)
			BSDF_GETHASSSS(BSDF) && BSDFWeightAvg > 0)
		{
			Classification.bHasSubsurface = true;

			SSS.BaseColor = SLAB_DIFFUSEALBEDO(BSDF);
			SSS.Specular  = SLAB_F0(BSDF);				// STRATA_TODO: Specular or F0?

			// Profile ID: we only keep the latest ID, and profile will overtake any non-profile based SSS
			SSS.bIsValid			= true;
			SSS.bIsProfile			= SSS.bIsProfile || BSDF_GETHASSSSPROFILE(BSDF);
			SSS.ProfileId			= StrataSubsurfaceProfileIdTo8bits(SLAB_SSSPROFILEID(BSDF));
			SSS.ProfileRadiusScale	= SLAB_SSSPROFILERADIUSSCALE(BSDF);

			// STRATA_TODO This should be accumulated according to the BSDF weight when we run this in the base pass.
			// We cannot do that now because BSDFWeight contains weight*coverage*trhougput all mixed together.
			SSS.DMFP				= SLAB_SSSDMFP(BSDF);
		}
		break;
	}

	case STRATA_BSDF_TYPE_HAIR:
	{
		Classification.bIsSimple = false;

		// No SSR
		if (bTopLayer)
		{
			TopLayerRoughness += BSDFWeightAvg * FullyRough; // Do not trigger SSR by using fully rough
			TopLayerRoughnessAverageCount += BSDFWeightAvg;
		}
		break;
	}

	case STRATA_BSDF_TYPE_SINGLELAYERWATER:
	{
		Classification.bIsSimple = false;

		if (bTopLayer)
		{
			const float Roughness = SLW_ROUGHNESS(BSDF);

			TopLayerRoughness += BSDFWeightAvg * Roughness;
			TopLayerRoughnessAverageCount += BSDFWeightAvg;
		}
		break;
	}

	} // switch		
}

void MainPS(
	float4 SVPos : SV_POSITION, 
	out uint  OutColor0 : SV_Target0,
	out uint  OutColor1 : SV_Target1,
	out uint2 OutColor2 : SV_Target2)
{
	const uint2 PixelPos = uint2(SVPos.xy);
	
	FStrataClassification ClassificationData	= (FStrataClassification)0;
	FStrataSubsurfaceData SSSData				= (FStrataSubsurfaceData)0;
	FStrataTopLayerData TopLayerData			= (FStrataTopLayerData)0;

	FStrataAddressing StrataAddressing = GetStrataPixelDataByteOffset(PixelPos, uint2(View.BufferSizeAndInvSize.xy), MaxBytesPerPixel);
	FStrataPixelHeader Header = UnpackStrataHeaderIn(MaterialLobesBuffer, StrataAddressing);

	ClassificationData.bIsStrataMaterial = Header.BSDFCount > 0;
	ClassificationData.MaterialAO = 1.0;
	ClassificationData.bIsSimple = true; // non strata material or sky pixels are considered as simple.
	if (ClassificationData.bIsStrataMaterial)
	{
		float2 ScreenPosition = SvPositionToScreenPosition(SVPos).xy;
		const float Depth = 1000000.0f;
		float4 WorldPos = mul(float4(ScreenPosition * Depth, Depth, 1), View.ScreenToWorld);
		const float3 V = normalize(View.WorldCameraOrigin - WorldPos.xyz);

		ClassificationData.MaterialAO = Header.MaterialAO;
		ClassificationData.bIsSimple = Header.BSDFCount == 1;

		float TopLayerRoughness = 0.0f;
		float TopLayerRoughnessAverageCount = 0.0f;
		for (uint BSDFIndex = 0; BSDFIndex < Header.BSDFCount; ++BSDFIndex)
		{
			const FStrataBSDF BSDF = UnpackStrataBSDFIn(MaterialLobesBuffer, StrataAddressing);
			ComputePostProcessBSDFData(Header, BSDF, StrataAddressing, V, ClassificationData, SSSData, TopLayerData, TopLayerRoughness, TopLayerRoughnessAverageCount);
		}

		TopLayerData.WorldNormal = normalize(TopLayerData.WorldNormal);
		ClassificationData.TopLayerRoughness = TopLayerRoughnessAverageCount > 0.0f ? TopLayerRoughness / TopLayerRoughnessAverageCount : 0.0f;
	}
	else
	{
		TopLayerData.WorldNormal = normalize(float3(-1.0f, -1.0f, -1.0f));
	}

	OutColor0 = StrataPackClassificationData(ClassificationData);
	OutColor1 = StrataPackTopLayerData(TopLayerData);
	OutColor2 = StrataPackSSSData(SSSData);
}
#endif // SHADER_CATEGORIZATION

////////////////////////////////////////////////////////////////////////////////////////////////////////////

uint EncodeTile(uint2 TileCoord)
{
	return TileCoord.x | (TileCoord.y << 16); // assumes 16bit is enough to represent a tiled resolution up to 65,535 :)
}

uint2 DecodeTile(uint In)
{
	return uint2(In & 0xFFFF, In >> 16);
}

#if SHADER_TILE_CATEGORIZATION
int TileSize;
int bRectPrimitive;
int2 ViewResolution;
Texture2D<uint> ClassificationTexture;
RWBuffer<uint> SimpleTileIndirectDataBuffer;
RWBuffer<uint> SimpleTileListDataBuffer;
RWBuffer<uint> ComplexTileIndirectDataBuffer;
RWBuffer<uint> ComplexTileListDataBuffer;

groupshared uint s_IsComplexIsMaterialTile[GROUP_THREAD_COUNT];

[numthreads(TILE_SIZE, TILE_SIZE, 1)]
void TileMainCS(uint2 DispatchThreadId : SV_DispatchThreadID, uint LinearIndex : SV_GroupIndex, uint3 GroupId : SV_GroupID)
{
	if (all(DispatchThreadId == 0))
	{
		SimpleTileIndirectDataBuffer[0] = bRectPrimitive > 0 ? 4 : 6;
		ComplexTileIndirectDataBuffer[0] = bRectPrimitive > 0 ? 4 : 6;
	}

	// TODO: add a SM6 permutation with ballot?
	const uint2 PixelCoord = DispatchThreadId;

	// Pixels outside of the view area are considered simple to enable screen borders to receive the simple permutation when not aligned to shader group size.
	bool bContainsComplexMaterial = false;
	bool bContainsStrataMaterial = false;
	if (all(PixelCoord < uint2(ViewResolution)))
	{
		const uint PackClassification = ClassificationTexture.Load(uint3(PixelCoord, 0));
		FStrataClassification Data = StrataUnpackClassificationData(PackClassification);
		bContainsComplexMaterial = !Data.bIsSimple;
		bContainsStrataMaterial = Data.bIsStrataMaterial;
	}
	s_IsComplexIsMaterialTile[LinearIndex] = (bContainsComplexMaterial ? 1 : 0) | (bContainsStrataMaterial ? 2 : 0);

	GroupMemoryBarrierWithGroupSync();
	if (LinearIndex < 32)
	{
		s_IsComplexIsMaterialTile[LinearIndex] = s_IsComplexIsMaterialTile[LinearIndex] | s_IsComplexIsMaterialTile[LinearIndex + 32];
	}
	GroupMemoryBarrierWithGroupSync();
	if (LinearIndex < 16)
	{
		s_IsComplexIsMaterialTile[LinearIndex] = s_IsComplexIsMaterialTile[LinearIndex] | s_IsComplexIsMaterialTile[LinearIndex + 16];
	}
	GroupMemoryBarrierWithGroupSync();

	if (LinearIndex < 8)
	{
		s_IsComplexIsMaterialTile[LinearIndex] = s_IsComplexIsMaterialTile[LinearIndex] | s_IsComplexIsMaterialTile[LinearIndex + 8];
	}
	GroupMemoryBarrierWithGroupSync();
	if (LinearIndex < 4)
	{
		s_IsComplexIsMaterialTile[LinearIndex] = s_IsComplexIsMaterialTile[LinearIndex] | s_IsComplexIsMaterialTile[LinearIndex + 4];
	}
	GroupMemoryBarrierWithGroupSync();
	if (LinearIndex < 2)
	{
		s_IsComplexIsMaterialTile[LinearIndex] = s_IsComplexIsMaterialTile[LinearIndex] | s_IsComplexIsMaterialTile[LinearIndex + 2];
	}
	GroupMemoryBarrierWithGroupSync();

	const uint FinalIsSimpleIsMaterialTile = s_IsComplexIsMaterialTile[LinearIndex] | s_IsComplexIsMaterialTile[LinearIndex + 1];
	bContainsComplexMaterial = (FinalIsSimpleIsMaterialTile & 1) > 0;
	bContainsStrataMaterial = (FinalIsSimpleIsMaterialTile & 2) > 0;

	if (LinearIndex < 1 && bContainsStrataMaterial)
	{
		if (bContainsComplexMaterial)
		{
			uint WriteToIndex;
			InterlockedAdd(ComplexTileIndirectDataBuffer[1], 1, WriteToIndex);
			ComplexTileListDataBuffer[WriteToIndex] = EncodeTile(GroupId.xy);
		}
		else
		{
			uint WriteToIndex;
			InterlockedAdd(SimpleTileIndirectDataBuffer[1], 1, WriteToIndex);
			SimpleTileListDataBuffer[WriteToIndex] = EncodeTile(GroupId.xy);
		}
	}
}
#endif // SHADER_TILE_CATEGORIZATION

////////////////////////////////////////////////////////////////////////////////////////////////////////////

#if SHADER_STENCIL_CATEGORIZATION

float4 OutputViewSizeAndInvSize;
float4 OutputBufferSizeAndInvSize;
float4x4 ViewScreenToTranslatedWorld;
Buffer<uint> TileListBuffer;
float4 DebugTileColor;

void StrataTilePassVS(
	in uint InVertexId : SV_VertexID,
	in uint InInstanceId : SV_InstanceID,
#if PERMUTATION_ENABLE_TEXCOORD_SCREENVECTOR
	out float2 OutTexCoord : TEXCOORD0,		// Used by DeferredLightPixelMain directional light
	out float3 OutScreenVector : TEXCOORD1,	// Used by DeferredLightPixelMain directional light
#endif
#if PERMUTATION_ENABLE_DEBUG
	out uint2 OutTileCoord : TILE_COORD,	// Used by StencilMainPS below
#endif
	out float4 OutPosition : SV_POSITION)	// Use by StencilMainPS
{
	const uint2 TileCoord = DecodeTile(TileListBuffer[InInstanceId]);

	uint2 TileVertex = TileCoord * TILE_SIZE;
	TileVertex.x += InVertexId == 1 || InVertexId == 2 || InVertexId == 4 ? TILE_SIZE : 0;
	TileVertex.y += InVertexId == 2 || InVertexId == 4 || InVertexId == 5 ? TILE_SIZE : 0;

	OutPosition = float4(float2(TileVertex) * OutputViewSizeAndInvSize.zw * float2(2.0f, -2.0f) + float2(-1.0, 1.0f), 0.5f, 1.0f);

#if PERMUTATION_ENABLE_TEXCOORD_SCREENVECTOR
	OutTexCoord = float2(TileVertex) * OutputBufferSizeAndInvSize.zw;
	OutScreenVector = mul(float4(OutPosition.xy, 1, 0), ViewScreenToTranslatedWorld).xyz;
#endif

#if PERMUTATION_ENABLE_DEBUG
	OutTileCoord = TileCoord;
#endif
}

void StencilMainPS(
	in uint2 InTileCoord : TILE_COORD,
	in float4 SVPos : SV_POSITION,
	out float4 OutColor0 : SV_Target0)
{
	const bool bTileX = (InTileCoord.x & 1) == 0;
	const bool bTileY = (InTileCoord.y & 1) == 0;
	const bool bChecker = (bTileX && bTileY) || (!bTileX && !bTileY);
	OutColor0 = DebugTileColor * float4(1.0f, 1.0f, 1.0f, bChecker ? 0.33f : 0.66f);
	OutColor0.rgb *= OutColor0.a; // premultiplied alpha blending assuming alpha is opacity==coverage
}

#endif //SHADER_STENCIL_CATEGORIZATION

////////////////////////////////////////////////////////////////////////////////////////////////////////////

#if SHADER_CLEAR_MATERIAL_BUFFER

RWByteAddressBuffer MaterialLobesBufferUAV;
uint MaxBytesPerPixel;
int2 TiledViewBufferResolution;

[numthreads(8, 8, 1)]
void ClearMaterialBufferMainCS(uint2 DispatchThreadId : SV_DispatchThreadID)
{
	if (any(int2(DispatchThreadId.xy) > TiledViewBufferResolution))
	{
		return;
	}

	// Custom clear of the Strata material buffer.
	// The first layer of tiled uints contains the header that we need to clear so we only write a single uint per pixel instead of clearing the entire buffer which would be too slow.

	FStrataPixelHeader StrataHeader = InitialiseStrataPixelHeader();
	uint BSDFCount = 0;
	uint PackedStrataHeader = PackStrataHeaderIn(BSDFCount, StrataHeader);

	FStrataAddressing StrataAddressing = GetStrataPixelDataByteOffset(DispatchThreadId.xy, TiledViewBufferResolution, MaxBytesPerPixel);
	MaterialLobesBufferUAV.Store(StrataAddressing.CurrentAddr, PackedStrataHeader);
}

#endif // SHADER_CLEAR_MATERIAL_BUFFER