cnq3/code/renderer/shaders/crp/transp_resolve.hlsl

/*
===========================================================================
Copyright (C) 2023-2024 Gian 'myT' Schellenbaum

This file is part of Challenge Quake 3 (CNQ3).

Challenge Quake 3 is free software; you can redistribute it
and/or modify it under the terms of the GNU General Public License as
published by the Free Software Foundation; either version 2 of the License,
or (at your option) any later version.

Challenge Quake 3 is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with Challenge Quake 3. If not, see <https://www.gnu.org/licenses/>.
===========================================================================
*/
// reads per-pixel fragment linked lists into arrays, sorts them and composites them


/*
OIT integration with volumetric lighting

Each texel has the final accumulated in-scattering (RGB) and transmittance (A)
from the near clip plane to the current voxel's depth.
The in-scattered light and transmittance values for any given depth range are therefore:
inScattering = farScatter - nearScatter
transmittance = farTrans / nearTrans

Definitions:
B : opaque/background surface color
C : color accumulator
D : color accumulator #2
T : transmittance from near plane to opaque surface
S : in-scattered light from near plane to opaque surface
T': transmittance from near plane to fragment
S': in-scattered light from near plane to fragment

If we use a single color accumulator C and apply transmittance T and in-scattering S post-blend,
we run the following logic for the opaque surface and every fragment:
C = Blend(C, ...)
C = CT + S
Doing this creates artifacts along the geometric edges of explosions (additive blend, black borders). But why?
Suppose we have no fragment, the final color F is therefore:
F = BT + S
What if we add a single fragment that contributes nothing (e.g. black and additive blended)?
C = BT/T' + S - S'
C = Blend(C, ...) // doesn't change C
F = (BT/T' + S - S') * T'/1 + S' - 0
F = BT + ST' - S'T' + S'
It's mismatched: we get (ST' - S'T' + S') instead of (S).

Let's try a separate color accumulator D just for the in-scattered light like so:
C = Blend(C, ...)
C = CT
D = D + S
F = C + D
No fragment case:
F = C + D == BT + S
Single useless fragment case:
C = BT/T'
D = S - S'
C = Blend(C, ...) // doesn't change C
C = CT'/1 == BT/T' * T'/1 == BT
D = D + S' - 0 == S - S' + S' - 0 == S
F = C + D == BT + S
Now it matches perfectly.
*/


#include "common.hlsli"
#include "fullscreen.hlsli"
#include "oit.h.hlsli"
#include "scene_view.h.hlsli"
#include "../common/state_bits.h.hlsli"


cbuffer RootConstants
{
	float2 scissorRectMin;
	float2 scissorRectMax;
	uint renderTargetTexture;
	uint shaderIndexBuffer;
	uint indexTexture;
	uint fragmentBuffer;
	uint centerPixel; // y: 16 - x: 16
	uint scatterTextureIndex;
	uint scatterSamplerIndex;
};

uint GetShaderStage(uint stateBits)
{
	return (stateBits & GLS_STAGEINDEX_BITS) >> GLS_STAGEINDEX_SHIFT;
}

bool IsBehind(float depthA, float depthB, uint stageA, uint stageB)
{
	if(depthA > depthB)
	{
		return true;
	}

	if(depthA == depthB && stageA < stageB)
	{
		return true;
	}

	return false;
}

// from NVIDIA's 2007 "Soft Particles" whitepaper by Tristan Lorach
float Contrast(float d, float power)
{
	bool aboveHalf = d > 0.5;
	float base = saturate(2.0 * (aboveHalf ? (1.0 - d) : d));
	float r = 0.5 * pow(base, power);

	return aboveHalf ? (1.0 - r) : r;
}

float GetBitAsFloat(uint bits, uint bitIndex)
{
	return (bits & (1u << bitIndex)) ? 1.0 : 0.0;
}

float4 DepthFadeFragmentColor(float4 color, OIT_Fragment fragment, float opaqueViewDepth)
{
	if(((fragment.depthFadeScaleBiasPO >> 8) & 1) == 0)
	{
		return color;
	}

#define BIT(Index) GetBitAsFloat(fragment.depthFadeScaleBiasPO, Index)
	float4 dst = color;
	float2 distOffset = UnpackHalf2(fragment.depthFadeDistOffset);
	float4 fadeColorScale = float4(BIT(0), BIT(1), BIT(2), BIT(3));
	float4 fadeColorBias = float4(BIT(4), BIT(5), BIT(6), BIT(7));
	float depthS = opaqueViewDepth; // stored depth, linear
	float depthP = fragment.depth - distOffset.y; // fragment depth, linear
	float fadeScale = Contrast((depthS - depthP) * distOffset.x, 2.0);
	dst = lerp(dst * fadeColorScale + fadeColorBias, dst, fadeScale);
#undef BIT

	return dst;
}

