Solari: Reduce temporal shadow lag

crates/bevy_solari/src/realtime/initial_path.wgsl

@@ -9,7 +9,7 @@ enable wgpu_ray_query;
 #import bevy_solari::brdf::{brdf_pdf, evaluate_and_sample_brdf, evaluate_brdf, F_AB}
 #import bevy_solari::presample_light_tiles::unpack_resolved_light_sample
 #import bevy_solari::realtime_bindings::{empty_reservoir, light_tile_resolved_samples, light_tile_samples, Reservoir, constants, view}
-#import bevy_solari::sampling::{calculate_resolved_light_contribution, isinf, LightSample, NULL_LIGHT_ID, power_heuristic, trace_light_visibility}
+#import bevy_solari::sampling::{calculate_resolved_light_contribution, isinf, LightSample, NULL_LIGHT_ID, power_heuristic, trace_visibility}
 #import bevy_solari::scene_bindings::{light_sources, MIRROR_ROUGHNESS_THRESHOLD, RAY_T_MAX, RAY_T_MIN, resolve_ray_hit_full, ResolvedMaterial, ResolvedRayHitFull, trace_ray}
 #import bevy_solari::world_cache::{get_cell_size, query_world_cache, WORLD_CACHE_CELL_LIFETIME}
 #ifdef DLSS_RR_GUIDE_BUFFERS
@@ -274,7 +274,7 @@ fn sample_light_ris(ray_origin: vec3<f32>, normal: vec3<f32>, wo: vec3<f32>, mat
     var unbiased_contribution_weight = 0.0;
     if selected_target_function > 0.0 {
         unbiased_contribution_weight = weight_sum / selected_target_function;
-        unbiased_contribution_weight *= trace_light_visibility(ray_origin, selected_world_position);
+        unbiased_contribution_weight *= trace_visibility(ray_origin, selected_world_position);
     }
 
     return DiSample(unbiased_contribution_weight, light_tile_samples[selected_tile_sample], selected_wi, selected_brdf_radiance, selected_inverse_solid_angle_pdf, selected_brdf_rays_can_hit);

crates/bevy_solari/src/realtime/node.rs

@@ -492,7 +492,7 @@ pub fn init_solari_lighting_pipelines(
             "spatial_and_shade",
             load_embedded_asset!(asset_server.as_ref(), "restir.wgsl"),
             None,
-            vec![],
+            vec!["SPATIAL_MERGE".into()],
         ),
         #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))]
         resolve_dlss_rr_textures_pipeline: create_pipeline(

crates/bevy_solari/src/realtime/restir.wgsl

@@ -8,7 +8,7 @@ enable wgpu_ray_query;
 #import bevy_solari::gbuffer_utils::{gpixel_resolve, permute_pixel, pixel_dissimilar}
 #import bevy_solari::initial_path::{generate_initial_reservoir, InitialSamplingResult}
 #import bevy_solari::realtime_bindings::{depth_buffer, empty_reservoir, gbuffer, motion_vectors, previous_depth_buffer, previous_gbuffer, previous_view, reservoirs_a, reservoirs_b, Reservoir, constants, view, view_output}
-#import bevy_solari::sampling::{balance_heuristic, calculate_resolved_light_contribution, isinf, isnan, LightSample, NULL_LIGHT_ID, power_heuristic, resolve_light_sample, ResolvedLightSample, trace_light_visibility}
+#import bevy_solari::sampling::{balance_heuristic, calculate_resolved_light_contribution, isinf, isnan, LightSample, NULL_LIGHT_ID, power_heuristic, resolve_light_sample, ResolvedLightSample, trace_visibility, trace_visibility_previous_frame}
 #import bevy_solari::scene_bindings::{light_sources, LIGHT_NOT_PRESENT_THIS_FRAME, previous_frame_light_id_translations, RAY_T_MAX, RAY_T_MIN, ResolvedMaterial}
 #import bevy_solari::world_cache::{query_world_cache, WORLD_CACHE_CELL_LIFETIME}
 
@@ -35,7 +35,7 @@ fn initial_and_temporal(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin
     let previous_camera_homogeneous = previous_view.world_from_clip * (previous_view.clip_from_view * vec4(0.0, 0.0, 0.0, 1.0));
     let previous_camera_world_position = previous_camera_homogeneous.xyz / previous_camera_homogeneous.w;
     let merge_result = merge_reservoirs(initial.reservoir, surface.world_position, surface.world_normal, surface.material,
-        temporal.reservoir, temporal.world_position, temporal.world_normal, temporal.material, previous_camera_world_position, false, &rng);
+        temporal.reservoir, temporal.world_position, temporal.world_normal, temporal.material, previous_camera_world_position, &rng);
 
     reservoirs_b[pixel_index] = merge_result.merged_reservoir;
 }
@@ -61,7 +61,7 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
     let spatial = load_spatial_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal, &rng);
     let merge_result = merge_reservoirs(input_reservoir, surface.world_position, surface.world_normal, surface.material,
-        spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.material, view.world_position, true, &rng);
+        spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.material, view.world_position, &rng);
 
     reservoirs_a[pixel_index] = merge_result.merged_reservoir;
 
