Solari: Improved mirror denoising

crates/bevy_pbr/src/prepass/mod.rs

@@ -205,7 +205,7 @@ pub fn update_previous_view_data(
     }
 }
 
-#[derive(Component, PartialEq, Default)]
+#[derive(Component, PartialEq, Clone, Default)]
 pub struct PreviousGlobalTransform(pub Affine3A);
 
 #[cfg(not(feature = "meshlet"))]

crates/bevy_solari/src/realtime/mod.rs

@@ -46,7 +46,7 @@ impl Plugin for SolariLightingPlugin {
         embedded_asset!(app, "world_cache_update.wgsl");
 
         #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
-        embedded_asset!(app, "resolve_dlss_rr_textures.wgsl");
+        load_shader_library!(app, "resolve_dlss_rr_textures.wgsl");
 
         app.insert_resource(DefaultOpaqueRendererMethod::deferred());
     }

crates/bevy_solari/src/realtime/node.rs

@@ -57,6 +57,8 @@ pub struct SolariLightingNode {
     gi_spatial_and_shade_pipeline: CachedComputePipelineId,
     specular_gi_pipeline: CachedComputePipelineId,
     #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
+    specular_gi_with_psr_pipeline: CachedComputePipelineId,
+    #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
     resolve_dlss_rr_textures_pipeline: CachedComputePipelineId,
 }
 
