1-bounce GI (no ReSTIR)

crates/bevy_solari/src/realtime/mod.rs

@@ -28,6 +28,7 @@ pub struct SolariLightingPlugin;
 impl Plugin for SolariLightingPlugin {
     fn build(&self, app: &mut App) {
         embedded_asset!(app, "restir_di.wgsl");
+        embedded_asset!(app, "restir_gi.wgsl");
 
         app.register_type::<SolariLighting>()
             .insert_resource(DefaultOpaqueRendererMethod::deferred());

crates/bevy_solari/src/realtime/node.rs

@@ -35,8 +35,9 @@ pub mod graph {
 
 pub struct SolariLightingNode {
     bind_group_layout: BindGroupLayout,
-    initial_and_temporal_pipeline: CachedComputePipelineId,
+    di_initial_and_temporal_pipeline: CachedComputePipelineId,
-    spatial_and_shade_pipeline: CachedComputePipelineId,
+    di_spatial_and_shade_pipeline: CachedComputePipelineId,
+    gi_initial_and_temporal_pipeline: CachedComputePipelineId,
 }
 
 impl ViewNode for SolariLightingNode {
@@ -71,8 +72,9 @@ impl ViewNode for SolariLightingNode {
         let previous_view_uniforms = world.resource::<PreviousViewUniforms>();
         let frame_count = world.resource::<FrameCount>();
         let (
-            Some(initial_and_temporal_pipeline),
+            Some(di_initial_and_temporal_pipeline),
-            Some(spatial_and_shade_pipeline),
+            Some(di_spatial_and_shade_pipeline),
+            Some(gi_initial_and_temporal_pipeline),
             Some(scene_bindings),
             Some(viewport),
             Some(gbuffer),
@@ -81,8 +83,9 @@ impl ViewNode for SolariLightingNode {
             Some(view_uniforms),
             Some(previous_view_uniforms),
         ) = (
-            pipeline_cache.get_compute_pipeline(self.initial_and_temporal_pipeline),
+            pipeline_cache.get_compute_pipeline(self.di_initial_and_temporal_pipeline),
-            pipeline_cache.get_compute_pipeline(self.spatial_and_shade_pipeline),
+            pipeline_cache.get_compute_pipeline(self.di_spatial_and_shade_pipeline),
+            pipeline_cache.get_compute_pipeline(self.gi_initial_and_temporal_pipeline),
             &scene_bindings.bind_group,
             camera.physical_viewport_size,
             view_prepass_textures.deferred_view(),
@@ -100,8 +103,12 @@ impl ViewNode for SolariLightingNode {
             &self.bind_group_layout,
             &BindGroupEntries::sequential((
                 view_target.get_unsampled_color_attachment().view,
-                solari_lighting_resources.reservoirs_a.as_entire_binding(),
+                solari_lighting_resources
-                solari_lighting_resources.reservoirs_b.as_entire_binding(),
+                    .di_reservoirs_a
+                    .as_entire_binding(),
+                solari_lighting_resources
+                    .di_reservoirs_b
+                    .as_entire_binding(),
                 gbuffer,
                 depth_buffer,
                 motion_vectors,
@@ -134,14 +141,17 @@ impl ViewNode for SolariLightingNode {
             ],
         );
 
-        pass.set_pipeline(initial_and_temporal_pipeline);
+        pass.set_pipeline(di_initial_and_temporal_pipeline);
         pass.set_push_constants(
             0,
             bytemuck::cast_slice(&[frame_index, solari_lighting.reset as u32]),
         );
         pass.dispatch_workgroups(viewport.x.div_ceil(8), viewport.y.div_ceil(8), 1);
 
-        pass.set_pipeline(spatial_and_shade_pipeline);
+        pass.set_pipeline(di_spatial_and_shade_pipeline);
+        pass.dispatch_workgroups(viewport.x.div_ceil(8), viewport.y.div_ceil(8), 1);
+
+        pass.set_pipeline(gi_initial_and_temporal_pipeline);
         pass.dispatch_workgroups(viewport.x.div_ceil(8), viewport.y.div_ceil(8), 1);
 
         pass_span.end(&mut pass);
@@ -188,7 +198,7 @@ impl FromWorld for SolariLightingNode {
                 (
                     texture_storage_2d(
                         ViewTarget::TEXTURE_FORMAT_HDR,
-                        StorageTextureAccess::WriteOnly,
+                        StorageTextureAccess::ReadWrite,
                     ),
                     storage_buffer_sized(false, None),
                     storage_buffer_sized(false, None),
@@ -203,9 +213,9 @@ impl FromWorld for SolariLightingNode {
             ),
         );
 
-        let initial_and_temporal_pipeline =
+        let di_initial_and_temporal_pipeline =
             pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
-                label: Some("solari_lighting_initial_and_temporal_pipeline".into()),
+                label: Some("solari_lighting_di_initial_and_temporal_pipeline".into()),
                 layout: vec![
                     scene_bindings.bind_group_layout.clone(),
                     bind_group_layout.clone(),
@@ -220,9 +230,9 @@ impl FromWorld for SolariLightingNode {
                 zero_initialize_workgroup_memory: false,
             });
 
-        let spatial_and_shade_pipeline =
+        let di_spatial_and_shade_pipeline =
             pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
-                label: Some("solari_lighting_spatial_and_shade_pipeline".into()),
+                label: Some("solari_lighting_di_spatial_and_shade_pipeline".into()),
                 layout: vec![
                     scene_bindings.bind_group_layout.clone(),
                     bind_group_layout.clone(),
@@ -237,10 +247,28 @@ impl FromWorld for SolariLightingNode {
                 zero_initialize_workgroup_memory: false,
             });
 
+        let gi_initial_and_temporal_pipeline =
+            pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
+                label: Some("solari_lighting_gi_initial_and_temporal_pipeline".into()),
+                layout: vec![
+                    scene_bindings.bind_group_layout.clone(),
+                    bind_group_layout.clone(),
+                ],
+                push_constant_ranges: vec![PushConstantRange {
+                    stages: ShaderStages::COMPUTE,
+                    range: 0..8,
+                }],
+                shader: load_embedded_asset!(world, "restir_gi.wgsl"),
+                shader_defs: vec![],
+                entry_point: "initial_and_temporal".into(),
+                zero_initialize_workgroup_memory: false,
+            });
+
         Self {
             bind_group_layout,
-            initial_and_temporal_pipeline,
+            di_initial_and_temporal_pipeline,
-            spatial_and_shade_pipeline,
+            di_spatial_and_shade_pipeline,
+            gi_initial_and_temporal_pipeline,
         }
     }
 }

crates/bevy_solari/src/realtime/prepare.rs

@@ -18,13 +18,13 @@ use bevy_render::{
 };
 
 /// Size of a Reservoir shader struct in bytes.
-const RESERVOIR_STRUCT_SIZE: u64 = 32;
+const DI_RESERVOIR_STRUCT_SIZE: u64 = 32;
 
 /// Internal rendering resources used for Solari lighting.
 #[derive(Component)]
 pub struct SolariLightingResources {
-    pub reservoirs_a: Buffer,
+    pub di_reservoirs_a: Buffer,
-    pub reservoirs_b: Buffer,
+    pub di_reservoirs_b: Buffer,
     pub previous_gbuffer: (Texture, TextureView),
     pub previous_depth: (Texture, TextureView),
     pub view_size: UVec2,
@@ -47,17 +47,17 @@ pub fn prepare_solari_lighting_resources(
             continue;
         }
 
-        let size = (view_size.x * view_size.y) as u64 * RESERVOIR_STRUCT_SIZE;
+        let size = (view_size.x * view_size.y) as u64 * DI_RESERVOIR_STRUCT_SIZE;
 
-        let reservoirs_a = render_device.create_buffer(&BufferDescriptor {
+        let di_reservoirs_a = render_device.create_buffer(&BufferDescriptor {
-            label: Some("solari_lighting_reservoirs_a"),
+            label: Some("solari_lighting_di_reservoirs_a"),
             size,
             usage: BufferUsages::STORAGE,
             mapped_at_creation: false,
         });
 
-        let reservoirs_b = render_device.create_buffer(&BufferDescriptor {
+        let di_reservoirs_b = render_device.create_buffer(&BufferDescriptor {
-            label: Some("solari_lighting_reservoirs_b"),
+            label: Some("solari_lighting_di_reservoirs_b"),
             size,
             usage: BufferUsages::STORAGE,
             mapped_at_creation: false,
@@ -88,8 +88,8 @@ pub fn prepare_solari_lighting_resources(
         let previous_depth_view = previous_depth.create_view(&TextureViewDescriptor::default());
 
         commands.entity(entity).insert(SolariLightingResources {
-            reservoirs_a,
+            di_reservoirs_a,
-            reservoirs_b,
+            di_reservoirs_b,
             previous_gbuffer: (previous_gbuffer, previous_gbuffer_view),
             previous_depth: (previous_depth, previous_depth_view),
             view_size,

crates/bevy_solari/src/realtime/restir_di.wgsl

@@ -10,9 +10,9 @@
 #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>;
+@group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, read_write>;
-@group(1) @binding(1) var<storage, read_write> reservoirs_a: array<Reservoir>;
+@group(1) @binding(1) var<storage, read_write> di_reservoirs_a: array<Reservoir>;
-@group(1) @binding(2) var<storage, read_write> reservoirs_b: array<Reservoir>;
+@group(1) @binding(2) var<storage, read_write> di_reservoirs_b: array<Reservoir>;
 @group(1) @binding(3) var gbuffer: texture_2d<u32>;
 @group(1) @binding(4) var depth_buffer: texture_depth_2d;
 @group(1) @binding(5) var motion_vectors: texture_2d<f32>;
@@ -38,7 +38,7 @@ fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
     let depth = textureLoad(depth_buffer, global_id.xy, 0);
     if depth == 0.0 {
-        reservoirs_b[pixel_index] = empty_reservoir();
+        di_reservoirs_b[pixel_index] = empty_reservoir();
         return;
     }
     let gpixel = textureLoad(gbuffer, global_id.xy, 0);
@@ -51,7 +51,7 @@ fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3<u32>) {
     let temporal_reservoir = load_temporal_reservoir(global_id.xy, depth, world_position, world_normal);
     let combined_reservoir = merge_reservoirs(initial_reservoir, temporal_reservoir, world_position, world_normal, diffuse_brdf, &rng);
 
-    reservoirs_b[pixel_index] = combined_reservoir.merged_reservoir;
+    di_reservoirs_b[pixel_index] = combined_reservoir.merged_reservoir;
 }
 
 @compute @workgroup_size(8, 8, 1)
@@ -63,7 +63,7 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
     let depth = textureLoad(depth_buffer, global_id.xy, 0);
     if depth == 0.0 {
-        reservoirs_a[pixel_index] = empty_reservoir();
+        di_reservoirs_a[pixel_index] = empty_reservoir();
         textureStore(view_output, global_id.xy, vec4(vec3(0.0), 1.0));
         return;
     }
@@ -74,12 +74,12 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
     let diffuse_brdf = base_color / PI;
     let emissive = rgb9e5_to_vec3_(gpixel.g);
 
-    let input_reservoir = reservoirs_b[pixel_index];
+    let input_reservoir = di_reservoirs_b[pixel_index];
     let spatial_reservoir = load_spatial_reservoir(global_id.xy, depth, world_position, world_normal, &rng);
     let merge_result = merge_reservoirs(input_reservoir, spatial_reservoir, world_position, world_normal, diffuse_brdf, &rng);
     let combined_reservoir = merge_result.merged_reservoir;
 
-    reservoirs_a[pixel_index] = combined_reservoir;
+    di_reservoirs_a[pixel_index] = combined_reservoir;
 
     var pixel_color = merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight * combined_reservoir.visibility;
     pixel_color *= view.exposure;
@@ -136,7 +136,7 @@ fn load_temporal_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3
     }
 
     let temporal_pixel_index = temporal_pixel_id.x + temporal_pixel_id.y * u32(view.viewport.z);
-    var temporal_reservoir = reservoirs_a[temporal_pixel_index];
+    var temporal_reservoir = di_reservoirs_a[temporal_pixel_index];
 
     temporal_reservoir.sample.light_id.x = previous_frame_light_id_translations[temporal_reservoir.sample.light_id.x];
     if temporal_reservoir.sample.light_id.x == LIGHT_NOT_PRESENT_THIS_FRAME {
@@ -160,7 +160,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.viewport.z);
-    var spatial_reservoir = reservoirs_b[spatial_pixel_index];
+    var spatial_reservoir = di_reservoirs_b[spatial_pixel_index];
 
     if reservoir_valid(spatial_reservoir) {
         spatial_reservoir.visibility = trace_light_visibility(spatial_reservoir.sample, world_position);
@@ -209,7 +209,7 @@ fn depth_ndc_to_view_z(ndc_depth: f32) -> f32 {
 #endif
 }
 
-// Don't adjust the size of this struct without also adjusting RESERVOIR_STRUCT_SIZE.
+// Don't adjust the size of this struct without also adjusting DI_RESERVOIR_STRUCT_SIZE.
 struct Reservoir {
     sample: LightSample,
     weight_sum: f32,

crates/bevy_solari/src/realtime/restir_gi.wgsl

@@ -0,0 +1,63 @@
+// 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, PI_2}
+#import bevy_render::view::View
+#import bevy_solari::sampling::{sample_uniform_hemisphere, sample_random_light}
+#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
+
+@group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, read_write>;
+// @group(1) @binding(1) var<storage, read_write> di_reservoirs_a: array<Reservoir>;
+// @group(1) @binding(2) var<storage, read_write> di_reservoirs_b: array<Reservoir>;
+@group(1) @binding(3) var gbuffer: texture_2d<u32>;
+@group(1) @binding(4) var depth_buffer: texture_depth_2d;
+@group(1) @binding(5) var motion_vectors: texture_2d<f32>;
+@group(1) @binding(6) var previous_gbuffer: texture_2d<u32>;
+@group(1) @binding(7) var previous_depth_buffer: texture_depth_2d;
+@group(1) @binding(8) var<uniform> view: View;
+@group(1) @binding(9) var<uniform> previous_view: PreviousViewUniforms;
+struct PushConstants { frame_index: u32, reset: u32 }
+var<push_constant> constants: PushConstants;
+
+@compute @workgroup_size(8, 8, 1)
+fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    if any(global_id.xy >= vec2u(view.viewport.zw)) { return; }
+
+    let pixel_index = global_id.x + global_id.y * u32(view.viewport.z);
+    var rng = pixel_index + constants.frame_index;
+
+    let depth = textureLoad(depth_buffer, global_id.xy, 0);
+    if depth == 0.0 { return; }
+    let gpixel = textureLoad(gbuffer, global_id.xy, 0);
+    let world_position = reconstruct_world_position(global_id.xy, depth);
+    let world_normal = octahedral_decode(unpack_24bit_normal(gpixel.a));
+    let base_color = pow(unpack4x8unorm(gpixel.r).rgb, vec3(2.2));
+    let diffuse_brdf = base_color / PI;
+
+    let ray_direction = sample_uniform_hemisphere(world_normal, &rng);
+    let ray_hit = trace_ray(world_position, ray_direction, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE);
+    if ray_hit.kind == RAY_QUERY_INTERSECTION_NONE { return; }
+    let sample_point = resolve_ray_hit_full(ray_hit);
+    if all(sample_point.material.emissive != vec3(0.0)) { return; }
+    let sample_point_diffuse_brdf = sample_point.material.base_color / PI;
+    let radiance = sample_random_light(sample_point.world_position, sample_point.world_normal, &rng);
+
+    let cos_theta = dot(ray_direction, world_normal);
+    let inverse_uniform_hemisphere_pdf = PI_2;
+    let contribution = (radiance * sample_point_diffuse_brdf * diffuse_brdf * cos_theta * inverse_uniform_hemisphere_pdf);
+
+    var pixel_color = textureLoad(view_output, global_id.xy);
+    pixel_color += vec4(contribution * view.exposure, 0.0);
+    textureStore(view_output, global_id.xy, pixel_color);
+}
+
+fn reconstruct_world_position(pixel_id: vec2<u32>, depth: f32) -> vec3<f32> {
+    let uv = (vec2<f32>(pixel_id) + 0.5) / view.viewport.zw;
+    let xy_ndc = (uv - vec2(0.5)) * vec2(2.0, -2.0);
+    let world_pos = view.world_from_clip * vec4(xy_ndc, depth, 1.0);
+    return world_pos.xyz / world_pos.w;
+}

crates/bevy_solari/src/scene/sampling.wgsl

@@ -15,6 +15,17 @@ fn sample_cosine_hemisphere(normal: vec3<f32>, rng: ptr<function, u32>) -> vec3<
     return vec3(x, y, z);
 }
 
+// https://www.pbr-book.org/3ed-2018/Monte_Carlo_Integration/2D_Sampling_with_Multidimensional_Transformations#UniformlySamplingaHemisphere
+fn sample_uniform_hemisphere(normal: vec3<f32>, rng: ptr<function, u32>) -> vec3<f32> {
+    let cos_theta = rand_f(rng);
+    let phi = PI_2 * rand_f(rng);
+    let sin_theta = sqrt(max(1.0 - cos_theta * cos_theta, 0.0));
+    let x = sin_theta * cos(phi);
+    let y = sin_theta * sin(phi);
+    let z = cos_theta;
+    return build_orthonormal_basis(normal) * vec3(x, y, z);
+}
+
 // https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec19%3A294
 fn sample_disk(disk_radius: f32, rng: ptr<function, u32>) -> vec2<f32> {
     let ab = 2.0 * rand_vec2f(rng) - 1.0;