float3 BlendInScatteredLight(float4 srcColor, float3 dstColor, uint blendBits)
{
	// source blend must not double source contributions, so we can't call BlendSource
	uint srcBlend = blendBits & GLS_SRCBLEND_BITS;
	float3 srcContrib = float3(0, 0, 0);
	if(srcBlend == GLS_SRCBLEND_DST_COLOR)
	{
		srcContrib = dstColor * (float3(1, 1, 1) - dstColor);
	}
	else if(srcBlend == GLS_SRCBLEND_ONE_MINUS_DST_COLOR)
	{
		srcContrib = dstColor * (float3(1, 1, 1) - dstColor);
	}

	uint dstBlend = blendBits & GLS_DSTBLEND_BITS;
	float3 dstContrib = BlendDest(srcColor, float4(dstColor, 1), dstBlend).rgb;
	float3 result = srcContrib + dstContrib;

	return result;
}

struct OIT_Resolve
{
	bool InitScissorRect(VOut input)
	{
		renderTarget = ResourceDescriptorHeap[renderTargetTexture];
		tcPx = int3(input.position.xy, 0);
		opaqueColor = renderTarget.Load(tcPx);
		if(any(input.position.xy < scissorRectMin) ||
			any(input.position.xy > scissorRectMax))
		{
			return true;
		}
		return false;
	}

	void Init(VOut input)
	{
		scene = GetSceneView();
#if defined(VOLUMETRIC_LIGHT)
		scatterTexture = ResourceDescriptorHeap[scatterTextureIndex];
		scatterSampler = SamplerDescriptorHeap[scatterSamplerIndex];
#endif

		RWTexture2D<uint> index = ResourceDescriptorHeap[indexTexture];
		RWStructuredBuffer<OIT_Fragment> fragments = ResourceDescriptorHeap[fragmentBuffer];
		Texture2D depthTex = ResourceDescriptorHeap[scene.depthTextureIndex];
		uint fragmentIndex = index[tcPx.xy];
		fragmentCount = 0;
		float storedDepthZW = depthTex.Load(tcPx).x; // stored depth, z/w
		opaqueViewDepth = scene.LinearDepth(storedDepthZW);

		// grab this pixel's fragments	
		while(fragmentIndex != 0 && fragmentCount < OIT_MAX_FRAGMENTS_PER_PIXEL)
		{
			sorted[fragmentCount] = fragments[fragmentIndex];
			fragmentIndex = sorted[fragmentCount].next;
			++fragmentCount;
		}

		// sort the fragments using an insertion sort
		for(uint i = 1; i < fragmentCount; ++i)
		{
			OIT_Fragment insert = sorted[i];
			uint stage = GetShaderStage(insert.stateBits);
			uint j = i;
			while(j > 0 && IsBehind(insert.depth, sorted[j - 1].depth, stage, GetShaderStage(sorted[j - 1].stateBits)))
			{
				sorted[j] = sorted[j - 1];
				--j;
			}
			sorted[j] = insert;
		}
	}

