From 8351da45f88ef4f4364a33d1fef37c00cba65e02 Mon Sep 17 00:00:00 2001 From: JMS55 <47158642+JMS55@users.noreply.github.com> Date: Mon, 30 Jun 2025 21:17:16 -0700 Subject: [PATCH] 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. --- crates/bevy_solari/src/realtime/mod.rs | 2 - .../bevy_solari/src/realtime/reservoir.wgsl | 88 ------------------ .../bevy_solari/src/realtime/restir_di.wgsl | 90 ++++++++++++++++++- 3 files changed, 86 insertions(+), 94 deletions(-) delete mode 100644 crates/bevy_solari/src/realtime/reservoir.wgsl diff --git a/crates/bevy_solari/src/realtime/mod.rs b/crates/bevy_solari/src/realtime/mod.rs index 9308ab5cf8..b6a8f27a31 100644 --- a/crates/bevy_solari/src/realtime/mod.rs +++ b/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::() .insert_resource(DefaultOpaqueRendererMethod::deferred()); diff --git a/crates/bevy_solari/src/realtime/reservoir.wgsl b/crates/bevy_solari/src/realtime/reservoir.wgsl deleted file mode 100644 index 16887667a5..0000000000 --- a/crates/bevy_solari/src/realtime/reservoir.wgsl +++ /dev/null @@ -1,88 +0,0 @@ -// https://intro-to-restir.cwyman.org/presentations/2023ReSTIR_Course_Notes.pdf - -#define_import_path bevy_solari::reservoir - -#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance -#import bevy_pbr::utils::rand_f -#import bevy_solari::sampling::{LightSample, calculate_light_contribution} - -const NULL_RESERVOIR_SAMPLE = 0xFFFFFFFFu; - -// 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, -} - -fn merge_reservoirs( - canonical_reservoir: Reservoir, - other_reservoir: Reservoir, - world_position: vec3, - world_normal: vec3, - diffuse_brdf: vec3, - rng: ptr, -) -> 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, world_normal: vec3, diffuse_brdf: vec3) -> vec4 { - 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); -} diff --git a/crates/bevy_solari/src/realtime/restir_di.wgsl b/crates/bevy_solari/src/realtime/restir_di.wgsl index fbc6d6820f..70de4564cc 100644 --- a/crates/bevy_solari/src/realtime/restir_di.wgsl +++ b/crates/bevy_solari/src/realtime/restir_di.wgsl @@ -1,3 +1,5 @@ +// 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 @@ -5,8 +7,7 @@ #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; @@ -24,7 +25,9 @@ var 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) { @@ -112,7 +115,7 @@ fn generate_initial_reservoir(world_position: vec3, world_normal: vec3 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, +} + +fn merge_reservoirs( + canonical_reservoir: Reservoir, + other_reservoir: Reservoir, + world_position: vec3, + world_normal: vec3, + diffuse_brdf: vec3, + rng: ptr, +) -> 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, world_normal: vec3, diffuse_brdf: vec3) -> vec4 { + 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); +}