Solari: Improve specular GI

crates/bevy_solari/src/pathtracer/pathtracer.wgsl

@@ -4,7 +4,7 @@
 #import bevy_render::maths::PI
 #import bevy_render::view::View
 #import bevy_solari::brdf::evaluate_brdf
-#import bevy_solari::sampling::{sample_random_light, random_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, balance_heuristic, power_heuristic}
+#import bevy_solari::sampling::{sample_random_light, random_emissive_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, power_heuristic}
 #import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX}
 
 @group(1) @binding(0) var accumulation_texture: texture_storage_2d<rgba32float, read_write>;
@@ -48,7 +48,7 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
             var mis_weight = 1.0;
             if !bounce_was_perfect_reflection {
-                let p_light = random_light_pdf(ray_hit);
+                let p_light = random_emissive_light_pdf(ray_hit);
                 mis_weight = power_heuristic(p_bounce, p_light);
             }
             radiance += mis_weight * throughput * ray_hit.material.emissive;
@@ -57,8 +57,13 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
             let is_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999;
             if !is_perfectly_specular {
                 let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng);
-                let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit);
-                mis_weight = power_heuristic(1.0 / direct_lighting.inverse_pdf, pdf_of_bounce);
+
+                mis_weight = 1.0;
+                if direct_lighting.brdf_rays_can_hit {
+                    let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit);
+                    mis_weight = power_heuristic(1.0 / direct_lighting.inverse_pdf, pdf_of_bounce);
+                }
+
                 let direct_lighting_brdf = evaluate_brdf(ray_hit.world_normal, wo, direct_lighting.wi, ray_hit.material);
                 radiance += mis_weight * throughput * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
             }

crates/bevy_solari/src/realtime/specular_gi.wgsl

@@ -3,8 +3,8 @@
 #import bevy_render::view::View
 #import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf}
 #import bevy_solari::gbuffer_utils::gpixel_resolve
-#import bevy_solari::sampling::{sample_ggx_vndf, ggx_vndf_pdf}
-#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
+#import bevy_solari::sampling::{sample_random_light, random_emissive_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, power_heuristic}
+#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX}
 #import bevy_solari::world_cache::{query_world_cache, WORLD_CACHE_CELL_LIFETIME}
 
 @group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, read_write>;
@@ -15,6 +15,9 @@
 struct PushConstants { frame_index: u32, reset: u32 }
 var<push_constant> constants: PushConstants;
 
+const DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD: f32 = 0.4;
+const WORLD_CACHE_TERMINATION_ROUGHNESS_THRESHOLD: f32 = 0.4;
+
 @compute @workgroup_size(8, 8, 1)
 fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
     if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
@@ -32,7 +35,7 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
     var radiance: vec3<f32>;
     var wi: vec3<f32>;
-    if surface.material.roughness > 0.1 {
+    if surface.material.roughness > DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD {
         // Surface is very rough, reuse the ReSTIR GI reservoir
         let gi_reservoir = gi_reservoirs_a[pixel_index];
         wi = normalize(gi_reservoir.sample_point_world_position - surface.world_position);
@@ -68,6 +71,8 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
 fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, rng: ptr<function, u32>) -> vec3<f32> {
     var ray_origin = initial_ray_origin;
     var wi = initial_wi;
+    var surface_perfectly_specular = false;
+    var p_bounce = 0.0;
 
     // Trace up to three bounces, getting the net throughput from them
     var radiance = vec3(0.0);
@@ -78,34 +83,72 @@ fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, rng:
         if ray.kind == RAY_QUERY_INTERSECTION_NONE { break; }
         let ray_hit = resolve_ray_hit_full(ray);
 
-        // Add world cache contribution
-        let diffuse_brdf = ray_hit.material.base_color / PI;
-        radiance += throughput * diffuse_brdf * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, WORLD_CACHE_CELL_LIFETIME, rng);
-
-        // Surface is very rough, terminate path in the world cache
-        if ray_hit.material.roughness > 0.1 && i != 0u { break; }
-
-        // Sample new ray direction from the GGX BRDF for next bounce
         let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent);
         let T = TBN[0];
         let B = TBN[1];
         let N = TBN[2];
+
         let wo = -wi;
         let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
