Solari: Better path termination heuristic

crates/bevy_solari/src/pathtracer/pathtracer.wgsl

@@ -39,7 +39,6 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
     var throughput = vec3(1.0);
     var p_bounce = 0.0;
     var bounce_was_perfect_reflection = true;
-    var previous_normal = vec3(0.0);
     loop {
         let ray = trace_ray(ray_origin, ray_direction, ray_t_min, RAY_T_MAX, RAY_FLAG_NONE);
         if ray.kind != RAY_QUERY_INTERSECTION_NONE {
@@ -75,7 +74,6 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
             ray_t_min = RAY_T_MIN;
             p_bounce = next_bounce.pdf;
             bounce_was_perfect_reflection = next_bounce.perfectly_specular_bounce;
-            previous_normal = ray_hit.world_normal;
 
             // Update throughput for next bounce
             let brdf = evaluate_brdf(ray_hit.world_normal, wo, next_bounce.wi, ray_hit.material);

crates/bevy_solari/src/realtime/specular_gi.wgsl

@@ -1,7 +1,6 @@
 #define_import_path bevy_solari::specular_gi
 
 #import bevy_pbr::pbr_functions::{calculate_tbn_mikktspace, calculate_diffuse_color, calculate_F0}
-#import bevy_pbr::prepass_bindings::PreviousViewUniforms
 #import bevy_render::maths::{orthonormalize, PI}
 #import bevy_render::view::View
 #import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf}
@@ -17,7 +16,6 @@
 
 const DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD: f32 = 0.4;
 const SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD: f32 = 0.0225;
-const TERMINATE_IN_WORLD_CACHE_THRESHOLD: f32 = 0.03;
 
 @compute @workgroup_size(8, 8, 1)
 fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
@@ -33,7 +31,8 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
     let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
 
     let wo_unnormalized = view.world_position - surface.world_position;
-    let wo = normalize(wo_unnormalized);
+    let wo_length = length(wo_unnormalized);
+    let wo = wo_unnormalized / wo_length;
 
     var radiance: vec3<f32>;
     var wi: vec3<f32>;
@@ -53,12 +52,7 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
         wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
         let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, surface.material.roughness);
 
-        // https://d1qx31qr3h6wln.cloudfront.net/publications/mueller21realtime.pdf#subsection.3.4, equation (4)
-        let cos_theta = saturate(dot(wo, surface.world_normal));
-        var a0 = dot(wo_unnormalized, wo_unnormalized) / (4.0 * PI * cos_theta);
-        a0 *= TERMINATE_IN_WORLD_CACHE_THRESHOLD;
-
-        radiance = trace_glossy_path(global_id.xy, surface, wi, pdf, a0, &rng) / pdf;
+        radiance = trace_glossy_path(global_id.xy, surface, wo_length, wi, pdf, &rng) / pdf;
     }
 
     let brdf = evaluate_specular_brdf(surface.world_normal, wo, wi, surface.material.base_color, surface.material.metallic,
@@ -74,15 +68,15 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
 #endif
 }
 
-fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initial_wi: vec3<f32>, initial_p_bounce: f32, a0: f32, rng: ptr<function, u32>) -> vec3<f32> {
+fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initial_ray_t: f32, initial_wi: vec3<f32>, initial_p_bounce: f32, rng: ptr<function, u32>) -> vec3<f32> {
     var radiance = vec3(0.0);
     var throughput = vec3(1.0);
 
     var ray_origin = primary_surface.world_position;
     var wi = initial_wi;
     var p_bounce = initial_p_bounce;
     var surface_perfect_mirror = false;
-    var path_spread = 0.0;
+    var path_spread = path_spread_heuristic(initial_ray_t, primary_surface.material.roughness);
 
 #ifdef DLSS_RR_GUIDE_BUFFERS
     var mirror_rotations = reflection_matrix(primary_surface.world_normal);
@@ -111,9 +105,6 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
         // Should not perform NEE for mirror-like surfaces
         surface_perfect_mirror = ray_hit.material.roughness <= MIRROR_ROUGHNESS_THRESHOLD && ray_hit.material.metallic > 0.9999;
 
-        // https://d1qx31qr3h6wln.cloudfront.net/publications/mueller21realtime.pdf#subsection.3.4, equation (3)
-        path_spread += sqrt((ray.t * ray.t) / (p_bounce * wo_tangent.z));
-
         // Primary surface replacement for perfect mirrors
         // https://developer.nvidia.com/blog/rendering-perfect-reflections-and-refractions-in-path-traced-games/#primary_surface_replacement
 #ifdef DLSS_RR_GUIDE_BUFFERS
@@ -127,8 +118,12 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
         }
 #endif
 
-        if path_spread * path_spread > a0 * get_cell_size(ray_hit.world_position, view.world_position) {
-            // Path spread is wide enough, terminate path in the world cache
+        // Terminate path in the world cache if the ray is long enough and the path spread is large enough
+        let world_cache_cell_size = get_cell_size(ray_hit.world_position, view.world_position);
+        let ray_longer_than_cell = ray.t > sqrt(3.0) * world_cache_cell_size;
+        let path_spread_large_enough = path_spread > world_cache_cell_size * world_cache_cell_size;
+
+        if ray_longer_than_cell && path_spread_large_enough {
             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, ray.t, WORLD_CACHE_CELL_LIFETIME, rng);
             break;
@@ -148,8 +143,10 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
         // Update throughput for next bounce
         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) / p_bounce;
+        throughput *= brdf / p_bounce;
+
+        // Path spread increase
+        path_spread += path_spread_heuristic(ray.t, ray_hit.material.roughness);
     }
 
     return radiance;
@@ -186,6 +183,12 @@ fn nee_mis_weight(inverse_p_light: f32, brdf_rays_can_hit: bool, wo_tangent: vec
     return power_heuristic(p_light, p_bounce);
 }
 
+fn path_spread_heuristic(ray_t: f32, roughness: f32) -> f32 {
+    let alpha_squared = min(roughness * roughness, 0.99);
+    let distance_squared = ray_t * ray_t;
+    return distance_squared * 0.5 * (alpha_squared / (1.0 - alpha_squared));
+}
+
 #ifdef DLSS_RR_GUIDE_BUFFERS
 // https://en.wikipedia.org/wiki/Householder_transformation
 fn reflection_matrix(plane_normal: vec3f) -> mat3x3<f32> {

crates/bevy_solari/src/realtime/world_cache_update.wgsl

@@ -99,13 +99,14 @@ fn sample_random_light_ris(world_position: vec3<f32>, world_normal: vec3<f32>, w
         let resolved_light_sample = unpack_resolved_light_sample(light_tile_resolved_samples[tile_sample], view.exposure);
         let light_contribution = calculate_resolved_light_contribution(resolved_light_sample, world_position, world_normal);
 
-        let target_function = luminance(light_contribution.radiance * saturate(dot(light_contribution.wi, world_normal)));
+        let contribution = light_contribution.radiance * saturate(dot(light_contribution.wi, world_normal));
+        let target_function = luminance(contribution);
         let resampling_weight = mis_weight * (target_function * light_contribution.inverse_pdf);
 
         weight_sum += resampling_weight;
 
         if rand_f(rng) < resampling_weight / weight_sum {
-            selected_sample_radiance = light_contribution.radiance;
+            selected_sample_radiance = contribution;
             selected_sample_target_function = target_function;
             selected_sample_world_position = resolved_light_sample.world_position;
         }