Various Solari improvements (#21649)

* Fix compile error when compiling with DLSS enabled after
#21205
* Use permutation sampling for ReSTIR DI temporal reuse to fix artifacts
under DLSS-RR
* For both DI and GI, removed the spatial raytrace, and moved it to the
final reservoir before shading.
* Reduced DI initial samples 32 -> 8 for better performance at the cost
of quality
* Various specular GI improvements and bugfixes (still kinda terrible
overall, I need to do some research on how people usually do this kind
of thing)
* Made the world cache adapt faster / be less stable
* Switched spatial hashing collisions from to linear probing

---------

Co-authored-by: Jasmine Schweitzer <[email protected]>

Author: JMS55

crates/bevy_solari/src/realtime/gbuffer_utils.wgsl

@@ -44,3 +44,12 @@ fn pixel_dissimilar(depth: f32, world_position: vec3<f32>, other_world_position:
 
     return tangent_plane_distance / view_z > 0.003 || dot(normal, other_normal) < 0.906;
 }
+
+fn permute_pixel(pixel_id: vec2<u32>, frame_index: u32, view_size: vec2<f32>) -> vec2<u32> {
+    let r = frame_index;
+    let offset = vec2(r & 3u, (r >> 2u) & 3u);
+    var shifted_pixel_id = pixel_id + offset;
+    shifted_pixel_id ^= vec2(3u);
+    shifted_pixel_id -= offset;
+    return min(shifted_pixel_id, vec2<u32>(view_size - 1.0));
+}

crates/bevy_solari/src/realtime/node.rs

@@ -204,7 +204,8 @@ impl ViewNode for SolariLightingNode {
         let bind_group_resolve_dlss_rr_textures = view_dlss_rr_textures.map(|d| {
             render_context.render_device().create_bind_group(
                 "solari_lighting_bind_group_resolve_dlss_rr_textures",
-                &self.bind_group_layout_resolve_dlss_rr_textures,
+                &pipeline_cache
+                    .get_bind_group_layout(&self.bind_group_layout_resolve_dlss_rr_textures),
                 &BindGroupEntries::sequential((
                     &d.diffuse_albedo.default_view,
                     &d.specular_albedo.default_view,

crates/bevy_solari/src/realtime/restir_di.wgsl

@@ -7,7 +7,7 @@
 #import bevy_render::maths::PI
 #import bevy_render::view::View
 #import bevy_solari::brdf::evaluate_brdf
-#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar}
+#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar, permute_pixel}
 #import bevy_solari::presample_light_tiles::{ResolvedLightSamplePacked, unpack_resolved_light_sample}
 #import bevy_solari::sampling::{LightSample, calculate_resolved_light_contribution, resolve_and_calculate_light_contribution, resolve_light_sample, trace_light_visibility}
 #import bevy_solari::scene_bindings::{light_sources, previous_frame_light_id_translations, LIGHT_NOT_PRESENT_THIS_FRAME}
@@ -27,7 +27,7 @@
 struct PushConstants { frame_index: u32, reset: u32 }
 var<push_constant> constants: PushConstants;
 
-const INITIAL_SAMPLES = 32u;
+const INITIAL_SAMPLES = 8u;
 const SPATIAL_REUSE_RADIUS_PIXELS = 30.0;
 const CONFIDENCE_WEIGHT_CAP = 20.0;
 
@@ -73,7 +73,12 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
     let input_reservoir = load_reservoir_b(global_id.xy);
     let spatial_reservoir = load_spatial_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal, &rng);
     let merge_result = merge_reservoirs(input_reservoir, spatial_reservoir, surface.world_position, surface.world_normal, diffuse_brdf, &rng);
-    let combined_reservoir = merge_result.merged_reservoir;
+    var combined_reservoir = merge_result.merged_reservoir;
+
+    if reservoir_valid(combined_reservoir) {
+        let resolved_light_sample = resolve_light_sample(combined_reservoir.sample, light_sources[combined_reservoir.sample.light_id >> 16u]);
+        combined_reservoir.unbiased_contribution_weight *= trace_light_visibility(surface.world_position, resolved_light_sample.world_position);
+    }
 
     store_reservoir_a(global_id.xy, combined_reservoir);
 
@@ -133,7 +138,7 @@ fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>
 fn load_temporal_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> Reservoir {
     let motion_vector = textureLoad(motion_vectors, pixel_id, 0).xy;
     let temporal_pixel_id_float = round(vec2<f32>(pixel_id) - (motion_vector * view.main_pass_viewport.zw));
-    let temporal_pixel_id = vec2<u32>(temporal_pixel_id_float);
+    let temporal_pixel_id = permute_pixel(vec2<u32>(temporal_pixel_id_float), constants.frame_index, view.viewport.zw);
 
     // Check if the current pixel was off screen during the previous frame (current pixel is newly visible),
     // or if all temporal history should assumed to be invalid
@@ -173,14 +178,7 @@ fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<
         return empty_reservoir();
     }
 
-    var spatial_reservoir = load_reservoir_b(spatial_pixel_id);
-
-    if reservoir_valid(spatial_reservoir) {
-        let resolved_light_sample = resolve_light_sample(spatial_reservoir.sample, light_sources[spatial_reservoir.sample.light_id >> 16u]);
-        spatial_reservoir.unbiased_contribution_weight *= trace_light_visibility(world_position, resolved_light_sample.world_position);
-    }
-
-    return spatial_reservoir;
+    return load_reservoir_b(spatial_pixel_id);
 }
 
 fn get_neighbor_pixel_id(center_pixel_id: vec2<u32>, rng: ptr<function, u32>) -> vec2<u32> {

crates/bevy_solari/src/realtime/restir_gi.wgsl

@@ -6,7 +6,7 @@
 #import bevy_render::maths::PI
 #import bevy_render::view::View
 #import bevy_solari::brdf::evaluate_diffuse_brdf
-#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar}
+#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar, permute_pixel}
 #import bevy_solari::sampling::{sample_random_light, trace_point_visibility}
 #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
@@ -67,7 +67,9 @@ 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.base_color / PI,
         spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.diffuse_brdf, &rng);
-    let combined_reservoir = merge_result.merged_reservoir;
+    var combined_reservoir = merge_result.merged_reservoir;
+
+    combined_reservoir.radiance *= trace_point_visibility(surface.world_position, combined_reservoir.sample_point_world_position);
 
     gi_reservoirs_a[pixel_index] = combined_reservoir;
 
@@ -76,10 +78,6 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
     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);
     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) * view.exposure, 1.0));
-#endif
 }
 
 fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> Reservoir {
@@ -107,7 +105,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);
+    reservoir.radiance = query_world_cache(sample_point.world_position, sample_point.geometric_world_normal, view.world_position, rng);
     reservoir.unbiased_contribution_weight = uniform_hemisphere_inverse_pdf();
 #endif
 
