Solari: Fix mirror artifacts

crates/bevy_solari/src/pathtracer/pathtracer.wgsl

@@ -5,7 +5,7 @@
 #import bevy_render::view::View
 #import bevy_solari::brdf::evaluate_brdf
 #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::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX, MIRROR_ROUGHNESS_THRESHOLD}
 
 @group(1) @binding(0) var accumulation_texture: texture_storage_2d<rgba32float, read_write>;
 @group(1) @binding(1) var view_output: texture_storage_2d<rgba16float, write>;
@@ -54,7 +54,7 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
             radiance += mis_weight * throughput * ray_hit.material.emissive;
 
             // Sample direct lighting, but only if the surface is not mirror-like
-            let is_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999;
+            let is_perfectly_specular = ray_hit.material.roughness <= MIRROR_ROUGHNESS_THRESHOLD && 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);
 
@@ -79,8 +79,7 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
             // Update throughput for next bounce
             let brdf = evaluate_brdf(ray_hit.world_normal, wo, next_bounce.wi, ray_hit.material);
-            let cos_theta = dot(next_bounce.wi, ray_hit.world_normal);
-            throughput *= (brdf * cos_theta) / next_bounce.pdf;
+            throughput *= brdf / next_bounce.pdf;
 
             // Russian roulette for early termination
             let p = luminance(throughput);
