refactor: impl multi scattering cache

features/PhysicalSky/Shaders/PhysicalSky/Common.hlsli

@@ -3,8 +3,14 @@
 
 #include "Common/Math.hlsli"
 
+#define TRANSMITTANCE_LUT_WIDTH 256
+#define TRANSMITTANCE_LUT_HEIGHT 64
+#define MULTI_SCATTERING_LUT_WIDTH 32
+#define MULTI_SCATTERING_LUT_HEIGHT 32
+
 cbuffer AtmosphereCB : register(b0)
 {
+	float PlanetRadius;
 	float AtmosphericRadius;
 	float AtmosphericDepth;
 
@@ -20,20 +26,22 @@ cbuffer AtmosphereCB : register(b0)
 	float4 OzoneAbsorption;
 	float OzoneCenterAltitude;
 	float OzoneWidth;
+
+    float3 GroundAlbedo;
 }
 
-float RayleighPhase(float nu)
+float RayleighPhaseFunction(float cosTheta)
 {
 	// BN08
-	return 3.0 / (16.0 * Math::PI) * (1.0 + nu * nu);
+	return 3.0 / (16.0 * Math::PI) * (1.0 + cosTheta * cosTheta);
 }
 
-float3 MiePhase(float nu)
+float3 MiePhaseFunction(float cosTheta)
 {
 	// GK99
 	float g = MieAnisotropy;
 	float k = 3.0 / (8.0 * Math::PI) * (1.0 - g * g) / (2.0 + g * g);
-	return k * (1.0 + nu * nu) / pow(1.0 + g * g - 2.0 * g * nu, 1.5);
+	return k * (1.0 + cosTheta * cosTheta) / pow(1.0 + g * g - 2.0 * g * cosTheta, 1.5);
 }
 
 float OzoneDensity(float h)
@@ -42,43 +50,137 @@ float OzoneDensity(float h)
 	return saturate(1.0 - 0.5 * abs(h - OzoneCenterAltitude) / OzoneWidth);
 }
 
-float3 Extinction(float h)
+float3 ComputeExtinction(float h)
 {
-	float3 e = RayleighAbsorption * exp(-h / RayleighHeight)
-				+ MieAbsorption.rgb * exp(-h / MieHeight)
+	float3 e = RayleighAbsorption.rgb * exp(-h / RayleighHeight)
+				+ MieAbsorption * exp(-h / MieHeight)
 				+ OzoneAbsorption.rgb * OzoneDensity(h);
 
-	return max(e, FLT_MIN);
+	return max(e, float3(1e-6));
 }
 
-float2 IntersectSphere(float r, float mu, float radialDistance)
+float2 IntersectSphere(float radius, float cosTheta, float r)
 {
-	const float s = r * rcp(radialDistance);
-	float d = s * s - saturate(1.0 - mu * mu);
+	const float s = radius * rcp(r);
+	float dis = s * s - saturate(1.0 - cosTheta * cosTheta);
 
-	return d < 0 ? d : radialDistance * float2(-mu - sqrt(d), -mu + sqrt(d));  // the closest hit & the farthest hit
+	return dis < 0 ? dis : r * float2(-cosTheta - sqrt(dis), -cosTheta + sqrt(dis));
 }
 
-float2 IntersectAtmosphere(float3 o, float3 v, out float3 n, out float r)
+float2 IntersectAtmosphere(float3 ro, float3 rd, out float3 n, out float r)
 {
-	n = normalize(o);
-	r = length(o);
+	n = normalize(ro);
+	r = length(ro);
 
-	float2 t = IntersectSphere(AtmosphericRadius, dot(n, -v), r);
+	float2 t = IntersectSphere(AtmosphericRadius, dot(n, -rd), r);
 
 	if (t.y >= 0)
 	{
 		t.x = max(t.x, 0.0);
 
 		if (t.x > 0)
 		{
-			float3 p = o + t.x * -v;
-			n = normalize(p);
+			n = normalize(ro + t.x * -rd);
 			r = AtmosphericRadius;
 		}
 	}
 
 	return t;
 }
 