+
+        // Add emissive contribution (but not on the first bounce, since ReSTIR DI handles that)
+        if i != 0u {
+            radiance += throughput * emissive_mis_weight(p_bounce, ray_hit, surface_perfectly_specular) * ray_hit.material.emissive;
+        }
+
+        // Should not perform NEE for mirror-like surfaces
+        surface_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999;
+
+        if ray_hit.material.roughness > WORLD_CACHE_TERMINATION_ROUGHNESS_THRESHOLD && i != 0u {
+            // Surface is very rough, terminate path in the world cache
+            let diffuse_brdf = ray_hit.material.base_color / PI;
+            radiance += throughput * diffuse_brdf * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, WORLD_CACHE_CELL_LIFETIME, rng);
+            break;
+        } else if !surface_perfectly_specular {
+            // Sample direct lighting (NEE)
+            let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, rng);
+            let direct_lighting_brdf = evaluate_brdf(ray_hit.world_normal, wo, direct_lighting.wi, ray_hit.material);
+            let mis_weight = nee_mis_weight(direct_lighting.inverse_pdf, direct_lighting.brdf_rays_can_hit, wo_tangent, direct_lighting.wi, ray_hit, TBN);
+            radiance += throughput * mis_weight * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
+        }
+
+        // Sample new ray direction from the GGX BRDF for next bounce
         let wi_tangent = sample_ggx_vndf(wo_tangent, ray_hit.material.roughness, rng);
         wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
         ray_origin = ray_hit.world_position;
 
         // Update throughput for next bounce
-        let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
+        p_bounce = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
         let brdf = evaluate_brdf(N, wo, wi, ray_hit.material);
         let cos_theta = saturate(dot(wi, N));
-        throughput *= (brdf * cos_theta) / pdf;
+        throughput *= (brdf * cos_theta) / p_bounce;
     }
 
     return radiance;
 }
 
+fn emissive_mis_weight(p_bounce: f32, ray_hit: ResolvedRayHitFull, previous_surface_perfectly_specular: bool) -> f32 {
+    if previous_surface_perfectly_specular { return 1.0; }
+
+    let p_light = random_emissive_light_pdf(ray_hit);
+    return power_heuristic(p_bounce, p_light);
+}
+
+fn nee_mis_weight(inverse_p_light: f32, brdf_rays_can_hit: bool, wo_tangent: vec3<f32>, wi: vec3<f32>, ray_hit: ResolvedRayHitFull, TBN: mat3x3<f32>) -> f32 {
+    if !brdf_rays_can_hit {
+        return 1.0;
+    }
+
+    let T = TBN[0];
+    let B = TBN[1];
+    let N = TBN[2];
+    let wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
+
+    let p_light = 1.0 / inverse_p_light;
+    let p_bounce = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
+    return max(0.0, power_heuristic(p_light, p_bounce));
+}
+
 // Don't adjust the size of this struct without also adjusting GI_RESERVOIR_STRUCT_SIZE.
 struct Reservoir {
     sample_point_world_position: vec3<f32>,

crates/bevy_solari/src/scene/sampling.wgsl

@@ -15,6 +15,10 @@ fn balance_heuristic(f: f32, g: f32) -> f32 {
 
 // https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 1)
 fn sample_ggx_vndf(wi_tangent: vec3<f32>, roughness: f32, rng: ptr<function, u32>) -> vec3<f32> {
+    if roughness <= 0.01 {
+        return vec3(-wi_tangent.xy, wi_tangent.z);
+    }
+
     let i = wi_tangent;
     let rand = rand_vec2f(rng);
     let i_std = normalize(vec3(i.xy * roughness, i.z));
@@ -69,6 +73,7 @@ struct LightContribution {
     radiance: vec3<f32>,
     inverse_pdf: f32,
     wi: vec3<f32>,
+    brdf_rays_can_hit: bool,
 }
 
 struct LightContributionNoPdf {
@@ -88,10 +93,9 @@ fn sample_random_light(ray_origin: vec3<f32>, origin_world_normal: vec3<f32>, rn
     return light_contribution;
 }
 
-fn random_light_pdf(hit: ResolvedRayHitFull) -> f32 {
+fn random_emissive_light_pdf(hit: ResolvedRayHitFull) -> f32 {
     let light_count = arrayLength(&light_sources);
-    let p_light = 1.0 / f32(light_count);
-    return p_light / (hit.triangle_area * f32(hit.triangle_count));
+    return 1.0 / (f32(light_count) * f32(hit.triangle_count) * hit.triangle_area);
 }
 
 fn generate_random_light_sample(rng: ptr<function, u32>) -> GenerateRandomLightSampleResult {
@@ -169,7 +173,7 @@ fn calculate_resolved_light_contribution(resolved_light_sample: ResolvedLightSam
 
     let radiance = resolved_light_sample.radiance * cos_theta_origin * (cos_theta_light / light_distance_squared);
 
-    return LightContribution(radiance, resolved_light_sample.inverse_pdf, wi);
+    return LightContribution(radiance, resolved_light_sample.inverse_pdf, wi, resolved_light_sample.world_position.w == 1.0);
 }
 
 fn resolve_and_calculate_light_contribution(light_sample: LightSample, ray_origin: vec3<f32>, origin_world_normal: vec3<f32>) -> LightContributionNoPdf {