bevy_solari ReSTIR DI

crates/bevy_core_pipeline/src/core_3d/mod.rs

@@ -1026,7 +1026,8 @@ pub fn prepare_prepass_textures(
                         format: CORE_3D_DEPTH_FORMAT,
                         usage: TextureUsages::COPY_DST
                             | TextureUsages::RENDER_ATTACHMENT
-                            | TextureUsages::TEXTURE_BINDING,
+                            | TextureUsages::TEXTURE_BINDING
+                            | TextureUsages::COPY_SRC, // TODO: Remove COPY_SRC, double buffer instead (for bevy_solari)
                         view_formats: &[],
                     };
                     texture_cache.get(&render_device, descriptor)
@@ -1092,7 +1093,8 @@ pub fn prepare_prepass_textures(
                             dimension: TextureDimension::D2,
                             format: DEFERRED_PREPASS_FORMAT,
                             usage: TextureUsages::RENDER_ATTACHMENT
-                                | TextureUsages::TEXTURE_BINDING,
+                                | TextureUsages::TEXTURE_BINDING
+                                | TextureUsages::COPY_SRC, // TODO: Remove COPY_SRC, double buffer instead (for bevy_solari)
                             view_formats: &[],
                         },
                     )

crates/bevy_core_pipeline/src/skybox/skybox_prepass.wgsl

@@ -5,6 +5,9 @@
 struct PreviousViewUniforms {
     view_from_world: mat4x4<f32>,
     clip_from_world: mat4x4<f32>,
+    clip_from_view: mat4x4<f32>,
+    world_from_clip: mat4x4<f32>,
+    view_from_clip: mat4x4<f32>,
 }
 
 @group(0) @binding(0) var<uniform> view: View;

crates/bevy_solari/src/realtime/node.rs

@@ -1,7 +1,9 @@
 use super::{prepare::SolariLightingResources, SolariLighting};
 use crate::scene::RaytracingSceneBindings;
 use bevy_asset::load_embedded_asset;
-use bevy_core_pipeline::prepass::ViewPrepassTextures;
+use bevy_core_pipeline::prepass::{
+    PreviousViewData, PreviousViewUniformOffset, PreviousViewUniforms, ViewPrepassTextures,
+};
 use bevy_diagnostic::FrameCount;
 use bevy_ecs::{
     query::QueryItem,
@@ -16,8 +18,8 @@ use bevy_render::{
             storage_buffer_sized, texture_2d, texture_depth_2d, texture_storage_2d, uniform_buffer,
         },
         BindGroupEntries, BindGroupLayout, BindGroupLayoutEntries, CachedComputePipelineId,
-        ComputePassDescriptor, ComputePipelineDescriptor, PipelineCache, PushConstantRange,
-        ShaderStages, StorageTextureAccess, TextureSampleType,
+        ComputePassDescriptor, ComputePipelineDescriptor, Extent3d, PipelineCache,
+        PushConstantRange, ShaderStages, StorageTextureAccess, TextureSampleType,
     },
     renderer::{RenderContext, RenderDevice},
     view::{ViewTarget, ViewUniform, ViewUniformOffset, ViewUniforms},