@@ -190,7 +190,6 @@ fn merge_reservoirs(
     other_world_normal: vec3<f32>,
     other_material: ResolvedMaterial,
     other_view_position: vec3<f32>,
-    is_spatial: bool,
     rng: ptr<function, u32>,
 ) -> ReservoirMergeResult {
     var canonical_resolved: ResolvedLightSample;
@@ -252,17 +251,24 @@ fn merge_reservoirs(
 
     // Visibility for the cross-domain targets
     if other_sample_at_canonical.target_function > 0.0 && other_sample_at_canonical_jacobian > 0.0 {
-        let visibility = trace_light_visibility(canonical_world_position + canonical_world_normal * RAY_T_MIN, other_sample_at_canonical.sample_world_position);
+        let visibility = trace_visibility(canonical_world_position + canonical_world_normal * RAY_T_MIN, other_sample_at_canonical.sample_world_position);
         other_sample_at_canonical.target_function *= visibility;
     }
     if canonical_sample_at_other.target_function > 0.0 && canonical_sample_at_other_jacobian > 0.0 {
-        let visibility = trace_light_visibility(other_world_position + other_world_normal * RAY_T_MIN, canonical_sample_at_other.sample_world_position);
+#ifdef SPATIAL_MERGE
+        let visibility = trace_visibility(other_world_position + other_world_normal * RAY_T_MIN, canonical_sample_at_other.sample_world_position);
+#else
+        let visibility = trace_visibility_previous_frame(other_world_position + other_world_normal * RAY_T_MIN, canonical_sample_at_other.sample_world_position);
+#endif
         canonical_sample_at_other.target_function *= visibility;
     }
 
     // Defensive balance heuristic MIS (for spatial reuse only)
     let total_confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
-    let defensive_t_c = f32(is_spatial) * select(1.0, canonical_reservoir.confidence_weight / total_confidence_weight, total_confidence_weight > 0.0);
+    var defensive_t_c = 0.0;
+#ifdef SPATIAL_MERGE
+    defensive_t_c = select(1.0, canonical_reservoir.confidence_weight / total_confidence_weight, total_confidence_weight > 0.0);
+#endif
 
     // Resampling weight for canonical sample
     let canonical_balance_mis_weight = balance_heuristic(

crates/bevy_solari/src/realtime/world_cache_update.wgsl

@@ -3,7 +3,7 @@ enable wgpu_ray_query;
 #import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
 #import bevy_pbr::utils::{rand_f, rand_range_u, sample_cosine_hemisphere}
 #import bevy_solari::presample_light_tiles::unpack_resolved_light_sample
-#import bevy_solari::sampling::{calculate_resolved_light_contribution, trace_light_visibility}
+#import bevy_solari::sampling::{calculate_resolved_light_contribution, trace_visibility}
 #import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN}
 #import bevy_solari::world_cache::query_world_cache
 #import bevy_solari::realtime_bindings::{
@@ -115,7 +115,7 @@ fn sample_random_light_ris(world_position: vec3<f32>, world_normal: vec3<f32>, w
         let inverse_target_function = select(0.0, 1.0 / selected_sample_target_function, selected_sample_target_function > 0.0);
         unbiased_contribution_weight = weight_sum * inverse_target_function;
 
-        unbiased_contribution_weight *= trace_light_visibility(world_position + (world_normal * RAY_T_MIN), selected_sample_world_position);
+        unbiased_contribution_weight *= trace_visibility(world_position + (world_normal * RAY_T_MIN), selected_sample_world_position);
     }
 
     return selected_sample_radiance * unbiased_contribution_weight;

crates/bevy_solari/src/scene/binder.rs

@@ -32,6 +32,7 @@ const LIGHT_NOT_PRESENT_THIS_FRAME: u32 = u32::MAX;
 pub struct RaytracingSceneBindings {
     pub bind_group: Option<BindGroup>,
     pub bind_group_layout: BindGroupLayoutDescriptor,
+    previous_frame_tlas: Option<Tlas>,
     previous_frame_light_entities: Vec<Entity>,
 }
 
@@ -58,6 +59,8 @@ pub fn prepare_raytracing_scene_bindings(
 ) {
     raytracing_scene_bindings.bind_group = None;
 
+    let previous_frame_tlas = raytracing_scene_bindings.previous_frame_tlas.take();
+
     let mut this_frame_entity_to_light_id = EntityHashMap::<u32>::default();
     let previous_frame_light_entities: Vec<_> = raytracing_scene_bindings
         .previous_frame_light_entities
@@ -293,6 +296,7 @@ pub fn prepare_raytracing_scene_bindings(
             samplers.as_slice(),
             materials.binding().unwrap(),
             tlas.as_binding(),
+            previous_frame_tlas.as_ref().unwrap_or(&tlas).as_binding(),
             transforms.binding().unwrap(),
             previous_frame_transforms.binding().unwrap(),
             geometry_ids.binding().unwrap(),
@@ -304,6 +308,8 @@ pub fn prepare_raytracing_scene_bindings(
             dfg_sampler,
         )),
     ));
+
+    raytracing_scene_bindings.previous_frame_tlas = Some(tlas);
 }
 
 impl RaytracingSceneBindings {
@@ -322,6 +328,7 @@ impl RaytracingSceneBindings {
                         sampler(SamplerBindingType::Filtering).count(MAX_TEXTURE_COUNT),
                         storage_buffer_read_only_sized(false, None),
                         acceleration_structure(),
+                        acceleration_structure(),
                         storage_buffer_read_only_sized(false, None),
                         storage_buffer_read_only_sized(false, None),
                         storage_buffer_read_only_sized(false, None),
@@ -334,6 +341,7 @@ impl RaytracingSceneBindings {
                     ),
                 ),
             ),
+            previous_frame_tlas: None,
             previous_frame_light_entities: Vec::new(),
         }
     }