@@ -120,43 +118,28 @@ fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>
 fn load_temporal_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> NeighborInfo {
     let motion_vector = textureLoad(motion_vectors, pixel_id, 0).xy;
     let temporal_pixel_id_float = round(vec2<f32>(pixel_id) - (motion_vector * view.main_pass_viewport.zw));
+    let temporal_pixel_id = permute_pixel(vec2<u32>(temporal_pixel_id_float), constants.frame_index, view.viewport.zw);
 
     // Check if the current pixel was off screen during the previous frame (current pixel is newly visible),
     // or if all temporal history should assumed to be invalid
     if any(temporal_pixel_id_float < vec2(0.0)) || any(temporal_pixel_id_float >= view.main_pass_viewport.zw) || bool(constants.reset) {
         return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
     }
 
-    let temporal_pixel_id_base = vec2<u32>(round(temporal_pixel_id_float));
-    for (var i = 0u; i < 4u; i++) {
-        let temporal_pixel_id = permute_pixel(temporal_pixel_id_base, i);
-
-        // Check if the pixel features have changed heavily between the current and previous frame
-        let temporal_depth = textureLoad(previous_depth_buffer, temporal_pixel_id, 0);
-        let temporal_surface = gpixel_resolve(textureLoad(previous_gbuffer, temporal_pixel_id, 0), temporal_depth, temporal_pixel_id, view.main_pass_viewport.zw, previous_view.world_from_clip);
-        let temporal_diffuse_brdf = temporal_surface.material.base_color / PI;
-        if pixel_dissimilar(depth, world_position, temporal_surface.world_position, world_normal, temporal_surface.world_normal, view) {
-            continue;
-        }
-
-        let temporal_pixel_index = temporal_pixel_id.x + temporal_pixel_id.y * u32(view.main_pass_viewport.z);
-        var temporal_reservoir = gi_reservoirs_a[temporal_pixel_index];
-
-        temporal_reservoir.confidence_weight = min(temporal_reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP);
-
-        return NeighborInfo(temporal_reservoir, temporal_surface.world_position, temporal_surface.world_normal, temporal_diffuse_brdf);
+    // Check if the pixel features have changed heavily between the current and previous frame
+    let temporal_depth = textureLoad(previous_depth_buffer, temporal_pixel_id, 0);
+    let temporal_surface = gpixel_resolve(textureLoad(previous_gbuffer, temporal_pixel_id, 0), temporal_depth, temporal_pixel_id, view.main_pass_viewport.zw, previous_view.world_from_clip);
+    let temporal_diffuse_brdf = temporal_surface.material.base_color / PI;
+    if pixel_dissimilar(depth, world_position, temporal_surface.world_position, world_normal, temporal_surface.world_normal, view) {
+        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
     }
 
-    return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
-}
+    let temporal_pixel_index = temporal_pixel_id.x + temporal_pixel_id.y * u32(view.main_pass_viewport.z);
+    var temporal_reservoir = gi_reservoirs_a[temporal_pixel_index];
 
-fn permute_pixel(pixel_id: vec2<u32>, i: u32) -> vec2<u32> {
-    let r = constants.frame_index + i;
-    let offset = vec2(r & 3u, (r >> 2u) & 3u);
-    var shifted_pixel_id = pixel_id + offset;
-    shifted_pixel_id ^= vec2(3u);
-    shifted_pixel_id -= offset;
-    return min(shifted_pixel_id, vec2<u32>(view.main_pass_viewport.zw - 1.0));
+    temporal_reservoir.confidence_weight = min(temporal_reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP);
+
+    return NeighborInfo(temporal_reservoir, temporal_surface.world_position, temporal_surface.world_normal, temporal_diffuse_brdf);
 }
 
 fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> NeighborInfo {
@@ -170,9 +153,7 @@ fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<
     }
 
     let spatial_pixel_index = spatial_pixel_id.x + spatial_pixel_id.y * u32(view.main_pass_viewport.z);
-    var spatial_reservoir = gi_reservoirs_b[spatial_pixel_index];
-
-    spatial_reservoir.radiance *= trace_point_visibility(world_position, spatial_reservoir.sample_point_world_position);
+    let spatial_reservoir = gi_reservoirs_b[spatial_pixel_index];
 
     return NeighborInfo(spatial_reservoir, spatial_surface.world_position, spatial_surface.world_normal, spatial_diffuse_brdf);
 }
