diff --git a/crates/bevy_solari/src/pathtracer/pathtracer.wgsl b/crates/bevy_solari/src/pathtracer/pathtracer.wgsl index c67b53e58e..be92b67bfa 100644 --- a/crates/bevy_solari/src/pathtracer/pathtracer.wgsl +++ b/crates/bevy_solari/src/pathtracer/pathtracer.wgsl @@ -47,7 +47,8 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3) { if ray_t_min == 0.0 { radiance = ray_hit.material.emissive; } // Sample direct lighting - radiance += throughput * diffuse_brdf * sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng); + let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng); + radiance += throughput * diffuse_brdf * direct_lighting.radiance * direct_lighting.inverse_pdf; // Sample new ray direction from the material BRDF for next bounce ray_direction = sample_cosine_hemisphere(ray_hit.world_normal, &rng); diff --git a/crates/bevy_solari/src/realtime/node.rs b/crates/bevy_solari/src/realtime/node.rs index 2561a897e6..056993d3c4 100644 --- a/crates/bevy_solari/src/realtime/node.rs +++ b/crates/bevy_solari/src/realtime/node.rs @@ -38,6 +38,7 @@ pub struct SolariLightingNode { di_initial_and_temporal_pipeline: CachedComputePipelineId, di_spatial_and_shade_pipeline: CachedComputePipelineId, gi_initial_and_temporal_pipeline: CachedComputePipelineId, + gi_spatial_and_shade_pipeline: CachedComputePipelineId, } impl ViewNode for SolariLightingNode { @@ -75,6 +76,7 @@ impl ViewNode for SolariLightingNode { Some(di_initial_and_temporal_pipeline), Some(di_spatial_and_shade_pipeline), Some(gi_initial_and_temporal_pipeline), + Some(gi_spatial_and_shade_pipeline), Some(scene_bindings), Some(viewport), Some(gbuffer), @@ -86,6 +88,7 @@ impl ViewNode for SolariLightingNode { 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), + pipeline_cache.get_compute_pipeline(self.gi_spatial_and_shade_pipeline), &scene_bindings.bind_group, camera.physical_viewport_size, view_prepass_textures.deferred_view(), @@ -109,6 +112,12 @@ impl ViewNode for SolariLightingNode { solari_lighting_resources .di_reservoirs_b .as_entire_binding(), + solari_lighting_resources + .gi_reservoirs_a + .as_entire_binding(), + solari_lighting_resources + .gi_reservoirs_b + .as_entire_binding(), gbuffer, depth_buffer, motion_vectors, @@ -154,6 +163,9 @@ impl ViewNode for SolariLightingNode { pass.set_pipeline(gi_initial_and_temporal_pipeline); pass.dispatch_workgroups(viewport.x.div_ceil(8), viewport.y.div_ceil(8), 1); + pass.set_pipeline(gi_spatial_and_shade_pipeline); + pass.dispatch_workgroups(viewport.x.div_ceil(8), viewport.y.div_ceil(8), 1); + pass_span.end(&mut pass); drop(pass); @@ -202,6 +214,8 @@ impl FromWorld for SolariLightingNode { ), storage_buffer_sized(false, None), storage_buffer_sized(false, None), + storage_buffer_sized(false, None), + storage_buffer_sized(false, None), texture_2d(TextureSampleType::Uint), texture_depth_2d(), texture_2d(TextureSampleType::Float { filterable: true }), @@ -264,11 +278,29 @@ impl FromWorld for SolariLightingNode { zero_initialize_workgroup_memory: false, }); + let gi_spatial_and_shade_pipeline = + pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("solari_lighting_gi_spatial_and_shade_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: "spatial_and_shade".into(), + zero_initialize_workgroup_memory: false, + }); + Self { bind_group_layout, di_initial_and_temporal_pipeline, di_spatial_and_shade_pipeline, gi_initial_and_temporal_pipeline, + gi_spatial_and_shade_pipeline, } } } diff --git a/crates/bevy_solari/src/realtime/prepare.rs b/crates/bevy_solari/src/realtime/prepare.rs index 3c454727bd..46a94a3ca2 100644 --- a/crates/bevy_solari/src/realtime/prepare.rs +++ b/crates/bevy_solari/src/realtime/prepare.rs @@ -17,14 +17,19 @@ use bevy_render::{ renderer::RenderDevice, }; -/// Size of a Reservoir shader struct in bytes. +/// Size of a DI Reservoir shader struct in bytes. const DI_RESERVOIR_STRUCT_SIZE: u64 = 32; +/// Size of a GI Reservoir shader struct in bytes. +const GI_RESERVOIR_STRUCT_SIZE: u64 = 48; + /// Internal rendering resources used for Solari lighting. #[derive(Component)] pub struct SolariLightingResources { pub di_reservoirs_a: Buffer, pub di_reservoirs_b: Buffer, + pub gi_reservoirs_a: Buffer, + pub gi_reservoirs_b: Buffer, pub previous_gbuffer: (Texture, TextureView), pub previous_depth: (Texture, TextureView), pub view_size: UVec2, @@ -47,18 +52,30 @@ pub fn prepare_solari_lighting_resources( continue; } - let size = (view_size.x * view_size.y) as u64 * DI_RESERVOIR_STRUCT_SIZE; - let di_reservoirs_a = render_device.create_buffer(&BufferDescriptor { label: Some("solari_lighting_di_reservoirs_a"), - size, + size: (view_size.x * view_size.y) as u64 * DI_RESERVOIR_STRUCT_SIZE, usage: BufferUsages::STORAGE, mapped_at_creation: false, }); let di_reservoirs_b = render_device.create_buffer(&BufferDescriptor { label: Some("solari_lighting_di_reservoirs_b"), - size, + size: (view_size.x * view_size.y) as u64 * DI_RESERVOIR_STRUCT_SIZE, + usage: BufferUsages::STORAGE, + mapped_at_creation: false, + }); + + let gi_reservoirs_a = render_device.create_buffer(&BufferDescriptor { + label: Some("solari_lighting_gi_reservoirs_a"), + size: (view_size.x * view_size.y) as u64 * GI_RESERVOIR_STRUCT_SIZE, + usage: BufferUsages::STORAGE, + mapped_at_creation: false, + }); + + let gi_reservoirs_b = render_device.create_buffer(&BufferDescriptor { + label: Some("solari_lighting_gi_reservoirs_b"), + size: (view_size.x * view_size.y) as u64 * GI_RESERVOIR_STRUCT_SIZE, usage: BufferUsages::STORAGE, mapped_at_creation: false, }); @@ -90,6 +107,8 @@ pub fn prepare_solari_lighting_resources( commands.entity(entity).insert(SolariLightingResources { di_reservoirs_a, di_reservoirs_b, + gi_reservoirs_a, + gi_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 c9efc6fb20..6dae89d21b 100644 --- a/crates/bevy_solari/src/realtime/restir_di.wgsl +++ b/crates/bevy_solari/src/realtime/restir_di.wgsl @@ -13,13 +13,13 @@ @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; +@group(1) @binding(5) var gbuffer: texture_2d; +@group(1) @binding(6) var depth_buffer: texture_depth_2d; +@group(1) @binding(7) var motion_vectors: texture_2d; +@group(1) @binding(8) var previous_gbuffer: texture_2d; +@group(1) @binding(9) var previous_depth_buffer: texture_depth_2d; +@group(1) @binding(10) var view: View; +@group(1) @binding(11) var previous_view: PreviousViewUniforms; struct PushConstants { frame_index: u32, reset: u32 } var constants: PushConstants; diff --git a/crates/bevy_solari/src/realtime/restir_gi.wgsl b/crates/bevy_solari/src/realtime/restir_gi.wgsl index 7ef919d845..8e090ca29c 100644 --- a/crates/bevy_solari/src/realtime/restir_gi.wgsl +++ b/crates/bevy_solari/src/realtime/restir_gi.wgsl @@ -11,18 +11,20 @@ #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; +@group(1) @binding(3) var gi_reservoirs_a: array; +@group(1) @binding(4) var gi_reservoirs_b: array; +@group(1) @binding(5) var gbuffer: texture_2d; +@group(1) @binding(6) var depth_buffer: texture_depth_2d; +@group(1) @binding(7) var motion_vectors: texture_2d; +@group(1) @binding(8) var previous_gbuffer: texture_2d; +@group(1) @binding(9) var previous_depth_buffer: texture_depth_2d; +@group(1) @binding(10) var view: View; +@group(1) @binding(11) var previous_view: PreviousViewUniforms; struct PushConstants { frame_index: u32, reset: u32 } var constants: PushConstants; +const CONFIDENCE_WEIGHT_CAP = 30.0; + @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; } @@ -31,33 +33,178 @@ fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3) { var rng = pixel_index + constants.frame_index; let depth = textureLoad(depth_buffer, global_id.xy, 0); - if depth == 0.0 { return; } + if depth == 0.0 { + gi_reservoirs_b[pixel_index] = empty_reservoir(); + 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 temporal_reservoir = load_temporal_reservoir(global_id.xy, depth, world_position, world_normal); + 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; } + if ray_hit.kind == RAY_QUERY_INTERSECTION_NONE { + gi_reservoirs_b[pixel_index] = temporal_reservoir; + return; + } let sample_point = resolve_ray_hit_full(ray_hit); - if all(sample_point.material.emissive != vec3(0.0)) { return; } + if all(sample_point.material.emissive != vec3(0.0)) { + gi_reservoirs_b[pixel_index] = temporal_reservoir; + 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 direct_lighting = sample_random_light(sample_point.world_position, sample_point.world_normal, &rng); + let sample_point_radiance = direct_lighting.radiance * sample_point_diffuse_brdf; 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 combined_reservoir = empty_reservoir(); + combined_reservoir.confidence_weight = 1.0 + temporal_reservoir.confidence_weight; + + let mis_weight_denominator = 1.0 / combined_reservoir.confidence_weight; + + let new_mis_weight = mis_weight_denominator; + let new_target_function = luminance(sample_point_radiance * diffuse_brdf * cos_theta); + let new_inverse_pdf = direct_lighting.inverse_pdf * inverse_uniform_hemisphere_pdf; + let new_resampling_weight = new_mis_weight * (new_target_function * new_inverse_pdf); + + let temporal_mis_weight = temporal_reservoir.confidence_weight * mis_weight_denominator; + let temporal_cos_theta = dot(normalize(temporal_reservoir.sample_point_world_position - world_position), world_normal); + let temporal_target_function = luminance(temporal_reservoir.radiance * diffuse_brdf * temporal_cos_theta); + let temporal_resampling_weight = temporal_mis_weight * (temporal_target_function * temporal_reservoir.unbiased_contribution_weight); + + combined_reservoir.weight_sum = new_resampling_weight + temporal_resampling_weight; + + if rand_f(&rng) < temporal_resampling_weight / combined_reservoir.weight_sum { + combined_reservoir.sample_point_world_position = temporal_reservoir.sample_point_world_position; + combined_reservoir.radiance = temporal_reservoir.radiance; + + let inverse_target_function = select(0.0, 1.0 / temporal_target_function, temporal_target_function > 0.0); + combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function; + } else { + combined_reservoir.sample_point_world_position = sample_point.world_position; + combined_reservoir.radiance = sample_point_radiance; + + let inverse_target_function = select(0.0, 1.0 / new_target_function, new_target_function > 0.0); + combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function; + } + + gi_reservoirs_b[pixel_index] = combined_reservoir; +} + + +@compute @workgroup_size(8, 8, 1) +fn spatial_and_shade(@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 { + gi_reservoirs_a[pixel_index] = empty_reservoir(); + 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 input_reservoir = gi_reservoirs_b[pixel_index]; + let cos_theta = dot(normalize(input_reservoir.sample_point_world_position - world_position), world_normal); + let radiance = input_reservoir.radiance * diffuse_brdf * cos_theta; + + gi_reservoirs_a[pixel_index] = input_reservoir; var pixel_color = textureLoad(view_output, global_id.xy); - pixel_color += vec4(contribution * view.exposure, 0.0); + pixel_color += vec4(radiance * input_reservoir.unbiased_contribution_weight * view.exposure, 0.0); textureStore(view_output, global_id.xy, pixel_color); } +fn load_temporal_reservoir(pixel_id: vec2, depth: f32, world_position: vec3, world_normal: vec3) -> Reservoir { + let motion_vector = textureLoad(motion_vectors, pixel_id, 0).xy; + let temporal_pixel_id_float = round(vec2(pixel_id) - (motion_vector * view.viewport.zw)); + let temporal_pixel_id = vec2(temporal_pixel_id_float); + if any(temporal_pixel_id_float < vec2(0.0)) || any(temporal_pixel_id_float >= view.viewport.zw) || bool(constants.reset) { + return empty_reservoir(); + } + + let temporal_depth = textureLoad(previous_depth_buffer, temporal_pixel_id, 0); + let temporal_gpixel = textureLoad(previous_gbuffer, temporal_pixel_id, 0); + let temporal_world_position = reconstruct_previous_world_position(temporal_pixel_id, temporal_depth); + let temporal_world_normal = octahedral_decode(unpack_24bit_normal(temporal_gpixel.a)); + if pixel_dissimilar(depth, world_position, temporal_world_position, world_normal, temporal_world_normal) { + return empty_reservoir(); + } + + let temporal_pixel_index = temporal_pixel_id.x + temporal_pixel_id.y * u32(view.viewport.z); + var temporal_reservoir = gi_reservoirs_a[temporal_pixel_index]; + + temporal_reservoir.confidence_weight = min(temporal_reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP); + + return temporal_reservoir; +} + 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; } + +fn reconstruct_previous_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 = previous_view.world_from_clip * vec4(xy_ndc, depth, 1.0); + return world_pos.xyz / world_pos.w; +} + +// Reject if tangent plane difference difference more than 0.3% or angle between normals more than 25 degrees +fn pixel_dissimilar(depth: f32, world_position: vec3, other_world_position: vec3, normal: vec3, other_normal: vec3) -> bool { + // https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s22699-fast-denoising-with-self-stabilizing-recurrent-blurs.pdf#page=45 + let tangent_plane_distance = abs(dot(normal, other_world_position - world_position)); + let view_z = -depth_ndc_to_view_z(depth); + + return tangent_plane_distance / view_z > 0.003 || dot(normal, other_normal) < 0.906; +} + +fn depth_ndc_to_view_z(ndc_depth: f32) -> f32 { +#ifdef VIEW_PROJECTION_PERSPECTIVE + return -view.clip_from_view[3][2]() / ndc_depth; +#else ifdef VIEW_PROJECTION_ORTHOGRAPHIC + return -(view.clip_from_view[3][2] - ndc_depth) / view.clip_from_view[2][2]; +#else + let view_pos = view.view_from_clip * vec4(0.0, 0.0, ndc_depth, 1.0); + return view_pos.z / view_pos.w; +#endif +} + +struct Reservoir { + sample_point_world_position: vec3, + weight_sum: f32, + radiance: vec3, + confidence_weight: f32, + unbiased_contribution_weight: f32, + padding1: f32, + padding2: f32, + padding3: f32, +} + +fn empty_reservoir() -> Reservoir { + return Reservoir( + vec3(0.0), + 0.0, + vec3(0.0), + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ); +} diff --git a/crates/bevy_solari/src/scene/sampling.wgsl b/crates/bevy_solari/src/scene/sampling.wgsl index c6ad92af49..a707fbb1b3 100644 --- a/crates/bevy_solari/src/scene/sampling.wgsl +++ b/crates/bevy_solari/src/scene/sampling.wgsl @@ -48,11 +48,16 @@ fn sample_disk(disk_radius: f32, rng: ptr) -> vec2 { return vec2(x, y); } -fn sample_random_light(ray_origin: vec3, origin_world_normal: vec3, rng: ptr) -> vec3 { +struct SampleRandomLightResult { + radiance: vec3, + inverse_pdf: f32, +} + +fn sample_random_light(ray_origin: vec3, origin_world_normal: vec3, rng: ptr) -> SampleRandomLightResult { let light_sample = generate_random_light_sample(rng); let light_contribution = calculate_light_contribution(light_sample, ray_origin, origin_world_normal); let visibility = trace_light_visibility(light_sample, ray_origin); - return light_contribution.radiance * visibility * light_contribution.inverse_pdf; + return SampleRandomLightResult(light_contribution.radiance * visibility, light_contribution.inverse_pdf); } struct LightSample {