+// Mapping
+float UnitRangeToTextureCoord(float x, int textureSize)
+{
+	return 0.5f / float(textureSize) + x * (1.0f - 1.0f / float(textureSize));
+}
+
+float TextureCoordToUnitRange(float u, int textureSize)
+{
+	return (u - 0.5f / float(textureSize)) / (1.0f - 1.0f / float(textureSize));
+}
+
+float3 SphericalToCartesian(float phi, float cosTheta)
+{
+	float sinPhi, cosPhi;
+	sincos(phi, sinPhi, cosPhi);
+	float sinTheta = sqrt(saturate(1 - cosPhi * cosPhi));
+
+	return float3(float2(cosPhi, sinPhi) * sinTheta, cosTheta);
+}
+
+float3 SampleSphereUniform(float u1, float u2)
+{
+	const float phi = Math::TAU * u2;
+	const float cosTheta = 1.0 - 2.0 * u1;
+
+	return SphericalToCartesian(phi, cosTheta);
+}
+
+float VanDerCorpus(uint bits)
+{
+	bits = (bits << 16) | (bits >> 16);
+	bits = ((bits & 0x00ff00ff) << 8) | ((bits & 0xff00ff00) >> 8);
+	bits = ((bits & 0x0f0f0f0f) << 4) | ((bits & 0xf0f0f0f0) >> 4);
+	bits = ((bits & 0x33333333) << 2) | ((bits & 0xcccccccc) >> 2);
+	bits = ((bits & 0x55555555) << 1) | ((bits & 0xaaaaaaaa) >> 1);
+	return bits * rcp(4294967296.0);
+}
+
+float Hammersley2d(uint i, uint sampleCount)
+{
+	return float2(float(i) / float(sampleCount), VanDerCorpus(i));
+}
+
+float2 MapTransmittance(float cosTheta, float r)
+{
+    float h = sqrt(AtmosphericRadius * AtmosphericRadius - PlanetRadius * PlanetRadius);
+    float rho = sqrt(max(r * r - PlanetRadius * PlanetRadius, 0.0));
+
+    float d = max(0, IntersectSphere(AtmosphericRadius, cosTheta, r).x);
+    float dMin = AtmosphericRadius - r;
+    float dMax = rho + h;
+
+    float u = UnitRangeToTextureCoord((d - dMin) / (dMax - dMin), TRANSMITTANCE_LUT_WIDTH);
+    float v = UnitRangeToTextureCoord(rho / h, TRANSMITTANCE_LUT_HEIGHT);
+
+    return float2(u, v);
+}
+
+float2 UnmapTransmittance(float2 coord)
+{
+    float h = sqrt(AtmosphericRadius * AtmosphericRadius - PlanetRadius * PlanetRadius);
+    float rho = h * TextureCoordToUnitRange(coord.y, TRANSMITTANCE_LUT_HEIGHT);
+
+    float r = sqrt(rho * rho + PlanetRadius * PlanetRadius);
+
+    float dMin = AtmosphericRadius - r;
+    float dMax = rho + h;
+    float d = dMin + TextureCoordToUnitRange(coord.x, TRANSMITTANCE_LUT_WIDTH) * (dMax - dMin);
+    float cosTheta = d == 0.0 ? 1.0 : clamp((h * h - rho * rho - d * d) / (2.0 * r * d), -1.0, 1.0);
+
+    return float2(cosTheta, r);
+}
+
+float3 SampleTransmittance(Texture2D<float3> t, SamplerState s, float cosTheta, float r)
+{
+    return t.SampleLevel(s, MapTransmittance(cosTheta, r), 0);
+}
+
+float2 MapMultipleScattering(float cosTheta, float r)
+{
+    return saturate(float2(cosTheta * 0.5f + 0.5f, r / AtmosphericDepth));
+}
+
+float2 UnmapMultipleScattering(float2 coord)
+{
+    const float2 uv = coord / (float2(MULTI_SCATTERING_LUT_WIDTH, MULTI_SCATTERING_LUT_HEIGHT) - 1.0);
+
+    return float2(uv.x * 2.0 - 1.0, lerp(PlanetRadius, AtmosphericRadius, uv.y));
+}
+
+float3 SampleMultipleScattering(Texture2D<float3> t, SamplerState s, float cosTheta, float r)
+{
+    return t.SampleLevel(s, MapMultipleScattering(cosTheta, r), 0);
+}
+
 #endif

features/PhysicalSky/Shaders/PhysicalSky/MultiScatteringLutCS.hlsl

@@ -1,11 +1,139 @@
 #include "./Common.hlsli"
 
