From 4bdae6a860a1389b024934f21deab06ea2bbb83d Mon Sep 17 00:00:00 2001 From: JMS55 <47158642+JMS55@users.noreply.github.com> Date: Tue, 1 Jul 2025 22:38:04 -0400 Subject: [PATCH] 1-bounce GI (no ReSTIR) --- crates/bevy_solari/src/realtime/mod.rs | 1 + crates/bevy_solari/src/realtime/node.rs | 62 +++++++++++++----- crates/bevy_solari/src/realtime/prepare.rs | 20 +++--- .../bevy_solari/src/realtime/restir_di.wgsl | 22 +++---- .../bevy_solari/src/realtime/restir_gi.wgsl | 63 +++++++++++++++++++ crates/bevy_solari/src/scene/sampling.wgsl | 11 ++++ 6 files changed, 141 insertions(+), 38 deletions(-) create mode 100644 crates/bevy_solari/src/realtime/restir_gi.wgsl diff --git a/crates/bevy_solari/src/realtime/mod.rs b/crates/bevy_solari/src/realtime/mod.rs index b6a8f27a31..96bedd0b97 100644 --- a/crates/bevy_solari/src/realtime/mod.rs +++ b/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::() .insert_resource(DefaultOpaqueRendererMethod::deferred()); diff --git a/crates/bevy_solari/src/realtime/node.rs b/crates/bevy_solari/src/realtime/node.rs index 8fb97f84b2..2561a897e6 100644 --- a/crates/bevy_solari/src/realtime/node.rs +++ b/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, - spatial_and_shade_pipeline: CachedComputePipelineId, + di_initial_and_temporal_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::(); let frame_count = world.resource::(); let ( - Some(initial_and_temporal_pipeline), - Some(spatial_and_shade_pipeline), + Some(di_initial_and_temporal_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.spatial_and_shade_pipeline), + pipeline_cache.get_compute_pipeline(self.di_initial_and_temporal_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.reservoirs_b.as_entire_binding(), + solari_lighting_resources + .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, - spatial_and_shade_pipeline, + di_initial_and_temporal_pipeline, + di_spatial_and_shade_pipeline, + gi_initial_and_temporal_pipeline, } } } diff --git a/crates/bevy_solari/src/realtime/prepare.rs b/crates/bevy_solari/src/realtime/prepare.rs index 992a75c451..3c454727bd 100644 --- a/crates/bevy_solari/src/realtime/prepare.rs +++ b/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 reservoirs_b: Buffer, + pub di_reservoirs_a: 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 { - label: Some("solari_lighting_reservoirs_a"), + let di_reservoirs_a = render_device.create_buffer(&BufferDescriptor { + label: Some("solari_lighting_di_reservoirs_a"), size, usage: BufferUsages::STORAGE, mapped_at_creation: false, }); - let reservoirs_b = render_device.create_buffer(&BufferDescriptor { - label: Some("solari_lighting_reservoirs_b"), + let di_reservoirs_b = render_device.create_buffer(&BufferDescriptor { + 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, - reservoirs_b, + di_reservoirs_a, + di_reservoirs_b, previous_gbuffer: (previous_gbuffer, previous_gbuffer_view), previous_depth: (previous_depth, previous_depth_view), view_size, diff --git a/crates/bevy_solari/src/realtime/restir_di.wgsl b/crates/bevy_solari/src/realtime/restir_di.wgsl index 70de4564cc..c9efc6fb20 100644 --- a/crates/bevy_solari/src/realtime/restir_di.wgsl +++ b/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; -@group(1) @binding(1) var reservoirs_a: array; -@group(1) @binding(2) var reservoirs_b: array; +@group(1) @binding(0) var view_output: texture_storage_2d; +@group(1) @binding(1) var di_reservoirs_a: array; +@group(1) @binding(2) var di_reservoirs_b: array; @group(1) @binding(3) var gbuffer: texture_2d; @group(1) @binding(4) var depth_buffer: texture_depth_2d; @group(1) @binding(5) var motion_vectors: texture_2d; @@ -38,7 +38,7 @@ fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3) { 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) { 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) { 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) { 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, 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, 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, diff --git a/crates/bevy_solari/src/realtime/restir_gi.wgsl b/crates/bevy_solari/src/realtime/restir_gi.wgsl new file mode 100644 index 0000000000..7ef919d845 --- /dev/null +++ b/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; +// @group(1) @binding(1) var di_reservoirs_a: array; +// @group(1) @binding(2) var di_reservoirs_b: array; +@group(1) @binding(3) var gbuffer: texture_2d; +@group(1) @binding(4) var depth_buffer: texture_depth_2d; +@group(1) @binding(5) var motion_vectors: texture_2d; +@group(1) @binding(6) var previous_gbuffer: texture_2d; +@group(1) @binding(7) var previous_depth_buffer: texture_depth_2d; +@group(1) @binding(8) var view: View; +@group(1) @binding(9) var previous_view: PreviousViewUniforms; +struct PushConstants { frame_index: u32, reset: u32 } +var constants: PushConstants; + +@compute @workgroup_size(8, 8, 1) +fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3) { + 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, depth: f32) -> vec3 { + let uv = (vec2(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; +} diff --git a/crates/bevy_solari/src/scene/sampling.wgsl b/crates/bevy_solari/src/scene/sampling.wgsl index be709f0bc8..c6ad92af49 100644 --- a/crates/bevy_solari/src/scene/sampling.wgsl +++ b/crates/bevy_solari/src/scene/sampling.wgsl @@ -15,6 +15,17 @@ fn sample_cosine_hemisphere(normal: vec3, rng: ptr) -> 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, rng: ptr) -> vec3 { + 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) -> vec2 { let ab = 2.0 * rand_vec2f(rng) - 1.0;