@@ -109,6 +111,16 @@ impl ViewNode for SolariLightingNode {
         let view_uniforms = world.resource::<ViewUniforms>();
         let previous_view_uniforms = world.resource::<PreviousViewUniforms>();
         let frame_count = world.resource::<FrameCount>();
+
+        #[cfg(not(all(feature = "dlss", not(feature = "force_disable_dlss"))))]
+        let specular_gi_pipeline = self.specular_gi_pipeline;
+        #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
+        let specular_gi_pipeline = if view_dlss_rr_textures.is_some() {
+            self.specular_gi_with_psr_pipeline
+        } else {
+            self.specular_gi_pipeline
+        };
+
         let (
             Some(decay_world_cache_pipeline),
             Some(compact_world_cache_single_block_pipeline),
@@ -143,7 +155,7 @@ impl ViewNode for SolariLightingNode {
             pipeline_cache.get_compute_pipeline(self.di_spatial_and_shade_pipeline),
             pipeline_cache.get_compute_pipeline(self.gi_initial_and_temporal_pipeline),
             pipeline_cache.get_compute_pipeline(self.gi_spatial_and_shade_pipeline),
-            pipeline_cache.get_compute_pipeline(self.specular_gi_pipeline),
+            pipeline_cache.get_compute_pipeline(specular_gi_pipeline),
             &scene_bindings.bind_group,
             view_prepass_textures.deferred_view(),
             view_prepass_textures.depth_view(),
@@ -255,8 +267,8 @@ impl ViewNode for SolariLightingNode {
         );
 
         #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
-        if let Some(bind_group_resolve_dlss_rr_textures) = bind_group_resolve_dlss_rr_textures {
-            pass.set_bind_group(2, &bind_group_resolve_dlss_rr_textures, &[]);
+        if let Some(bind_group_resolve_dlss_rr_textures) = &bind_group_resolve_dlss_rr_textures {
+            pass.set_bind_group(2, bind_group_resolve_dlss_rr_textures, &[]);
             pass.set_pipeline(resolve_dlss_rr_textures_pipeline);
             pass.dispatch_workgroups(dx, dy, 1);
         }
@@ -343,6 +355,10 @@ impl ViewNode for SolariLightingNode {
         d.end(&mut pass);
 
         let d = diagnostics.time_span(&mut pass, "solari_lighting/specular_indirect_lighting");
+        #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
+        if let Some(bind_group_resolve_dlss_rr_textures) = &bind_group_resolve_dlss_rr_textures {
+            pass.set_bind_group(2, bind_group_resolve_dlss_rr_textures, &[]);
+        }
         pass.set_pipeline(specular_gi_pipeline);
         pass.set_push_constants(
             0,
@@ -543,12 +559,20 @@ impl FromWorld for SolariLightingNode {
                 vec![],
             ),
             #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
+            specular_gi_with_psr_pipeline: create_pipeline(
+                "solari_lighting_specular_gi_with_psr_pipeline",
+                "specular_gi",
+                load_embedded_asset!(world, "specular_gi.wgsl"),
+                Some(&bind_group_layout_resolve_dlss_rr_textures),
+                vec!["DLSS_RR_GUIDE_BUFFERS".into()],
+            ),
+            #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
             resolve_dlss_rr_textures_pipeline: create_pipeline(
                 "solari_lighting_resolve_dlss_rr_textures_pipeline",
                 "resolve_dlss_rr_textures",
                 load_embedded_asset!(world, "resolve_dlss_rr_textures.wgsl"),
                 Some(&bind_group_layout_resolve_dlss_rr_textures),
-                vec![],
+                vec!["DLSS_RR_GUIDE_BUFFERS".into()],
             ),
         }
     }

crates/bevy_solari/src/realtime/realtime_bindings.wgsl

@@ -18,6 +18,7 @@
 @group(1) @binding(11) var previous_depth_buffer: texture_depth_2d;
 @group(1) @binding(12) var<uniform> view: View;
 @group(1) @binding(13) var<uniform> previous_view: PreviousViewUniforms;
+
 @group(1) @binding(14) var<storage, read_write> world_cache_checksums: array<atomic<u32>, #{WORLD_CACHE_SIZE}>;
 #ifdef WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER
 @group(1) @binding(15) var<storage, read_write> world_cache_life: array<u32, #{WORLD_CACHE_SIZE}>;
@@ -32,6 +33,14 @@
 @group(1) @binding(21) var<storage, read_write> world_cache_b: array<u32, 1024u>;
 @group(1) @binding(22) var<storage, read_write> world_cache_active_cell_indices: array<u32, #{WORLD_CACHE_SIZE}>;
 @group(1) @binding(23) var<storage, read_write> world_cache_active_cells_count: u32;
+
+#ifdef DLSS_RR_GUIDE_BUFFERS
+@group(2) @binding(0) var diffuse_albedo: texture_storage_2d<rgba8unorm, write>;
+@group(2) @binding(1) var specular_albedo: texture_storage_2d<rgba8unorm, write>;
+@group(2) @binding(2) var normal_roughness: texture_storage_2d<rgba16float, write>;
+@group(2) @binding(3) var specular_motion_vectors: texture_storage_2d<rg16float, write>;
+#endif
+
 struct PushConstants { frame_index: u32, reset: u32 }
 var<push_constant> constants: PushConstants;
 

crates/bevy_solari/src/realtime/resolve_dlss_rr_textures.wgsl

@@ -1,24 +1,22 @@
+#define_import_path bevy_solari::resolve_dlss_rr_textures
+
 #import bevy_pbr::pbr_functions::{calculate_diffuse_color, calculate_F0}
 #import bevy_render::view::View
 #import bevy_solari::gbuffer_utils::gpixel_resolve
-#import bevy_solari::realtime_bindings::{gbuffer, depth_buffer, view}
-
-@group(2) @binding(0) var diffuse_albedo: texture_storage_2d<rgba8unorm, write>;
-@group(2) @binding(1) var specular_albedo: texture_storage_2d<rgba8unorm, write>;
-@group(2) @binding(2) var normal_roughness: texture_storage_2d<rgba16float, write>;
-@group(2) @binding(3) var specular_motion_vectors: texture_storage_2d<rg16float, write>;
+#import bevy_solari::realtime_bindings::{gbuffer, depth_buffer, view, diffuse_albedo, specular_albedo, normal_roughness, specular_motion_vectors}
 
 @compute @workgroup_size(8, 8, 1)
 fn resolve_dlss_rr_textures(@builtin(global_invocation_id) global_id: vec3<u32>) {
     let pixel_id = global_id.xy;
     if any(pixel_id >= vec2u(view.main_pass_viewport.zw)) { return; }
 
+    textureStore(specular_motion_vectors, pixel_id, vec4(0.0));
+
     let depth = textureLoad(depth_buffer, global_id.xy, 0);
     if depth == 0.0 {
         textureStore(diffuse_albedo, pixel_id, vec4(0.0));
         textureStore(specular_albedo, pixel_id, vec4(0.5));
         textureStore(normal_roughness, pixel_id, vec4(0.0));
-        textureStore(specular_motion_vectors, pixel_id, vec4(0.0));
         return;
     }
 
@@ -29,7 +27,6 @@ fn resolve_dlss_rr_textures(@builtin(global_invocation_id) global_id: vec3<u32>)
     textureStore(diffuse_albedo, pixel_id, vec4(calculate_diffuse_color(surface.material.base_color, surface.material.metallic, 0.0, 0.0), 0.0));
     textureStore(specular_albedo, pixel_id, vec4(env_brdf_approx2(F0, surface.material.roughness, surface.world_normal, wo), 0.0));
     textureStore(normal_roughness, pixel_id, vec4(surface.world_normal, surface.material.perceptual_roughness));
-    textureStore(specular_motion_vectors, pixel_id, vec4(0.0)); // TODO
 }
 
 fn env_brdf_approx2(specular_color: vec3<f32>, alpha: f32, N: vec3<f32>, V: vec3<f32>) -> vec3<f32> {

crates/bevy_solari/src/realtime/specular_gi.wgsl

@@ -1,14 +1,19 @@
 #define_import_path bevy_solari::specular_gi
 
-#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace
+#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}
-#import bevy_solari::gbuffer_utils::gpixel_resolve
+#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::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
+#import bevy_solari::realtime_bindings::{diffuse_albedo, specular_albedo, normal_roughness, specular_motion_vectors, previous_view}
+#import bevy_solari::resolve_dlss_rr_textures::env_brdf_approx2
+#endif
 
 const DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD: f32 = 0.4;
 const SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD: f32 = 0.0225;
@@ -53,7 +58,7 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
         var a0 = dot(wo_unnormalized, wo_unnormalized) / (4.0 * PI * cos_theta);
         a0 *= TERMINATE_IN_WORLD_CACHE_THRESHOLD;
 
-        radiance = trace_glossy_path(surface.world_position, wi, surface.material.roughness, pdf, a0, &rng) / pdf;
+        radiance = trace_glossy_path(global_id.xy, surface, wi, pdf, a0, &rng) / pdf;
     }
 
     let brdf = evaluate_specular_brdf(surface.world_normal, wo, wi, surface.material.base_color, surface.material.metallic,
@@ -70,16 +75,22 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
 #endif
 }
 
-fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, initial_roughness: f32, initial_p_bounce: f32, a0: f32, rng: ptr<function, u32>) -> vec3<f32> {
-    var ray_origin = initial_ray_origin;
+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> {
+    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_perfectly_specular = false;
+    var surface_perfect_mirror = false;
     var path_spread = 0.0;
 
-    // Trace up to three bounces, getting the net throughput from them
-    var radiance = vec3(0.0);
-    var throughput = vec3(1.0);
+#ifdef DLSS_RR_GUIDE_BUFFERS
+    var mirror_rotations = reflection_matrix(primary_surface.world_normal);
+    var psr_finished = false;
+#endif
+
+    // Trace up to three bounces
     for (var i = 0u; i < 3u; i += 1u) {
         // Trace ray
         let ray = trace_ray(ray_origin, wi, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE);
@@ -95,21 +106,47 @@ fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, initi
         let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
 
         // Add emissive contribution
-        let mis_weight = emissive_mis_weight(i, initial_roughness, p_bounce, ray_hit, surface_perfectly_specular);
+        let mis_weight = emissive_mis_weight(i, primary_surface.material.roughness, p_bounce, ray_hit, surface_perfect_mirror);
         radiance += throughput * mis_weight * 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;
+        surface_perfect_mirror = ray_hit.material.roughness <= 0.001 && 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/#DLSS_RR_GUIDE_BUFFERS
+#ifdef DLSS_RR_GUIDE_BUFFERS
+        if !psr_finished && primary_surface.material.roughness <= 0.001 && primary_surface.material.metallic > 0.9999 {
+            if surface_perfect_mirror {
+                mirror_rotations = mirror_rotations * reflection_matrix(ray_hit.world_normal);
+            } else {
+                psr_finished = true;
+
+                // Simplification: Apply all rotations in the chain around the first mirror, rather than applying each rotation around its respective mirror
+                let virtual_position = (mirror_rotations * (ray_hit.world_position - primary_surface.world_position)) + primary_surface.world_position;
+                let virtual_previous_frame_position = (mirror_rotations * (ray_hit.previous_frame_world_position - primary_surface.world_position)) + primary_surface.world_position;
+                let specular_motion_vector = calculate_motion_vector(virtual_position, virtual_previous_frame_position);
+
+                let F0 = calculate_F0(ray_hit.material.base_color, ray_hit.material.metallic, ray_hit.material.reflectance);
+                let wo = normalize(view.world_position - virtual_position);
+                let virtual_normal = normalize(mirror_rotations * ray_hit.world_normal);
+
+                textureStore(specular_motion_vectors, pixel_id, vec4(specular_motion_vector, vec2(0.0)));
+                textureStore(diffuse_albedo, pixel_id, vec4(calculate_diffuse_color(ray_hit.material.base_color, ray_hit.material.metallic, 0.0, 0.0), 0.0));
+                textureStore(specular_albedo, pixel_id, vec4(env_brdf_approx2(F0, ray_hit.material.roughness, ray_hit.world_normal, wo), 0.0));
+                textureStore(normal_roughness, pixel_id, vec4(virtual_normal, ray_hit.material.perceptual_roughness));
+            }
+        }
+#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
             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 {
+        } else if !surface_perfect_mirror {
             // 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);
@@ -132,9 +169,9 @@ fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, initi
     return radiance;
 }
 
-fn emissive_mis_weight(i: u32, initial_roughness: f32, p_bounce: f32, ray_hit: ResolvedRayHitFull, previous_surface_perfectly_specular: bool) -> f32 {
+fn emissive_mis_weight(i: u32, initial_roughness: f32, p_bounce: f32, ray_hit: ResolvedRayHitFull, previous_surface_perfect_mirror: bool) -> f32 {
     if i != 0u {
-        if previous_surface_perfectly_specular { return 1.0; }
+        if previous_surface_perfect_mirror { return 1.0; }
 
         let p_light = random_emissive_light_pdf(ray_hit);
         return power_heuristic(p_bounce, p_light);
@@ -163,12 +200,30 @@ fn nee_mis_weight(inverse_p_light: f32, brdf_rays_can_hit: bool, wo_tangent: vec
     return 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>,
-    weight_sum: f32,
-    radiance: vec3<f32>,
-    confidence_weight: f32,
-    sample_point_world_normal: vec3<f32>,
-    unbiased_contribution_weight: f32,
+#ifdef DLSS_RR_GUIDE_BUFFERS
+// https://en.wikipedia.org/wiki/Householder_transformation
+fn reflection_matrix(plane_normal: vec3f) -> mat3x3<f32> {
+    // N times Nᵀ.
+    let n_nt = mat3x3<f32>(
+        plane_normal * plane_normal.x,
+        plane_normal * plane_normal.y,
+        plane_normal * plane_normal.z,
+    );
+    let identity_matrix = mat3x3<f32>(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0);
+    return identity_matrix - n_nt * 2.0;
 }
+
+fn calculate_motion_vector(world_position: vec3<f32>, previous_world_position: vec3<f32>) -> vec2<f32> {
+    let clip_position_t = view.unjittered_clip_from_world * vec4(world_position, 1.0);
+    let clip_position = clip_position_t.xy / clip_position_t.w;
+    let previous_clip_position_t = previous_view.clip_from_world * vec4(previous_world_position, 1.0);
+    let previous_clip_position = previous_clip_position_t.xy / previous_clip_position_t.w;
+    // These motion vectors are used as offsets to UV positions and are stored
+    // in the range -1,1 to allow offsetting from the one corner to the
+    // diagonally-opposite corner in UV coordinates, in either direction.
+    // A difference between diagonally-opposite corners of clip space is in the
+    // range -2,2, so this needs to be scaled by 0.5. And the V direction goes
+    // down where clip space y goes up, so y needs to be flipped.
+    return (clip_position - previous_clip_position) * vec2(0.5, -0.5);
+}
+#endif