Solari: Merge reservoir file, reformulate confidence weight (#19895)

Some misc cleanup in preparation for future PRs.

* Merge reservoir.wgsl with restir_di.wgsl, as the reservoir is going to
be DI-specific and won't be reused for GI
* Reformulate confidence weights to not multiply by INITIAL_SAMPLES. The
multiplication cancels out, it doesn't matter.

Author: JMS55

crates/bevy_solari/src/realtime/mod.rs

@@ -13,7 +13,6 @@ use bevy_ecs::{component::Component, reflect::ReflectComponent, schedule::IntoSc
 use bevy_pbr::DefaultOpaqueRendererMethod;
 use bevy_reflect::{std_traits::ReflectDefault, Reflect};
 use bevy_render::{
-    load_shader_library,
     render_graph::{RenderGraphApp, ViewNodeRunner},
     renderer::RenderDevice,
     view::Hdr,
@@ -29,7 +28,6 @@ pub struct SolariLightingPlugin;
 impl Plugin for SolariLightingPlugin {
     fn build(&self, app: &mut App) {
         embedded_asset!(app, "restir_di.wgsl");
-        load_shader_library!(app, "reservoir.wgsl");
 
         app.register_type::<SolariLighting>()
             .insert_resource(DefaultOpaqueRendererMethod::deferred());

crates/bevy_solari/src/realtime/reservoir.wgsl

@@ -1,12 +1,13 @@
+// https://intro-to-restir.cwyman.org/presentations/2023ReSTIR_Course_Notes.pdf
+
 #import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
 #import bevy_pbr::pbr_deferred_types::unpack_24bit_normal
 #import bevy_pbr::prepass_bindings::PreviousViewUniforms
 #import bevy_pbr::rgb9e5::rgb9e5_to_vec3_
 #import bevy_pbr::utils::{rand_f, octahedral_decode}
 #import bevy_render::maths::PI
 #import bevy_render::view::View
-#import bevy_solari::reservoir::{Reservoir, empty_reservoir, reservoir_valid, merge_reservoirs}
-#import bevy_solari::sampling::{generate_random_light_sample, calculate_light_contribution, trace_light_visibility, sample_disk}
+#import bevy_solari::sampling::{LightSample, generate_random_light_sample, calculate_light_contribution, trace_light_visibility, sample_disk}
 #import bevy_solari::scene_bindings::{previous_frame_light_id_translations, LIGHT_NOT_PRESENT_THIS_FRAME}
 
 @group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, write>;
@@ -24,7 +25,9 @@ var<push_constant> constants: PushConstants;
 
 const INITIAL_SAMPLES = 32u;
 const SPATIAL_REUSE_RADIUS_PIXELS = 30.0;
-const CONFIDENCE_WEIGHT_CAP = 20.0 * f32(INITIAL_SAMPLES);
+const CONFIDENCE_WEIGHT_CAP = 20.0;
+
+const NULL_RESERVOIR_SAMPLE = 0xFFFFFFFFu;
 
 @compute @workgroup_size(8, 8, 1)
 fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3<u32>) {
@@ -112,7 +115,7 @@ fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>
         reservoir.unbiased_contribution_weight *= reservoir.visibility;
     }
 
-    reservoir.confidence_weight = f32(INITIAL_SAMPLES);
+    reservoir.confidence_weight = 1.0;
     return reservoir;
 }
 
@@ -205,3 +208,82 @@ fn depth_ndc_to_view_z(ndc_depth: f32) -> f32 {
     return view_pos.z / view_pos.w;
 #endif
 }
+
+// Don't adjust the size of this struct without also adjusting RESERVOIR_STRUCT_SIZE.
+struct Reservoir {
+    sample: LightSample,
+    weight_sum: f32,
+    confidence_weight: f32,
+    unbiased_contribution_weight: f32,
+    visibility: f32,
+}
+
+fn empty_reservoir() -> Reservoir {
+    return Reservoir(
+        LightSample(vec2(NULL_RESERVOIR_SAMPLE, 0u), vec2(0.0)),
+        0.0,
+        0.0,
+        0.0,
+        0.0
+    );
+}
+
+fn reservoir_valid(reservoir: Reservoir) -> bool {
+    return reservoir.sample.light_id.x != NULL_RESERVOIR_SAMPLE;
+}
+
+struct ReservoirMergeResult {
+    merged_reservoir: Reservoir,
+    selected_sample_radiance: vec3<f32>,
+}
+
+fn merge_reservoirs(
+    canonical_reservoir: Reservoir,
+    other_reservoir: Reservoir,
+    world_position: vec3<f32>,
+    world_normal: vec3<f32>,
+    diffuse_brdf: vec3<f32>,
+    rng: ptr<function, u32>,
+) -> ReservoirMergeResult {
+    // TODO: Balance heuristic MIS weights
+    let mis_weight_denominator = 1.0 / (canonical_reservoir.confidence_weight + other_reservoir.confidence_weight);
+
+    let canonical_mis_weight = canonical_reservoir.confidence_weight * mis_weight_denominator;
+    let canonical_target_function = reservoir_target_function(canonical_reservoir, world_position, world_normal, diffuse_brdf);
+    let canonical_resampling_weight = canonical_mis_weight * (canonical_target_function.a * canonical_reservoir.unbiased_contribution_weight);
+
+    let other_mis_weight = other_reservoir.confidence_weight * mis_weight_denominator;
+    let other_target_function = reservoir_target_function(other_reservoir, world_position, world_normal, diffuse_brdf);
+    let other_resampling_weight = other_mis_weight * (other_target_function.a * other_reservoir.unbiased_contribution_weight);
+
+    var combined_reservoir = empty_reservoir();
+    combined_reservoir.weight_sum = canonical_resampling_weight + other_resampling_weight;
+    combined_reservoir.confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
+
+    // https://yusuketokuyoshi.com/papers/2024/Efficient_Visibility_Reuse_for_Real-time_ReSTIR_(Supplementary_Document).pdf
+    combined_reservoir.visibility = max(0.0, (canonical_reservoir.visibility * canonical_resampling_weight
+        + other_reservoir.visibility * other_resampling_weight) / combined_reservoir.weight_sum);
+
+    if rand_f(rng) < other_resampling_weight / combined_reservoir.weight_sum {
+        combined_reservoir.sample = other_reservoir.sample;
+
+        let inverse_target_function = select(0.0, 1.0 / other_target_function.a, other_target_function.a > 0.0);
+        combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
+
+        return ReservoirMergeResult(combined_reservoir, other_target_function.rgb);
+    } else {
+        combined_reservoir.sample = canonical_reservoir.sample;
+
+        let inverse_target_function = select(0.0, 1.0 / canonical_target_function.a, canonical_target_function.a > 0.0);
+        combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
+
+        return ReservoirMergeResult(combined_reservoir, canonical_target_function.rgb);
+    }
+}
+
+fn reservoir_target_function(reservoir: Reservoir, world_position: vec3<f32>, world_normal: vec3<f32>, diffuse_brdf: vec3<f32>) -> vec4<f32> {
+    if !reservoir_valid(reservoir) { return vec4(0.0); }
+    let light_contribution = calculate_light_contribution(reservoir.sample, world_position, world_normal).radiance;
+    let target_function = luminance(light_contribution * diffuse_brdf);
+    return vec4(light_contribution, target_function);
+}