@@ -105,7 +104,7 @@ struct NextBounce {
 }
 
 fn importance_sample_next_bounce(wo: vec3<f32>, ray_hit: ResolvedRayHitFull, rng: ptr<function, u32>) -> NextBounce {
-    let is_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999;
+    let is_perfectly_specular = ray_hit.material.roughness <= MIRROR_ROUGHNESS_THRESHOLD && ray_hit.material.metallic > 0.9999;
     if is_perfectly_specular {
         return NextBounce(reflect(-wo, ray_hit.world_normal), 1.0, true);
     }

crates/bevy_solari/src/realtime/gbuffer_utils.wgsl

@@ -18,8 +18,7 @@ fn gpixel_resolve(gpixel: vec4<u32>, depth: f32, pixel_id: vec2<u32>, view_size:
 
     let base_rough = unpack4x8unorm(gpixel.r);
     let base_color = pow(base_rough.rgb, vec3(2.2));
-    // Clamp roughness to prevent NaNs
-    let perceptual_roughness = clamp(base_rough.a, 0.0316227766, 1.0); // Clamp roughness to 0.001
+    let perceptual_roughness = base_rough.a;
     let roughness = perceptual_roughness * perceptual_roughness;
     let props = unpack4x8unorm(gpixel.b);
     let reflectance = vec3(props.r);

crates/bevy_solari/src/realtime/restir_di.wgsl

@@ -83,7 +83,7 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
     var brdf: vec3<f32>;
     // If the surface is very smooth, let specular GI handle the specular lobe
     if surface.material.roughness <= SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD {
-        brdf = evaluate_diffuse_brdf(surface.material.base_color, surface.material.metallic);
+        brdf = evaluate_diffuse_brdf(surface.world_normal, merge_result.wi, surface.material.base_color, surface.material.metallic);
     } else {
         brdf = evaluate_brdf(surface.world_normal, wo, merge_result.wi, surface.material);
     }
@@ -111,7 +111,7 @@ fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>
         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 * diffuse_brdf);
+        let target_function = luminance(light_contribution.radiance * diffuse_brdf * saturate(dot(light_contribution.wi, world_normal)));
         let resampling_weight = mis_weight * (target_function * light_contribution.inverse_pdf);
 
         weight_sum += resampling_weight;
@@ -329,6 +329,6 @@ struct ReservoirContribution {
 fn reservoir_contribution(reservoir: Reservoir, world_position: vec3<f32>, world_normal: vec3<f32>, diffuse_brdf: vec3<f32>) -> ReservoirContribution {
     if !reservoir_valid(reservoir) { return ReservoirContribution(vec3(0.0), 0.0, vec3(0.0)); }
     let light_contribution = resolve_and_calculate_light_contribution(reservoir.sample, world_position, world_normal);
-    let target_function = luminance(light_contribution.radiance * diffuse_brdf);
+    let target_function = luminance(light_contribution.radiance * diffuse_brdf * saturate(dot(light_contribution.wi, world_normal)));
     return ReservoirContribution(light_contribution.radiance, target_function, light_contribution.wi);
 }

crates/bevy_solari/src/realtime/restir_gi.wgsl

@@ -69,7 +69,7 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
     gi_reservoirs_a[pixel_index] = combined_reservoir;
 #endif
 
-    let brdf = evaluate_diffuse_brdf(surface.material.base_color, surface.material.metallic);
+    let brdf = evaluate_diffuse_brdf(surface.world_normal, merge_result.wi, surface.material.base_color, surface.material.metallic);
 
     var pixel_color = textureLoad(view_output, global_id.xy);
     pixel_color += vec4(merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight * view.exposure * brdf, 0.0);
@@ -98,7 +98,7 @@ fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>
 
 #ifdef NO_WORLD_CACHE
     let direct_lighting = sample_random_light(sample_point.world_position, sample_point.world_normal, rng);
-    reservoir.radiance = direct_lighting.radiance;
+    reservoir.radiance = direct_lighting.radiance * saturate(dot(direct_lighting.wi, sample_point.world_normal));
     reservoir.unbiased_contribution_weight = direct_lighting.inverse_pdf * uniform_hemisphere_inverse_pdf();
 #else
     reservoir.radiance = query_world_cache(sample_point.world_position, sample_point.geometric_world_normal, view.world_position, ray.t, WORLD_CACHE_CELL_LIFETIME, rng);
@@ -219,6 +219,7 @@ fn empty_reservoir() -> Reservoir {
 struct ReservoirMergeResult {
     merged_reservoir: Reservoir,
     selected_sample_radiance: vec3<f32>,
+    wi: vec3<f32>,
 }
 
 fn merge_reservoirs(
@@ -233,8 +234,10 @@ fn merge_reservoirs(
     rng: ptr<function, u32>,
 ) -> ReservoirMergeResult {
     // Radiances for resampling
-    let canonical_sample_radiance = canonical_reservoir.radiance * saturate(dot(normalize(canonical_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal));
-    let other_sample_radiance = other_reservoir.radiance * saturate(dot(normalize(other_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal));
+    let canonical_sample_wi = normalize(canonical_reservoir.sample_point_world_position - canonical_world_position);
+    let other_sample_wi = normalize(other_reservoir.sample_point_world_position - canonical_world_position);
+    let canonical_sample_radiance = canonical_reservoir.radiance * saturate(dot(canonical_sample_wi, canonical_world_normal));
+    let other_sample_radiance = other_reservoir.radiance * saturate(dot(other_sample_wi, canonical_world_normal));
 
     // Target functions for resampling and MIS
     let canonical_target_function_canonical_sample = luminance(canonical_sample_radiance * canonical_diffuse_brdf);
@@ -264,7 +267,7 @@ fn merge_reservoirs(
 
     // Don't merge samples with huge jacobians, as it explodes the variance
     if canonical_target_function_other_sample_jacobian > 1.2 {
-        return ReservoirMergeResult(canonical_reservoir, canonical_sample_radiance);
+        return ReservoirMergeResult(canonical_reservoir, canonical_sample_radiance, canonical_sample_wi);
     }
 
     // Resampling weight for canonical sample
@@ -294,7 +297,7 @@ fn merge_reservoirs(
         let inverse_target_function = select(0.0, 1.0 / canonical_target_function_other_sample, canonical_target_function_other_sample > 0.0);
         combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
 
-        return ReservoirMergeResult(combined_reservoir, other_sample_radiance);
+        return ReservoirMergeResult(combined_reservoir, other_sample_radiance, other_sample_wi);
     } else {
         combined_reservoir.sample_point_world_position = canonical_reservoir.sample_point_world_position;
         combined_reservoir.sample_point_world_normal = canonical_reservoir.sample_point_world_normal;
@@ -303,6 +306,6 @@ fn merge_reservoirs(
         let inverse_target_function = select(0.0, 1.0 / canonical_target_function_canonical_sample, canonical_target_function_canonical_sample > 0.0);
         combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
 
-        return ReservoirMergeResult(combined_reservoir, canonical_sample_radiance);
+        return ReservoirMergeResult(combined_reservoir, canonical_sample_radiance, canonical_sample_wi);
     }
 }

crates/bevy_solari/src/realtime/specular_gi.wgsl

@@ -7,7 +7,7 @@
 #import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf}
 #import bevy_solari::gbuffer_utils::{gpixel_resolve, ResolvedGPixel}
 #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::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX, MIRROR_ROUGHNESS_THRESHOLD}
 #import bevy_solari::world_cache::{query_world_cache, get_cell_size, WORLD_CACHE_CELL_LIFETIME}
 #import bevy_solari::realtime_bindings::{view_output, gi_reservoirs_a, gbuffer, depth_buffer, view, constants}
 #ifdef DLSS_RR_GUIDE_BUFFERS
@@ -63,8 +63,7 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
     let brdf = evaluate_specular_brdf(surface.world_normal, wo, wi, surface.material.base_color, surface.material.metallic,
         surface.material.reflectance, surface.material.perceptual_roughness, surface.material.roughness);
-    let cos_theta = saturate(dot(wi, surface.world_normal));
-    radiance *= brdf * cos_theta * view.exposure;
+    radiance *= brdf * view.exposure;
 
     var pixel_color = textureLoad(view_output, global_id.xy);
     pixel_color += vec4(radiance, 0.0);
@@ -110,15 +109,15 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
         radiance += throughput * mis_weight * ray_hit.material.emissive;
 
         // Should not perform NEE for mirror-like surfaces
-        surface_perfect_mirror = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999;
+        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
-        if !psr_finished && primary_surface.material.roughness <= 0.001 && primary_surface.material.metallic > 0.9999 {
+        if !psr_finished && primary_surface.material.roughness <= MIRROR_ROUGHNESS_THRESHOLD && primary_surface.material.metallic > 0.9999 {
             if surface_perfect_mirror {
                 mirror_rotations = mirror_rotations * reflection_matrix(ray_hit.world_normal);
             } else {

crates/bevy_solari/src/realtime/world_cache_update.wgsl

@@ -99,7 +99,7 @@ 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);
+        let target_function = luminance(light_contribution.radiance * saturate(dot(light_contribution.wi, world_normal)));
         let resampling_weight = mis_weight * (target_function * light_contribution.inverse_pdf);
 
         weight_sum += resampling_weight;

crates/bevy_solari/src/scene/brdf.wgsl

@@ -3,15 +3,15 @@
 #import bevy_pbr::lighting::{F_AB, D_GGX, V_SmithGGXCorrelated, fresnel, specular_multiscatter}
 #import bevy_pbr::pbr_functions::{calculate_diffuse_color, calculate_F0}
 #import bevy_render::maths::PI
-#import bevy_solari::scene_bindings::ResolvedMaterial
+#import bevy_solari::scene_bindings::{ResolvedMaterial, MIRROR_ROUGHNESS_THRESHOLD}
 
 fn evaluate_brdf(
     world_normal: vec3<f32>,
     wo: vec3<f32>,
     wi: vec3<f32>,
     material: ResolvedMaterial,
 ) -> vec3<f32> {
-    let diffuse_brdf = evaluate_diffuse_brdf(material.base_color, material.metallic);
+    let diffuse_brdf = evaluate_diffuse_brdf(world_normal, wi, material.base_color, material.metallic);
     let specular_brdf = evaluate_specular_brdf(
         world_normal,
         wo,
@@ -25,9 +25,9 @@ fn evaluate_brdf(
     return diffuse_brdf + specular_brdf;
 }
 
-fn evaluate_diffuse_brdf(base_color: vec3<f32>, metallic: f32) -> vec3<f32> {
-    let diffuse_color = calculate_diffuse_color(base_color, metallic, 0.0, 0.0);
-    return diffuse_color / PI;
+fn evaluate_diffuse_brdf(N: vec3<f32>, L: vec3<f32>, base_color: vec3<f32>, metallic: f32) -> vec3<f32> {
+    let diffuse_color = calculate_diffuse_color(base_color, metallic, 0.0, 0.0) / PI;
+    return diffuse_color * saturate(dot(N, L));
 }
 
 fn evaluate_specular_brdf(
@@ -47,10 +47,14 @@ fn evaluate_specular_brdf(
     let NdotV = max(dot(N, V), 0.0001);
 
     let F0 = calculate_F0(base_color, metallic, reflectance);
-    let F_ab = F_AB(perceptual_roughness, NdotV);
+    let F = fresnel(F0, LdotH);
+
+    if roughness <= MIRROR_ROUGHNESS_THRESHOLD {
+        return F;
+    }
 
     let D = D_GGX(roughness, NdotH);
     let Vs = V_SmithGGXCorrelated(roughness, NdotV, NdotL);
-    let F = fresnel(F0, LdotH);
-    return specular_multiscatter(D, Vs, F, F0, F_ab, 1.0);
+    let F_ab = F_AB(perceptual_roughness, NdotV);
+    return specular_multiscatter(D, Vs, F, F0, F_ab, 1.0) * saturate(dot(N, L));
 }

crates/bevy_solari/src/scene/raytracing_scene_bindings.wgsl

@@ -53,6 +53,8 @@ struct Material {
 
 const TEXTURE_MAP_NONE = 0xFFFFFFFFu;
 
+const MIRROR_ROUGHNESS_THRESHOLD = 0.001f;
+
 struct LightSource {
     kind: u32, // 1 bit for kind, 31 bits for extra data
     id: u32,
@@ -145,8 +147,6 @@ fn resolve_material(material: Material, uv: vec2<f32>) -> ResolvedMaterial {
         m.metallic *= metallic_roughness.b;
     }
 
-    // Clamp roughness to prevent NaNs
-    m.perceptual_roughness = clamp(m.perceptual_roughness, 0.0316227766, 1.0); // Clamp roughness to 0.001
     m.roughness = m.perceptual_roughness * m.perceptual_roughness;
 
     return m;

crates/bevy_solari/src/scene/sampling.wgsl

@@ -3,7 +3,7 @@
 #import bevy_pbr::lighting::D_GGX
 #import bevy_pbr::utils::{rand_f, rand_vec2f, rand_u, rand_range_u}
 #import bevy_render::maths::{PI_2, orthonormalize}
-#import bevy_solari::scene_bindings::{trace_ray, RAY_T_MIN, RAY_T_MAX, light_sources, directional_lights, LightSource, LIGHT_SOURCE_KIND_DIRECTIONAL, resolve_triangle_data_full, ResolvedRayHitFull}
+#import bevy_solari::scene_bindings::{trace_ray, RAY_T_MIN, RAY_T_MAX, light_sources, directional_lights, LightSource, LIGHT_SOURCE_KIND_DIRECTIONAL, resolve_triangle_data_full, ResolvedRayHitFull, MIRROR_ROUGHNESS_THRESHOLD}
 
 fn power_heuristic(f: f32, g: f32) -> f32 {
     return balance_heuristic(f * f, g * g);
@@ -19,7 +19,8 @@ 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.001 {
+    // Mirror BRDF case
+    if roughness <= MIRROR_ROUGHNESS_THRESHOLD {
         return vec3(-wi_tangent.xy, wi_tangent.z);
     }
 
@@ -43,6 +44,12 @@ fn sample_ggx_vndf(wi_tangent: vec3<f32>, roughness: f32, rng: ptr<function, u32
 
 // https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 2)
 fn ggx_vndf_pdf(wi_tangent: vec3<f32>, wo_tangent: vec3<f32>, roughness: f32) -> f32 {
+    // Mirror BRDF case
+    if roughness <= MIRROR_ROUGHNESS_THRESHOLD {
+        let mirror_wo = vec3(-wi_tangent.xy, wi_tangent.z);
+        return f32(all(abs(mirror_wo - wo_tangent) < vec3(0.0001)));
+    }
+
     let i = wi_tangent;
     let o = wo_tangent;
     let m = normalize(i + o);
@@ -171,11 +178,10 @@ fn calculate_resolved_light_contribution(resolved_light_sample: ResolvedLightSam
     let light_distance = length(ray);
     let wi = ray / light_distance;
 
-    let cos_theta_origin = saturate(dot(wi, origin_world_normal));
     let cos_theta_light = saturate(dot(-wi, resolved_light_sample.world_normal));
     let light_distance_squared = light_distance * light_distance;
 
-    let radiance = resolved_light_sample.radiance * cos_theta_origin * (cos_theta_light / light_distance_squared);
+    let radiance = resolved_light_sample.radiance * (cos_theta_light / light_distance_squared);
 
     return LightContribution(radiance, resolved_light_sample.inverse_pdf, wi, resolved_light_sample.world_position.w == 1.0);
 }