crates/bevy_solari/src/scene/raytracing_scene_bindings.wgsl

@@ -81,15 +81,16 @@ const LIGHT_NOT_PRESENT_THIS_FRAME = 0xFFFFFFFFu;
 @group(0) @binding(3) var samplers: binding_array<sampler>;
 @group(0) @binding(4) var<storage> materials: array<Material>;
 @group(0) @binding(5) var tlas: acceleration_structure;
-@group(0) @binding(6) var<storage> transforms: array<mat3x4<f32>>;
-@group(0) @binding(7) var<storage> previous_frame_transforms: array<mat3x4<f32>>;
-@group(0) @binding(8) var<storage> geometry_ids: array<InstanceGeometryIds>;
-@group(0) @binding(9) var<storage> material_ids: array<u32>; // TODO: Store material_id in instance_custom_index instead?
-@group(0) @binding(10) var<storage> light_sources: array<LightSource>;
-@group(0) @binding(11) var<storage> directional_lights: array<DirectionalLight>;
-@group(0) @binding(12) var<storage> previous_frame_light_id_translations: array<u32>;
-@group(0) @binding(13) var brdf_dfg_lut: texture_2d<f32>;
-@group(0) @binding(14) var brdf_dfg_lut_sampler: sampler;
+@group(0) @binding(6) var previous_frame_tlas: acceleration_structure;
+@group(0) @binding(7) var<storage> transforms: array<mat3x4<f32>>;
+@group(0) @binding(8) var<storage> previous_frame_transforms: array<mat3x4<f32>>;
+@group(0) @binding(9) var<storage> geometry_ids: array<InstanceGeometryIds>;
+@group(0) @binding(10) var<storage> material_ids: array<u32>;
+@group(0) @binding(11) var<storage> light_sources: array<LightSource>;
+@group(0) @binding(12) var<storage> directional_lights: array<DirectionalLight>;
+@group(0) @binding(13) var<storage> previous_frame_light_id_translations: array<u32>;
+@group(0) @binding(14) var brdf_dfg_lut: texture_2d<f32>;
+@group(0) @binding(15) var brdf_dfg_lut_sampler: sampler;
 
 const RAY_T_MIN = 0.001f;
 const RAY_T_MAX = 100000.0f;
