Solari: Fix MIS not using consistent PDF measures

crates/bevy_solari/src/pathtracer/pathtracer.wgsl

@@ -50,7 +50,7 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
             // Emissive contribution
             var mis_weight = 1.0;
             if p_bounce != 0.0 { // Not first bounce
-                let p_light = random_emissive_light_pdf(ray_hit);
+                let p_light = random_emissive_light_pdf(ray_hit, ray.t, NdotV);
                 mis_weight = power_heuristic(p_bounce, p_light);
             }
             radiance += mis_weight * throughput * ray_hit.material.emissive;
@@ -64,7 +64,7 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
                 mis_weight = 1.0;
                 if direct_lighting.brdf_rays_can_hit {
                     let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material, F_ab);
-                    mis_weight = power_heuristic(1.0 / direct_lighting.inverse_pdf, pdf_of_bounce);
+                    mis_weight = power_heuristic(1.0 / direct_lighting.inverse_solid_angle_pdf, pdf_of_bounce);
                 }
 
                 let direct_lighting_brdf = evaluate_brdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material, F_ab);
@@ -95,4 +95,3 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
     textureStore(accumulation_texture, global_id.xy, vec4(new_color, old_color.a + 1.0));
     textureStore(view_output, global_id.xy, vec4(new_color, 1.0));
 }
-

crates/bevy_solari/src/realtime/restir_gi.wgsl

@@ -8,7 +8,7 @@ enable wgpu_ray_query;
 #import bevy_render::view::View
 #import bevy_solari::brdf::{evaluate_diffuse_brdf, F_AB}
 #import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar, permute_pixel}
-#import bevy_solari::sampling::{sample_random_light, trace_point_visibility, balance_heuristic, isnan}
+#import bevy_solari::sampling::{sample_random_light, trace_point_visibility, balance_heuristic, isinf, isnan}
 #import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
 #import bevy_solari::world_cache::{query_world_cache, WORLD_CACHE_CELL_LIFETIME}
 #import bevy_solari::realtime_bindings::{view_output, gi_reservoirs_a, gi_reservoirs_b, gbuffer, depth_buffer, motion_vectors, previous_gbuffer, previous_depth_buffer, view, previous_view, constants, Reservoir}
@@ -208,10 +208,6 @@ fn jacobian(
     return select(jacobian, 0.0, isinf(jacobian) || isnan(jacobian));
 }
 
-fn isinf(x: f32) -> bool {
-    return (bitcast<u32>(x) & 0x7fffffffu) == 0x7f800000u;
-}
-
 fn empty_reservoir() -> Reservoir {
     return Reservoir(
         vec3(0.0),

crates/bevy_solari/src/realtime/specular_gi.wgsl

@@ -113,7 +113,7 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
         let F_ab = F_AB(ray_hit.material.perceptual_roughness, NdotV);
 
         // Add emissive contribution
-        let mis_weight = emissive_mis_weight(i, primary_surface.material.roughness, p_bounce, ray_hit);
+        let mis_weight = emissive_mis_weight(i, primary_surface.material.roughness, p_bounce, ray_hit, ray.t, NdotV);
         radiance += throughput * mis_weight * ray_hit.material.emissive;
 
         // Should not perform NEE for mirror-like surfaces
@@ -145,7 +145,7 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
             // 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(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material, F_ab);
-            let mis_weight = nee_mis_weight(direct_lighting.inverse_pdf, direct_lighting.brdf_rays_can_hit, wo_tangent, direct_lighting.wi, ray_hit, TBN);
+            let mis_weight = nee_mis_weight(direct_lighting.inverse_solid_angle_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;
         }
 
@@ -172,9 +172,9 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
     return radiance;
 }
 
-fn emissive_mis_weight(i: u32, initial_roughness: f32, p_bounce: f32, ray_hit: ResolvedRayHitFull) -> f32 {
+fn emissive_mis_weight(i: u32, initial_roughness: f32, p_bounce: f32, ray_hit: ResolvedRayHitFull, ray_distance: f32, NdotV: f32) -> f32 {
     if i != 0u {
-        let p_light = random_emissive_light_pdf(ray_hit);
+        let p_light = random_emissive_light_pdf(ray_hit, ray_distance, NdotV);
         return power_heuristic(p_bounce, p_light);
     } else {
         // The first bounce gets MIS weight 0.0 or 1.0 depending on if ReSTIR DI shaded using the specular lobe or not

crates/bevy_solari/src/scene/sampling.wgsl

@@ -83,6 +83,10 @@ fn ggx_vndf_pdf(wi_tangent: vec3<f32>, wo_tangent: vec3<f32>, roughness: f32) ->
     return select(pdf, 0.0, isnan(pdf));
 }
 
+fn isinf(x: f32) -> bool {
+    return (bitcast<u32>(x) & 0x7fffffffu) == 0x7f800000u;
+}
+
 fn isnan(x: f32) -> bool {
     return (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
 }
@@ -104,6 +108,7 @@ struct ResolvedLightSample {
 struct LightContribution {
     radiance: vec3<f32>,
     inverse_pdf: f32,
+    inverse_solid_angle_pdf: f32,
     wi: vec3<f32>,
     brdf_rays_can_hit: bool,
 }
@@ -125,9 +130,10 @@ fn sample_random_light(ray_origin: vec3<f32>, origin_world_normal: vec3<f32>, rn
     return light_contribution;
 }
 
-fn random_emissive_light_pdf(hit: ResolvedRayHitFull) -> f32 {
+fn random_emissive_light_pdf(hit: ResolvedRayHitFull, ray_distance: f32, NdotV: f32) -> f32 {
     let light_count = arrayLength(&light_sources);
-    return 1.0 / (f32(light_count) * f32(hit.triangle_count) * hit.triangle_area);
+    let area_pdf = 1.0 / (f32(light_count) * f32(hit.triangle_count) * hit.triangle_area);
+    return area_pdf * (ray_distance * ray_distance) / NdotV;
 }
 
 fn generate_random_light_sample(rng: ptr<function, u32>) -> GenerateRandomLightSampleResult {
@@ -201,10 +207,12 @@ fn calculate_resolved_light_contribution(resolved_light_sample: ResolvedLightSam
 
     let cos_theta_light = saturate(dot(-wi, resolved_light_sample.world_normal));
     let light_distance_squared = light_distance * light_distance;
+    let denominator = cos_theta_light / light_distance_squared;
 
-    let radiance = resolved_light_sample.radiance * (cos_theta_light / light_distance_squared);
+    let radiance = resolved_light_sample.radiance * denominator;
+    let inverse_solid_angle_pdf = resolved_light_sample.inverse_pdf * denominator;
 
-    return LightContribution(radiance, resolved_light_sample.inverse_pdf, wi, resolved_light_sample.world_position.w == 1.0);
+    return LightContribution(radiance, resolved_light_sample.inverse_pdf, inverse_solid_angle_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 {