	float4 Resolve(VOut input)
	{
		float4 color = saturate(opaqueColor);

#if defined(VOLUMETRIC_LIGHT)
		float3 inScatterAccum = float3(0, 0, 0);
		// initialize volume traversal
		float3 volumeSize = GetTextureSize(scatterTexture);
		float opaqueFroxelDepth01 = scene.FroxelViewDepthToZ01(opaqueViewDepth, volumeSize.z);
		// @TODO: do the depth bias only when global fog is enabled or a CVar is set
		opaqueFroxelDepth01 = max(opaqueFroxelDepth01 - 1.0 / volumeSize.z, 0.0);
		float3 scatterTC = float3(input.texCoords, opaqueFroxelDepth01);
		float4 scatterData = scatterTexture.SampleLevel(scatterSampler, scatterTC, 0);
		{
			float4 closerScatterData = FindCloserScatterData(0, -1.0, input, volumeSize);
			float3 inScattering = scatterData.rgb - closerScatterData.rgb;
			float transmittance = min(scatterData.a / max(closerScatterData.a, 0.000001), 1.0);
			inScatterAccum = inScattering;
			color.rgb = color.rgb * transmittance;
			scatterData = closerScatterData;
		}
#endif

		// blend the results
		lastFragmentIndex = -1;
		float dstDepth = -1.0;
		for(uint i = 0; i < fragmentCount; ++i)
		{
			OIT_Fragment frag = sorted[i];
			uint stateBits = frag.stateBits;
			float fragDepth = frag.depth;
			if((stateBits & (GLS_DEPTHFUNC_EQUAL | GLS_DEPTHTEST_DISABLE)) == GLS_DEPTHFUNC_EQUAL &&
				fragDepth != dstDepth)
			{
				continue;
			}

			float4 fragColor = UnpackColor(frag.color);
			float4 prevColor = color;
			fragColor = DepthFadeFragmentColor(fragColor, frag, opaqueViewDepth);
			color = Blend(fragColor, color, stateBits);
			color = saturate(color);
			if((stateBits & GLS_DEPTHMASK_TRUE) != 0u &&
				fragDepth != dstDepth)
			{
				dstDepth = fragDepth;
			}

			// we have to not include the alpha channel in this test for it to be correct
			if(any(color.rgb != prevColor.rgb))
			{
				lastFragmentIndex = (int)i;
			}

#if defined(VOLUMETRIC_LIGHT)
			float4 closerScatterData = FindCloserScatterData(i + 1, dstDepth, input, volumeSize);
			float3 inScattering = scatterData.rgb - closerScatterData.rgb;
			float transmittance = min(scatterData.a / max(closerScatterData.a, 0.000001), 1.0);
			inScatterAccum = inScattering + BlendInScatteredLight(fragColor, inScatterAccum, stateBits & GLS_BLEND_BITS);
			color.rgb *= transmittance;
			scatterData = closerScatterData;
#endif
		}

#if defined(VOLUMETRIC_LIGHT)
		color.rgb += inScatterAccum;
#endif

		return color;
	}

#if defined(VOLUMETRIC_LIGHT)
	float4 FindCloserScatterData(uint startIndex, float dstDepth, VOut input, float3 volumeSize)
	{
		float4 closerScatterData = float4(0, 0, 0, 1);
		for(uint j = startIndex; j < fragmentCount; ++j)
		{
			OIT_Fragment frag = sorted[j];
			uint stateBits = frag.stateBits;
			float fragDepth = frag.depth;
			if((stateBits & (GLS_DEPTHFUNC_EQUAL | GLS_DEPTHTEST_DISABLE)) == GLS_DEPTHFUNC_EQUAL &&
				fragDepth != dstDepth)
			{
				continue;
			}

			float froxelDepth01 = scene.FroxelViewDepthToZ01(fragDepth, volumeSize.z);
			float3 scatterTC = float3(input.texCoords, froxelDepth01);
			closerScatterData = scatterTexture.SampleLevel(scatterSampler, scatterTC, 0);
			break;
		}

		return closerScatterData;
	}
#endif

	void WriteShaderID(VOut input)
	{
		// write out the fragment shader ID of the closest visible fragment of the center pixel
		if(lastFragmentIndex >= 0)
		{
			OIT_Fragment closest = sorted[lastFragmentIndex];
			uint shaderTrace = closest.shaderTrace;
			if(shaderTrace & 1)
			{
				uint2 fragmentCoords = uint2(input.position.xy);
				uint2 centerCoords = uint2(centerPixel & 0xFFFF, centerPixel >> 16);
				if(all(fragmentCoords == centerCoords))
				{
					RWByteAddressBuffer shaderIdBuf = ResourceDescriptorHeap[shaderIndexBuffer];
					uint shaderIndex = shaderTrace >> 1;
					shaderIdBuf.Store(0, shaderIndex);
				}
			}
		}
	}

	SceneView scene;
	Texture2D renderTarget;
	int3 tcPx;
	float4 opaqueColor;
	OIT_Fragment sorted[OIT_MAX_FRAGMENTS_PER_PIXEL];
	uint fragmentCount;
	int lastFragmentIndex;
	float opaqueViewDepth;
#if defined(VOLUMETRIC_LIGHT)
	Texture3D<float4> scatterTexture;
	SamplerState scatterSampler;
#endif
};

float4 ps(VOut input) : SV_Target
{
	OIT_Resolve resolve;
	if(resolve.InitScissorRect(input))
	{
		return resolve.opaqueColor;
	}

	resolve.Init(input);
	float4 color = resolve.Resolve(input);
	resolve.WriteShaderID(input);

	return color;
}