Solari: Light leak reduction and stochastic WC update

crates/bevy_solari/src/realtime/node.rs

@@ -48,7 +48,8 @@ pub struct SolariLightingNode {
     compact_world_cache_single_block_pipeline: CachedComputePipelineId,
     compact_world_cache_blocks_pipeline: CachedComputePipelineId,
     compact_world_cache_write_active_cells_pipeline: CachedComputePipelineId,
-    sample_for_world_cache_pipeline: CachedComputePipelineId,
+    sample_di_for_world_cache_pipeline: CachedComputePipelineId,
+    sample_gi_for_world_cache_pipeline: CachedComputePipelineId,
     blend_new_world_cache_samples_pipeline: CachedComputePipelineId,
     presample_light_tiles_pipeline: CachedComputePipelineId,
     di_initial_and_temporal_pipeline: CachedComputePipelineId,
@@ -114,7 +115,8 @@ impl ViewNode for SolariLightingNode {
             Some(compact_world_cache_single_block_pipeline),
             Some(compact_world_cache_blocks_pipeline),
             Some(compact_world_cache_write_active_cells_pipeline),
-            Some(sample_for_world_cache_pipeline),
+            Some(sample_di_for_world_cache_pipeline),
+            Some(sample_gi_for_world_cache_pipeline),
             Some(blend_new_world_cache_samples_pipeline),
             Some(presample_light_tiles_pipeline),
             Some(di_initial_and_temporal_pipeline),
@@ -136,7 +138,8 @@ impl ViewNode for SolariLightingNode {
             pipeline_cache.get_compute_pipeline(self.compact_world_cache_blocks_pipeline),
             pipeline_cache
                 .get_compute_pipeline(self.compact_world_cache_write_active_cells_pipeline),
-            pipeline_cache.get_compute_pipeline(self.sample_for_world_cache_pipeline),
+            pipeline_cache.get_compute_pipeline(self.sample_di_for_world_cache_pipeline),
+            pipeline_cache.get_compute_pipeline(self.sample_gi_for_world_cache_pipeline),
             pipeline_cache.get_compute_pipeline(self.blend_new_world_cache_samples_pipeline),
             pipeline_cache.get_compute_pipeline(self.presample_light_tiles_pipeline),
             pipeline_cache.get_compute_pipeline(self.di_initial_and_temporal_pipeline),
@@ -288,7 +291,17 @@ impl ViewNode for SolariLightingNode {
 
         pass.set_bind_group(2, None, &[]);
 
-        pass.set_pipeline(sample_for_world_cache_pipeline);
+        pass.set_pipeline(sample_di_for_world_cache_pipeline);
+        pass.set_push_constants(
+            0,
+            bytemuck::cast_slice(&[frame_index, solari_lighting.reset as u32]),
+        );
+        pass.dispatch_workgroups_indirect(
+            &solari_lighting_resources.world_cache_active_cells_dispatch,
+            0,
+        );
+
+        pass.set_pipeline(sample_gi_for_world_cache_pipeline);
         pass.set_push_constants(
             0,
             bytemuck::cast_slice(&[frame_index, solari_lighting.reset as u32]),
@@ -351,6 +364,14 @@ impl ViewNode for SolariLightingNode {
         pass.dispatch_workgroups(dx, dy, 1);
         d.end(&mut pass);
 
+        drop(pass);
+
+        diagnostics.record_u32(
+            render_context.command_encoder(),
+            &s.world_cache_active_cells_count.slice(..),
+            "solari_lighting/world_cache_active_cells_count",
+        );
+
         Ok(())
     }
 }
@@ -486,9 +507,16 @@ impl FromWorld for SolariLightingNode {
                 Some(&bind_group_layout_world_cache_active_cells_dispatch),
                 vec!["WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER".into()],
             ),
-            sample_for_world_cache_pipeline: create_pipeline(
-                "solari_lighting_sample_for_world_cache_pipeline",
-                "sample_radiance",
+            sample_di_for_world_cache_pipeline: create_pipeline(
+                "solari_lighting_sample_di_for_world_cache_pipeline",
+                "sample_di",
+                load_embedded_asset!(world, "world_cache_update.wgsl"),
+                None,
+                vec![],
+            ),
+            sample_gi_for_world_cache_pipeline: create_pipeline(
+                "solari_lighting_sample_gi_for_world_cache_pipeline",
+                "sample_gi",
                 load_embedded_asset!(world, "world_cache_update.wgsl"),
                 None,
                 vec!["WORLD_CACHE_QUERY_ATOMIC_MAX_LIFETIME".into()],
@@ -526,7 +554,7 @@ impl FromWorld for SolariLightingNode {
                 "initial_and_temporal",
                 load_embedded_asset!(world, "restir_gi.wgsl"),
                 None,
-                vec![],
+                vec!["WORLD_CACHE_FIRST_BOUNCE_LIGHT_LEAK_PREVENTION".into()],
             ),
             gi_spatial_and_shade_pipeline: create_pipeline(
                 "solari_lighting_gi_spatial_and_shade_pipeline",

crates/bevy_solari/src/realtime/prepare.rs

@@ -214,7 +214,7 @@ pub fn prepare_solari_lighting_resources(
         let world_cache_active_cells_count = render_device.create_buffer(&BufferDescriptor {
             label: Some("solari_lighting_world_cache_active_cells_count"),
             size: size_of::<u32>() as u64,
-            usage: BufferUsages::STORAGE,
+            usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC,
             mapped_at_creation: false,
         });
 

crates/bevy_solari/src/realtime/restir_gi.wgsl

@@ -101,7 +101,7 @@ fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>
     reservoir.radiance = direct_lighting.radiance;
     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, WORLD_CACHE_CELL_LIFETIME, rng);
+    reservoir.radiance = query_world_cache(sample_point.world_position, sample_point.geometric_world_normal, view.world_position, ray.t, WORLD_CACHE_CELL_LIFETIME, rng);
     reservoir.unbiased_contribution_weight = uniform_hemisphere_inverse_pdf();
 #endif
 

crates/bevy_solari/src/realtime/specular_gi.wgsl

@@ -66,7 +66,7 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
     textureStore(view_output, global_id.xy, pixel_color);
 
 #ifdef VISUALIZE_WORLD_CACHE
-    textureStore(view_output, global_id.xy, vec4(query_world_cache(surface.world_position, surface.world_normal, view.world_position, WORLD_CACHE_CELL_LIFETIME, &rng) * view.exposure, 1.0));
+    textureStore(view_output, global_id.xy, vec4(query_world_cache(surface.world_position, surface.world_normal, view.world_position, RAY_T_MAX, WORLD_CACHE_CELL_LIFETIME, &rng) * view.exposure, 1.0));
 #endif
 }
 
@@ -107,7 +107,7 @@ fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, initi
         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);
+            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;
         } else if !surface_perfectly_specular {
             // Sample direct lighting (NEE)

crates/bevy_solari/src/realtime/world_cache_query.wgsl

@@ -21,27 +21,37 @@ const WORLD_CACHE_MAX_TEMPORAL_SAMPLES: f32 = 32.0;
 const WORLD_CACHE_DIRECT_LIGHT_SAMPLE_COUNT: u32 = 32u;
 /// Maximum amount of distance to trace GI rays between two cache cells
 const WORLD_CACHE_MAX_GI_RAY_DISTANCE: f32 = 50.0;
+/// Soft upper limit on the amount of cache cells to update each frame
+const WORLD_CACHE_CELL_UPDATES_SOFT_CAP: u32 = 40000u;
 
 /// Maximum amount of frames a cell can live for without being queried
-const WORLD_CACHE_CELL_LIFETIME: u32 = 30u;
+const WORLD_CACHE_CELL_LIFETIME: u32 = 10u;
 /// Maximum amount of attempts to find a cache entry after a hash collision
 const WORLD_CACHE_MAX_SEARCH_STEPS: u32 = 3u;
 
 /// Size of a cache cell at the lowest LOD in meters
-const WORLD_CACHE_POSITION_BASE_CELL_SIZE: f32 = 0.25;
+const WORLD_CACHE_POSITION_BASE_CELL_SIZE: f32 = 0.15;
 /// How fast the world cache transitions between LODs as a function of distance to the camera
-const WORLD_CACHE_POSITION_LOD_SCALE: f32 = 8.0;
+const WORLD_CACHE_POSITION_LOD_SCALE: f32 = 15.0;
 
 /// Marker value for an empty cell
 const WORLD_CACHE_EMPTY_CELL: u32 = 0u;
 
 #ifndef WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER
-fn query_world_cache(world_position_in: vec3<f32>, world_normal: vec3<f32>, view_position: vec3<f32>, cell_lifetime: u32, rng: ptr<function, u32>) -> vec3<f32> {
+fn query_world_cache(world_position_in: vec3<f32>, world_normal: vec3<f32>, view_position: vec3<f32>, ray_t: f32, cell_lifetime: u32, rng: ptr<function, u32>) -> vec3<f32> {
     var world_position = world_position_in;
     var cell_size = get_cell_size(world_position, view_position);
 
+#ifdef WORLD_CACHE_FIRST_BOUNCE_LIGHT_LEAK_PREVENTION
+    if ray_t < cell_size {
+        // Prevent light leaks
+        cell_size = WORLD_CACHE_POSITION_BASE_CELL_SIZE;
+    }
+#endif
+
+#ifndef NO_JITTER_WORLD_CACHE
+    // Jitter query point, which essentially blurs the cache a bit so it's not so grid-like
     // https://tomclabault.github.io/blog/2025/regir, jitter_world_position_tangent_plane
-#ifdef JITTER_WORLD_CACHE
     let TBN = orthonormalize(world_normal);
     let offset = (rand_vec2f(rng) * 2.0 - 1.0) * cell_size * 0.5;
     world_position += offset.x * TBN[0] + offset.y * TBN[1];

crates/bevy_solari/src/realtime/world_cache_update.wgsl

@@ -8,6 +8,7 @@
     WORLD_CACHE_MAX_TEMPORAL_SAMPLES,
     WORLD_CACHE_DIRECT_LIGHT_SAMPLE_COUNT,
     WORLD_CACHE_MAX_GI_RAY_DISTANCE,
+    WORLD_CACHE_CELL_UPDATES_SOFT_CAP,
     query_world_cache,
 }
 #import bevy_solari::realtime_bindings::{
@@ -24,26 +25,38 @@
 }
 
 @compute @workgroup_size(64, 1, 1)
-fn sample_radiance(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) active_cell_id: vec3<u32>) {
+fn sample_di(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) active_cell_id: vec3<u32>) {
     if active_cell_id.x >= world_cache_active_cells_count { return; }
 
     let cell_index = world_cache_active_cell_indices[active_cell_id.x];
     let geometry_data = world_cache_geometry_data[cell_index];
     var rng = cell_index + constants.frame_index;
 
-    var new_radiance = sample_random_light_ris(geometry_data.world_position, geometry_data.world_normal, workgroup_id.xy, &rng);
+    if rand_f(&rng) >= f32(WORLD_CACHE_CELL_UPDATES_SOFT_CAP) / f32(world_cache_active_cells_count) { return; }
+
+    let new_radiance = sample_random_light_ris(geometry_data.world_position, geometry_data.world_normal, workgroup_id.xy, &rng);
+
+    world_cache_active_cells_new_radiance[active_cell_id.x] = new_radiance;
+}
+
+@compute @workgroup_size(64, 1, 1)
+fn sample_gi(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) active_cell_id: vec3<u32>) {
+    if active_cell_id.x >= world_cache_active_cells_count { return; }
+
+    let cell_index = world_cache_active_cell_indices[active_cell_id.x];
+    let geometry_data = world_cache_geometry_data[cell_index];
+    var rng = cell_index + constants.frame_index;
+
+    if rand_f(&rng) >= f32(WORLD_CACHE_CELL_UPDATES_SOFT_CAP) / f32(world_cache_active_cells_count) { return; }
 
-#ifndef NO_MULTIBOUNCE
     let ray_direction = sample_cosine_hemisphere(geometry_data.world_normal, &rng);
     let ray = trace_ray(geometry_data.world_position, ray_direction, RAY_T_MIN, WORLD_CACHE_MAX_GI_RAY_DISTANCE, RAY_FLAG_NONE);
     if ray.kind != RAY_QUERY_INTERSECTION_NONE {
         let ray_hit = resolve_ray_hit_full(ray);
         let cell_life = atomicLoad(&world_cache_life[cell_index]);
-        new_radiance += ray_hit.material.base_color * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, cell_life, &rng);
+        let radiance = query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, ray.t, cell_life, &rng);
+        world_cache_active_cells_new_radiance[active_cell_id.x] += ray_hit.material.base_color * radiance;
     }
-#endif
-
-    world_cache_active_cells_new_radiance[active_cell_id.x] = new_radiance;
 }
 
 @compute @workgroup_size(64, 1, 1)
@@ -53,6 +66,8 @@ fn blend_new_samples(@builtin(global_invocation_id) active_cell_id: vec3<u32>) {
     let cell_index = world_cache_active_cell_indices[active_cell_id.x];
     var rng = cell_index + constants.frame_index;
 
+    if rand_f(&rng) >= f32(WORLD_CACHE_CELL_UPDATES_SOFT_CAP) / f32(world_cache_active_cells_count) { return; }
+
     let old_radiance = world_cache_radiance[cell_index];
     let new_radiance = world_cache_active_cells_new_radiance[active_cell_id.x];
     let luminance_delta = world_cache_luminance_deltas[cell_index];
@@ -64,7 +79,7 @@ fn blend_new_samples(@builtin(global_invocation_id) active_cell_id: vec3<u32>) {
     let blend_amount = 1.0 / min(sample_count, max_sample_count);
 
     let blended_radiance = mix(old_radiance.rgb, new_radiance, blend_amount);
-    let blended_luminance_delta = mix(luminance_delta, luminance(blended_radiance - old_radiance.rgb), 1.0 / 8.0);
+    let blended_luminance_delta = mix(luminance_delta, luminance(blended_radiance) - luminance(old_radiance.rgb), 1.0 / 8.0);
 
     world_cache_radiance[cell_index] = vec4(blended_radiance, sample_count);
     world_cache_luminance_deltas[cell_index] = blended_luminance_delta;

examples/3d/solari.rs

@@ -602,5 +602,18 @@ fn update_performance_text(
         "render/solari_lighting/specular_indirect_lighting/elapsed_gpu",
     );
     text.push_str(&format!("{:17}     TODO\n", "DLSS-RR"));
-    text.push_str(&format!("\n{:17}  {total:.2} ms", "Total"));
+    text.push_str(&format!("{:17}  {total:.2} ms\n", "Total"));
+
+    if let Some(world_cache_active_cells_count) = diagnostics
+        .get(&DiagnosticPath::new(
+            "render/solari_lighting/world_cache_active_cells_count",
+        ))
+        .and_then(Diagnostic::average)
+    {
+        text.push_str(&format!(
+            "\nWorld cache cells {} ({:.0}%)",
+            world_cache_active_cells_count as u32,
+            (world_cache_active_cells_count * 100.0) / (2u64.pow(20) as f64)
+        ));
+    }
 }