@@ -104,6 +105,14 @@ fn trace_ray(ray_origin: vec3<f32>, ray_direction: vec3<f32>, ray_t_min: f32, ra
     return rayQueryGetCommittedIntersection(&rq);
 }
 
+fn trace_ray_previous_frame(ray_origin: vec3<f32>, ray_direction: vec3<f32>, ray_t_min: f32, ray_t_max: f32, ray_flag: u32) -> RayIntersection {
+    let ray = RayDesc(ray_flag, RAY_NO_CULL, ray_t_min, ray_t_max, ray_origin, ray_direction);
+    var rq: ray_query;
+    rayQueryInitialize(&rq, previous_frame_tlas, ray);
+    rayQueryProceed(&rq);
+    return rayQueryGetCommittedIntersection(&rq);
+}
+
 fn sample_texture(id: u32, uv: vec2<f32>) -> vec3<f32> {
     return textureSampleLevel(textures[id], samplers[id], uv, 0.0).rgb; // TODO: Mipmap
 }

crates/bevy_solari/src/scene/sampling.wgsl

@@ -5,7 +5,7 @@ enable wgpu_ray_query;
 #import bevy_pbr::lighting::D_GGX
 #import bevy_pbr::utils::{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, MIRROR_ROUGHNESS_THRESHOLD}
+#import bevy_solari::scene_bindings::{trace_ray, trace_ray_previous_frame, 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);
@@ -124,7 +124,7 @@ struct GenerateRandomLightSampleResult {
 fn sample_random_light(ray_origin: vec3<f32>, origin_world_normal: vec3<f32>, rng: ptr<function, u32>) -> LightContribution {
     let sample = generate_random_light_sample(rng);
     var light_contribution = calculate_resolved_light_contribution(sample.resolved_light_sample, ray_origin, origin_world_normal);
-    light_contribution.radiance *= trace_light_visibility(ray_origin, sample.resolved_light_sample.world_position);
+    light_contribution.radiance *= trace_visibility(ray_origin, sample.resolved_light_sample.world_position);
     return light_contribution;
 }
 
@@ -213,11 +213,11 @@ fn calculate_resolved_light_contribution(resolved_light_sample: ResolvedLightSam
     return LightContribution(radiance, resolved_light_sample.inverse_pdf, inverse_solid_angle_pdf, wi, resolved_light_sample.world_position.w == 1.0);
 }
 
-fn trace_light_visibility(ray_origin: vec3<f32>, light_sample_world_position: vec4<f32>) -> f32 {
-    var ray_direction = light_sample_world_position.xyz;
+fn trace_visibility(ray_origin: vec3<f32>, point: vec4<f32>) -> f32 {
+    var ray_direction = point.xyz;
     var ray_t_max = RAY_T_MAX;
 
-    if light_sample_world_position.w == 1.0 {
+    if point.w == 1.0 {
         let ray = ray_direction - ray_origin;
         let dist = length(ray);
         ray_direction = ray / dist;
@@ -230,6 +230,23 @@ fn trace_light_visibility(ray_origin: vec3<f32>, light_sample_world_position: ve
     return f32(ray_hit.kind == RAY_QUERY_INTERSECTION_NONE);
 }
 
+fn trace_visibility_previous_frame(ray_origin: vec3<f32>, point: vec4<f32>) -> f32 {
+    var ray_direction = point.xyz;
+    var ray_t_max = RAY_T_MAX;
+
+    if point.w == 1.0 {
+        let ray = ray_direction - ray_origin;
+        let dist = length(ray);
+        ray_direction = ray / dist;
+        ray_t_max = dist - RAY_T_MIN;
+    }
+
+    if ray_t_max < RAY_T_MIN { return 0.0; }
+
+    let ray_hit = trace_ray_previous_frame(ray_origin, ray_direction, RAY_T_MIN, ray_t_max, RAY_FLAG_TERMINATE_ON_FIRST_HIT);
+    return f32(ray_hit.kind == RAY_QUERY_INTERSECTION_NONE);
+}
+
 // https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec22%3A297
 fn triangle_barycentrics(seed: u32) -> vec3<f32> {
     var rng = seed;