-#define MULTI_SCATTERING_LUT_WIDTH 32
-#define MULTI_SCATTERING_LUT_HEIGHT 32
 #define SAMPLE_COUNT 64
 
+groupshared float3 gsRadianceMS[SAMPLE_COUNT];
+groupshared float3 gsRadiance[SAMPLE_COUNT];
+
 RWTexture2D<float3> MultiScatteringLUT : register(u0);
+Texture2D<float3> Transmittance : register(t0);
+
+SamplerState TransmittanceSampler;
+
+void ComputeIntegrateResult(float3 ro, float3 rd, float end, float3 lightDir,
+	out float3 skyMultiScattering, out float3 skyColor, out float3 skyTransmittance)
+{
+	skyColor = 0.0f;
+	skyTransmittance = 1.0f;
+    skyMultiScattering = 0.0f;
+
+	const uint sampleCount = 16;
+
+	for (uint s = 0; s < sampleCount; s++)
+	{
+        float t0 = (s) / (float) sampleCount, t1 = (s + 1.0f) / (float) sampleCount;
+        t0 = t0 * t0 * end;
+        t1 = t1 * t1 * end;
+        float t = lerp(t0, t1, 0.5f);
+        float dt = t1 - t0;
+
+        const float3 p = ro + t * rd;
+        const float r = max(length(p), PlanetRadius);
+        const float3 n = p * rcp(r);
+        const float h = r - PlanetRadius;
+
+        const float3 sigmaE = ComputeExtinction(h);
+        const float3 sigmaS = RayleighScattering.rgb * exp(-h / RayleighHeight) + MieScattering.rgb * exp(-h / MieHeight);
+        const float3 transmittance = exp(-sigmaE * dt);
+
+        skyMultiScattering += skyTransmittance * ((sigmaS - sigmaS * transmittance) / sigmaE);
+
+        float cosTheta = dot(n, lightDir);
+        float sinThetaH = PlanetRadius * rcp(r);
+        float cosThetaH = -sqrt(saturate(1 - sinThetaH * sinThetaH));
+        if (cosTheta >= cosThetaH) // Above horizon
+        {
+            const float3 phaseScatter = sigmaS * rcp(4.0 * Math::PI); // Isotropic
+            float3 ms = SampleTransmittance(Transmittance, TransmittanceSampler, cosTheta, r) * phaseScatter;
+            skyColor += skyTransmittance * ((ms - ms * transmittance) / sigmaE);
+        }
+
+        skyTransmittance *= transmittance;
+    }
+}
+
+float3 DrawPlanet(float3 pos, float3 lightDir)
+{
+    float3 n = normalize(pos);
+
+    float3 albedo = GroundAlbedo.xyz;
+    float3 gBrdf = albedo * rcp(Math::PI);
+
+    float cosTheta = dot(n, lightDir);
+
+    float3 intensity = 0.0f;
+    if (cosTheta >= 0)
+    {
+        intensity = SampleTransmittance(Transmittance, TransmittanceSampler, cosTheta, PlanetRadius);
+    }
+
+    return gBrdf * (saturate(cosTheta) * intensity);
+}
+
+void ParallelSum(uint threadIdx, inout float3 radiance, inout float3 radianceMS)
+{
+	gsRadiance[threadIdx] = radiance;
+	gsRadianceMS[threadIdx] = radianceMS;
+	GroupMemoryBarrierWithGroupSync();
+
+	[unroll]
+	for (uint s = SAMPLE_COUNT / 2u; s > 0u; s >>= 1u)
+	{
+		if (threadIdx < s)
+		{
+			gsRadiance[threadIdx] += gsRadiance[threadIdx + s];
+			gsRadianceMS[threadIdx] += gsRadianceMS[threadIdx + s];
+		}
+
+		GroupMemoryBarrierWithGroupSync();
+	}
+
+	radiance = gsRadiance[0];
+	radianceMS = gsRadianceMS[0];
+}
 
 [numthreads(1, 1, SAMPLE_COUNT)] void main(uint3 coord
 											: SV_DispatchThreadID) {
+    const float2 data = UnmapMultipleScattering(coord.xy);
+	const float2 uv = coord.xy / (float2(MULTI_SCATTERING_LUT_WIDTH, MULTI_SCATTERING_LUT_HEIGHT) - 1);
+	const uint threadIdx = coord.z;
+
+    float3 ld = float3(0.0, data.x, sqrt(saturate(1 - data.x * data.x)));
+    float3 ro = float3(0.0f, data.y, 0.0f);
+
+	float2 sample = Hammersley2d(threadIdx, SAMPLE_COUNT);
+	float3 rd = SampleSphereUniform(sample.x, sample.y);
+
+	float3 n; float r;
+	float2 t = IntersectAtmosphere(ro, -rd, n, r);
+	float entry = t.x, exit = t.y;
+
+	float cosChi = dot(n, rd);
+	float sinThetaH = PlanetRadius * rcp(r);
+	float cosThetaH = -sqrt(saturate(1 - sinThetaH * sinThetaH));
+
+	bool targetGround = entry >= 0 && cosChi < cosThetaH;
+
+	if (targetGround)
+		exit = entry + IntersectSphere(PlanetRadius, cosChi, r).x;
+
+	float3 skyMultiScattering = 0.0f, skyColor = 0.0f, skyTransmittance = 1.0f;
+	if (exit > 0.0f)
+        ComputeIntegrateResult(ro, rd, exit, ld, skyMultiScattering, skyColor, skyTransmittance);
+
+	if (targetGround)
+        skyColor += DrawPlanet(ro + exit * rd, ld) * skyTransmittance;
+
+	const float dS = 1 / SAMPLE_COUNT;
+	float3 radiance = skyColor * dS;
+    float3 radianceMS = skyMultiScattering * dS;
+
+	ParallelSum(threadIdx, radiance, radianceMS);
+	if (threadIdx > 0)
+		return;
+
+	const float3 fms = 1.0f * rcp(1.0 - radianceMS);  // Equation 9
+	const float3 ms = radiance * fms;  // Equation 10
+
+	MultiScatteringLUT[coord.xy] = ms;
 }