@@ -44,6 +46,7 @@ impl ViewNode for SolariLightingNode {
         &'static ViewTarget,
         &'static ViewPrepassTextures,
         &'static ViewUniformOffset,
+        &'static PreviousViewUniformOffset,
     );
 
     fn run(
@@ -57,12 +60,14 @@ impl ViewNode for SolariLightingNode {
             view_target,
             view_prepass_textures,
             view_uniform_offset,
+            previous_view_uniform_offset,
         ): QueryItem<Self::ViewQuery>,
         world: &World,
     ) -> Result<(), NodeRunError> {
         let pipeline_cache = world.resource::<PipelineCache>();
         let scene_bindings = world.resource::<RaytracingSceneBindings>();
         let view_uniforms = world.resource::<ViewUniforms>();
+        let previous_view_uniforms = world.resource::<PreviousViewUniforms>();
         let frame_count = world.resource::<FrameCount>();
         let (
             Some(initial_and_temporal_pipeline),
@@ -73,6 +78,7 @@ impl ViewNode for SolariLightingNode {
             Some(depth_buffer),
             Some(motion_vectors),
             Some(view_uniforms),
+            Some(previous_view_uniforms),
         ) = (
             pipeline_cache.get_compute_pipeline(self.initial_and_temporal_pipeline),
             pipeline_cache.get_compute_pipeline(self.spatial_and_shade_pipeline),
@@ -82,6 +88,7 @@ impl ViewNode for SolariLightingNode {
             view_prepass_textures.depth_view(),
             view_prepass_textures.motion_vectors_view(),
             view_uniforms.uniforms.binding(),
+            previous_view_uniforms.uniforms.binding(),
         )
         else {
             return Ok(());
@@ -97,7 +104,10 @@ impl ViewNode for SolariLightingNode {
                 gbuffer,
                 depth_buffer,
                 motion_vectors,
+                &solari_lighting_resources.previous_gbuffer.1,
+                &solari_lighting_resources.previous_depth.1,
                 view_uniforms,
+                previous_view_uniforms,
             )),
         );
 
@@ -114,7 +124,14 @@ impl ViewNode for SolariLightingNode {
         let pass_span = diagnostics.pass_span(&mut pass, "solari_lighting");
 
         pass.set_bind_group(0, scene_bindings, &[]);
-        pass.set_bind_group(1, &bind_group, &[view_uniform_offset.offset]);
+        pass.set_bind_group(
+            1,
+            &bind_group,
+            &[
+                view_uniform_offset.offset,
+                previous_view_uniform_offset.offset,
+            ],
+        );
 
         pass.set_pipeline(initial_and_temporal_pipeline);
         pass.set_push_constants(
@@ -127,6 +144,39 @@ impl ViewNode for SolariLightingNode {
         pass.dispatch_workgroups(viewport.x.div_ceil(8), viewport.y.div_ceil(8), 1);
 
         pass_span.end(&mut pass);
+        drop(pass);
+
+        // TODO: Remove these copies, and double buffer instead
+        command_encoder.copy_texture_to_texture(
+            view_prepass_textures
+                .deferred
+                .clone()
+                .unwrap()
+                .texture
+                .texture
+                .as_image_copy(),
+            solari_lighting_resources.previous_gbuffer.0.as_image_copy(),
+            Extent3d {
+                width: viewport.x,
+                height: viewport.y,
+                depth_or_array_layers: 1,
+            },
+        );
+        command_encoder.copy_texture_to_texture(
+            view_prepass_textures
+                .depth
+                .clone()
+                .unwrap()
+                .texture
+                .texture
+                .as_image_copy(),
+            solari_lighting_resources.previous_depth.0.as_image_copy(),
+            Extent3d {
+                width: viewport.x,
+                height: viewport.y,
+                depth_or_array_layers: 1,
+            },
+        );
 
         Ok(())
     }
@@ -152,7 +202,10 @@ impl FromWorld for SolariLightingNode {
                     texture_2d(TextureSampleType::Uint),
                     texture_depth_2d(),
                     texture_2d(TextureSampleType::Float { filterable: true }),
+                    texture_2d(TextureSampleType::Uint),
+                    texture_depth_2d(),
                     uniform_buffer::<ViewUniform>(true),
+                    uniform_buffer::<PreviousViewData>(true),
                 ),
             ),
         );

crates/bevy_solari/src/realtime/prepare.rs

@@ -1,14 +1,19 @@
 use super::SolariLighting;
+use bevy_core_pipeline::{core_3d::CORE_3D_DEPTH_FORMAT, deferred::DEFERRED_PREPASS_FORMAT};
 use bevy_ecs::{
     component::Component,
     entity::Entity,
     query::With,
     system::{Commands, Query, Res},
 };
+use bevy_image::ToExtents;
 use bevy_math::UVec2;
 use bevy_render::{
     camera::ExtractedCamera,
-    render_resource::{Buffer, BufferDescriptor, BufferUsages},
+    render_resource::{
+        Buffer, BufferDescriptor, BufferUsages, Texture, TextureDescriptor, TextureDimension,
+        TextureUsages, TextureView, TextureViewDescriptor,
+    },
     renderer::RenderDevice,
 };
 
@@ -20,6 +25,8 @@ const RESERVOIR_STRUCT_SIZE: u64 = 32;
 pub struct SolariLightingResources {
     pub reservoirs_a: Buffer,
     pub reservoirs_b: Buffer,
+    pub previous_gbuffer: (Texture, TextureView),
+    pub previous_depth: (Texture, TextureView),
     pub view_size: UVec2,
 }
 
@@ -56,9 +63,35 @@ pub fn prepare_solari_lighting_resources(
             mapped_at_creation: false,
         });
 
+        let previous_gbuffer = render_device.create_texture(&TextureDescriptor {
+            label: Some("solari_lighting_previous_gbuffer"),
+            size: view_size.to_extents(),
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: TextureDimension::D2,
+            format: DEFERRED_PREPASS_FORMAT,
+            usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
+            view_formats: &[],
+        });
+        let previous_gbuffer_view = previous_gbuffer.create_view(&TextureViewDescriptor::default());
+
+        let previous_depth = render_device.create_texture(&TextureDescriptor {
+            label: Some("solari_lighting_previous_depth"),
+            size: view_size.to_extents(),
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: TextureDimension::D2,
+            format: CORE_3D_DEPTH_FORMAT,
+            usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
+            view_formats: &[],
+        });
+        let previous_depth_view = previous_depth.create_view(&TextureViewDescriptor::default());
+
         commands.entity(entity).insert(SolariLightingResources {
             reservoirs_a,
             reservoirs_b,
+            previous_gbuffer: (previous_gbuffer, previous_gbuffer_view),
+            previous_depth: (previous_depth, previous_depth_view),
             view_size,
         });
     }

crates/bevy_solari/src/realtime/reservoir.wgsl

@@ -2,7 +2,9 @@
 
 #define_import_path bevy_solari::reservoir
 
-#import bevy_solari::sampling::LightSample
+#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;
 
@@ -12,7 +14,7 @@ struct Reservoir {
     weight_sum: f32,
     confidence_weight: f32,
     unbiased_contribution_weight: f32,
-    _padding: f32,
+    visibility: f32,
 }
 
 fn empty_reservoir() -> Reservoir {
@@ -28,3 +30,59 @@ fn empty_reservoir() -> Reservoir {
 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);
+}