@@ -252,27 +233,19 @@ 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_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));
 
     // Target functions for resampling and MIS
     let canonical_target_function_canonical_sample = luminance(canonical_sample_radiance * canonical_diffuse_brdf);
     let canonical_target_function_other_sample = luminance(other_sample_radiance * canonical_diffuse_brdf);
 
     // Extra target functions for MIS
     let other_target_function_canonical_sample = luminance(
-        canonical_reservoir.radiance *
-        saturate(dot(normalize(canonical_reservoir.sample_point_world_position - other_world_position), other_world_normal)) *
-        other_diffuse_brdf
+        canonical_reservoir.radiance * saturate(dot(normalize(canonical_reservoir.sample_point_world_position - other_world_position), other_world_normal)) * other_diffuse_brdf
     );
     let other_target_function_other_sample = luminance(
-        other_reservoir.radiance *
-        saturate(dot(normalize(other_reservoir.sample_point_world_position - other_world_position), other_world_normal)) *
-        other_diffuse_brdf
+        other_reservoir.radiance * saturate(dot(normalize(other_reservoir.sample_point_world_position - other_world_position), other_world_normal)) * other_diffuse_brdf
     );
 
     // Jacobians for resampling and MIS
@@ -299,19 +272,14 @@ fn merge_reservoirs(
         canonical_reservoir.confidence_weight * canonical_target_function_canonical_sample,
         other_reservoir.confidence_weight * other_target_function_canonical_sample * other_target_function_canonical_sample_jacobian,
     );