features/PhysicalSky/Shaders/PhysicalSky/TransmittanceLutCS.hlsl

@@ -0,0 +1,38 @@
+#include "./Common.hlsli"
+
+#define SAMPLE_COUNT 500
+
+RWTexture2D<float3> TransmittanceLUT : register(u0);
+
+[numthreads(8, 8, 1)] void main(uint2 coord
+								: SV_DispatchThreadID) {
+    const float2 data = UnmapTransmittance(float2(coord));
+    float cosTheta = data.x, r = data.y;
+
+    float3 ro = float3(0, r, 0);
+    float3 rd = float3(sqrt(1 - cosTheta * cosTheta), cosTheta, 0);
+
+    float2 t = IntersectSphere(AtmosphericRadius, cosTheta, r);
+
+    float3 transmittance = 1.0f;
+	if (t.y > 0.0)
+	{
+		t.x = max(t.x, 0.0);
+        float3 start = ro + rd * t.x;
+        float3 end = ro + rd * t.y;
+	    float len = t.y - t.x;
+
+        float3 h = 0.0f;
+        for(int i = 0; i < SAMPLE_COUNT; ++i)
+        {
+            float3 p = lerp(start, end, float(i) / SAMPLE_COUNT);
+            float e = ComputeExtinction(length(p) - PlanetRadius);
+
+            h += e * (i == 0 || i == SAMPLE_COUNT) ? 0.5f : 1.0f;
+        }
+
+        transmittance = exp(-h * (len / SAMPLE_COUNT));
+    }
+
+	TransmittanceLUT[coord.xy] = transmittance;
+}

src/Features/PhysicalSky.cpp

@@ -45,6 +45,7 @@ void PhysicalSky::CompileComputeShaders()
 	std::vector<ShaderCompileInfo>
 		shaderInfos = {
 			{ &indirectIrradianceCS, "IndirectIrradianceCS.hlsl", {} },
+			{ &transmittanceLutCS, "TransmittanceLutCS.hlsl", {} },
 			{ &multiScatteringLutCS, "MultiScatteringLutCS.hlsl", {} },
 		};
 

src/Features/PhysicalSky.h

@@ -45,6 +45,7 @@ struct PhysicalSky : Feature
 	eastl::unique_ptr<ConstantBuffer> atmosphereCB = nullptr;
 
 	winrt::com_ptr<ID3D11ComputeShader> indirectIrradianceCS = nullptr;
+	winrt::com_ptr<ID3D11ComputeShader> transmittanceLutCS = nullptr;
 	winrt::com_ptr<ID3D11ComputeShader> multiScatteringLutCS = nullptr;
 
 	virtual bool SupportsVR() override { return true; };