-    let canonical_sample_resampling_weight = canonical_sample_mis_weight *
-        canonical_target_function_canonical_sample *
-        canonical_reservoir.unbiased_contribution_weight;
+    let canonical_sample_resampling_weight = canonical_sample_mis_weight * canonical_target_function_canonical_sample * canonical_reservoir.unbiased_contribution_weight;
 
     // Resampling weight for other sample
     let other_sample_mis_weight = balance_heuristic(
         other_reservoir.confidence_weight * other_target_function_other_sample,
         canonical_reservoir.confidence_weight * canonical_target_function_other_sample * canonical_target_function_other_sample_jacobian,
     );
-    let other_sample_resampling_weight = other_sample_mis_weight *
-        canonical_target_function_other_sample *
-        other_reservoir.unbiased_contribution_weight *
-        canonical_target_function_other_sample_jacobian;
+    let other_sample_resampling_weight = other_sample_mis_weight * canonical_target_function_other_sample * other_reservoir.unbiased_contribution_weight * canonical_target_function_other_sample_jacobian;
 
     // Perform resampling
     var combined_reservoir = empty_reservoir();

crates/bevy_solari/src/realtime/specular_gi.wgsl

@@ -32,7 +32,7 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
     var radiance: vec3<f32>;
     var wi: vec3<f32>;
-    if surface.material.roughness > 0.04 {
+    if surface.material.roughness > 0.1 {
         // Surface is very rough, reuse the ReSTIR GI reservoir
         let gi_reservoir = gi_reservoirs_a[pixel_index];
         wi = normalize(gi_reservoir.sample_point_world_position - surface.world_position);
@@ -59,25 +59,31 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
     var pixel_color = textureLoad(view_output, global_id.xy);
     pixel_color += vec4(radiance, 0.0);
     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, &rng) * view.exposure, 1.0));
+#endif
 }
 
 fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, rng: ptr<function, u32>) -> vec3<f32> {
     var ray_origin = initial_ray_origin;
     var wi = initial_wi;
 
     // Trace up to three bounces, getting the net throughput from them
+    var radiance = vec3(0.0);
     var throughput = vec3(1.0);
     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);
         if ray.kind == RAY_QUERY_INTERSECTION_NONE { break; }
         let ray_hit = resolve_ray_hit_full(ray);
 
+        // Add world cache contribution
+        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, rng);
+
         // Surface is very rough, terminate path in the world cache
-        if ray_hit.material.roughness > 0.04 || i == 2u {
-            let diffuse_brdf = ray_hit.material.base_color / PI;
-            return throughput * diffuse_brdf * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position);
-        }
+        if ray_hit.material.roughness > 0.1 && i != 0u { break; }
 
         // Sample new ray direction from the GGX BRDF for next bounce
         let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent);
@@ -93,11 +99,11 @@ fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, rng:
         // Update throughput for next bounce
         let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
         let brdf = evaluate_brdf(N, wo, wi, ray_hit.material);
-        let cos_theta = dot(wi, N);
+        let cos_theta = saturate(dot(wi, N));
         throughput *= (brdf * cos_theta) / pdf;
     }
 
-    return vec3(0.0);
+    return radiance;
 }
 
 // Don't adjust the size of this struct without also adjusting GI_RESERVOIR_STRUCT_SIZE.