diff --git a/crates/bevy_solari/src/pathtracer/pathtracer.wgsl b/crates/bevy_solari/src/pathtracer/pathtracer.wgsl index da216da959201..68cb45e207ffb 100644 --- a/crates/bevy_solari/src/pathtracer/pathtracer.wgsl +++ b/crates/bevy_solari/src/pathtracer/pathtracer.wgsl @@ -54,7 +54,7 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3) { radiance += mis_weight * throughput * ray_hit.material.emissive; // Sample direct lighting, but only if the surface is not mirror-like - let is_perfectly_specular = ray_hit.material.roughness < 0.0001 && ray_hit.material.metallic > 0.9999; + let is_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999; if !is_perfectly_specular { let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng); let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit); @@ -100,7 +100,7 @@ struct NextBounce { } fn importance_sample_next_bounce(wo: vec3, ray_hit: ResolvedRayHitFull, rng: ptr) -> NextBounce { - let is_perfectly_specular = ray_hit.material.roughness < 0.0001 && ray_hit.material.metallic > 0.9999; + let is_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999; if is_perfectly_specular { return NextBounce(reflect(-wo, ray_hit.world_normal), 1.0, true); } diff --git a/crates/bevy_solari/src/realtime/gbuffer_utils.wgsl b/crates/bevy_solari/src/realtime/gbuffer_utils.wgsl new file mode 100644 index 0000000000000..4a6037a050657 --- /dev/null +++ b/crates/bevy_solari/src/realtime/gbuffer_utils.wgsl @@ -0,0 +1,57 @@ +#define_import_path bevy_solari::gbuffer_utils + +#import bevy_pbr::pbr_deferred_types::unpack_24bit_normal +#import bevy_pbr::rgb9e5::rgb9e5_to_vec3_ +#import bevy_pbr::utils::octahedral_decode +#import bevy_render::view::View +#import bevy_solari::scene_bindings::ResolvedMaterial + +struct ResolvedGPixel { + world_position: vec3, + world_normal: vec3, + material: ResolvedMaterial, +} + +fn gpixel_resolve(gpixel: vec4, depth: f32, pixel_id: vec2, view_size: vec2, world_from_clip: mat4x4) -> ResolvedGPixel { + let world_position = reconstruct_world_position(pixel_id, depth, view_size, world_from_clip); + let world_normal = octahedral_decode(unpack_24bit_normal(gpixel.a)); + + let base_rough = unpack4x8unorm(gpixel.r); + let base_color = pow(base_rough.rgb, vec3(2.2)); + let perceptual_roughness = base_rough.a; + let roughness = clamp(perceptual_roughness * perceptual_roughness, 0.001, 1.0); + let props = unpack4x8unorm(gpixel.b); + let reflectance = vec3(props.r); + let metallic = props.g; + let emissive = rgb9e5_to_vec3_(gpixel.g); + let material = ResolvedMaterial(base_color, emissive, reflectance, perceptual_roughness, roughness, metallic); + + return ResolvedGPixel(world_position, world_normal, material); +} + +fn reconstruct_world_position(pixel_id: vec2, depth: f32, view_size: vec2, world_from_clip: mat4x4) -> vec3 { + let uv = (vec2(pixel_id) + 0.5) / view_size; + let xy_ndc = (uv - vec2(0.5)) * vec2(2.0, -2.0); + let world_pos = 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, view: View) -> 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, view); + + return tangent_plane_distance / view_z > 0.003 || dot(normal, other_normal) < 0.906; +} + +fn depth_ndc_to_view_z(ndc_depth: f32, view: View) -> 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 +} diff --git a/crates/bevy_solari/src/realtime/mod.rs b/crates/bevy_solari/src/realtime/mod.rs index 04e3528e63004..3d39f109a6dca 100644 --- a/crates/bevy_solari/src/realtime/mod.rs +++ b/crates/bevy_solari/src/realtime/mod.rs @@ -32,9 +32,11 @@ pub struct SolariLightingPlugin; impl Plugin for SolariLightingPlugin { fn build(&self, app: &mut App) { + load_shader_library!(app, "gbuffer_utils.wgsl"); load_shader_library!(app, "presample_light_tiles.wgsl"); embedded_asset!(app, "restir_di.wgsl"); embedded_asset!(app, "restir_gi.wgsl"); + embedded_asset!(app, "specular_trace.wgsl"); load_shader_library!(app, "world_cache_query.wgsl"); embedded_asset!(app, "world_cache_compact.wgsl"); embedded_asset!(app, "world_cache_update.wgsl"); diff --git a/crates/bevy_solari/src/realtime/node.rs b/crates/bevy_solari/src/realtime/node.rs index 9d3e8e5d0d9c9..9809b06a62ef2 100644 --- a/crates/bevy_solari/src/realtime/node.rs +++ b/crates/bevy_solari/src/realtime/node.rs @@ -55,6 +55,7 @@ pub struct SolariLightingNode { di_spatial_and_shade_pipeline: CachedComputePipelineId, gi_initial_and_temporal_pipeline: CachedComputePipelineId, gi_spatial_and_shade_pipeline: CachedComputePipelineId, + specular_trace_pipeline: CachedComputePipelineId, #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] resolve_dlss_rr_textures_pipeline: CachedComputePipelineId, } @@ -120,6 +121,7 @@ impl ViewNode for SolariLightingNode { Some(di_spatial_and_shade_pipeline), Some(gi_initial_and_temporal_pipeline), Some(gi_spatial_and_shade_pipeline), + Some(specular_trace_pipeline), Some(scene_bindings), Some(gbuffer), Some(depth_buffer), @@ -139,6 +141,7 @@ impl ViewNode for SolariLightingNode { 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), + pipeline_cache.get_compute_pipeline(self.specular_trace_pipeline), &scene_bindings.bind_group, view_prepass_textures.deferred_view(), view_prepass_textures.depth_view(), @@ -318,6 +321,13 @@ impl ViewNode for SolariLightingNode { ); pass.dispatch_workgroups(dx, dy, 1); + pass.set_pipeline(specular_trace_pipeline); + pass.set_push_constants( + 0, + bytemuck::cast_slice(&[frame_index, solari_lighting.reset as u32]), + ); + pass.dispatch_workgroups(dx, dy, 1); + pass_span.end(&mut pass); drop(pass); @@ -530,6 +540,13 @@ impl FromWorld for SolariLightingNode { None, vec![], ), + specular_trace_pipeline: create_pipeline( + "solari_lighting_specular_trace_pipeline", + "specular_trace", + load_embedded_asset!(world, "specular_trace.wgsl"), + None, + vec![], + ), #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] resolve_dlss_rr_textures_pipeline: create_pipeline( "solari_lighting_resolve_dlss_rr_textures_pipeline", diff --git a/crates/bevy_solari/src/realtime/restir_di.wgsl b/crates/bevy_solari/src/realtime/restir_di.wgsl index 4512214a01541..a1484d9e2349f 100644 --- a/crates/bevy_solari/src/realtime/restir_di.wgsl +++ b/crates/bevy_solari/src/realtime/restir_di.wgsl @@ -2,12 +2,12 @@ // https://d1qx31qr3h6wln.cloudfront.net/publications/ReSTIR%20GI.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, rand_range_u, octahedral_decode, sample_disk} +#import bevy_pbr::utils::{rand_f, rand_range_u, sample_disk} #import bevy_render::maths::PI #import bevy_render::view::View +#import bevy_solari::brdf::evaluate_brdf +#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar} #import bevy_solari::presample_light_tiles::{ResolvedLightSamplePacked, unpack_resolved_light_sample} #import bevy_solari::sampling::{LightSample, calculate_resolved_light_contribution, resolve_and_calculate_light_contribution, resolve_light_sample, trace_light_visibility} #import bevy_solari::scene_bindings::{light_sources, previous_frame_light_id_translations, LIGHT_NOT_PRESENT_THIS_FRAME} @@ -45,15 +45,12 @@ fn initial_and_temporal(@builtin(workgroup_id) workgroup_id: vec3, @builtin store_reservoir_b(global_id.xy, 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 surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip); - let initial_reservoir = generate_initial_reservoir(world_position, world_normal, diffuse_brdf, workgroup_id.xy, &rng); - let temporal_reservoir = load_temporal_reservoir(global_id.xy, depth, world_position, world_normal); - let merge_result = merge_reservoirs(initial_reservoir, temporal_reservoir, world_position, world_normal, diffuse_brdf, &rng); + let diffuse_brdf = surface.material.base_color / PI; + let initial_reservoir = generate_initial_reservoir(surface.world_position, surface.world_normal, diffuse_brdf, workgroup_id.xy, &rng); + let temporal_reservoir = load_temporal_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal); + let merge_result = merge_reservoirs(initial_reservoir, temporal_reservoir, surface.world_position, surface.world_normal, diffuse_brdf, &rng); store_reservoir_b(global_id.xy, merge_result.merged_reservoir); } @@ -70,24 +67,23 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3) { store_reservoir_a(global_id.xy, 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 emissive = rgb9e5_to_vec3_(gpixel.g); + let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip); + let diffuse_brdf = surface.material.base_color / PI; let input_reservoir = load_reservoir_b(global_id.xy); - 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 spatial_reservoir = load_spatial_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal, &rng); + let merge_result = merge_reservoirs(input_reservoir, spatial_reservoir, surface.world_position, surface.world_normal, diffuse_brdf, &rng); let combined_reservoir = merge_result.merged_reservoir; store_reservoir_a(global_id.xy, combined_reservoir); + let wo = normalize(view.world_position - surface.world_position); + let brdf = evaluate_brdf(surface.world_normal, wo, merge_result.wi, surface.material); + var pixel_color = merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight; pixel_color *= view.exposure; - pixel_color *= diffuse_brdf; - pixel_color += emissive; + pixel_color *= brdf; + pixel_color += surface.material.emissive; textureStore(view_output, global_id.xy, vec4(pixel_color, 1.0)); } @@ -147,10 +143,8 @@ fn load_temporal_reservoir(pixel_id: vec2, depth: f32, world_position: vec3 // Check if the pixel features have changed heavily between the current and previous frame 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) { + let temporal_surface = gpixel_resolve(textureLoad(previous_gbuffer, temporal_pixel_id, 0), temporal_depth, temporal_pixel_id, view.main_pass_viewport.zw, previous_view.world_from_clip); + if pixel_dissimilar(depth, world_position, temporal_surface.world_position, world_normal, temporal_surface.world_normal, view) { return empty_reservoir(); } @@ -174,10 +168,8 @@ fn load_spatial_reservoir(pixel_id: vec2, depth: f32, world_position: vec3< let spatial_pixel_id = get_neighbor_pixel_id(pixel_id, rng); let spatial_depth = textureLoad(depth_buffer, spatial_pixel_id, 0); - let spatial_gpixel = textureLoad(gbuffer, spatial_pixel_id, 0); - let spatial_world_position = reconstruct_world_position(spatial_pixel_id, spatial_depth); - let spatial_world_normal = octahedral_decode(unpack_24bit_normal(spatial_gpixel.a)); - if pixel_dissimilar(depth, world_position, spatial_world_position, world_normal, spatial_world_normal) { + let spatial_surface = gpixel_resolve(textureLoad(gbuffer, spatial_pixel_id, 0), spatial_depth, spatial_pixel_id, view.main_pass_viewport.zw, view.world_from_clip); + if pixel_dissimilar(depth, world_position, spatial_surface.world_position, world_normal, spatial_surface.world_normal, view) { return empty_reservoir(); } @@ -197,40 +189,6 @@ fn get_neighbor_pixel_id(center_pixel_id: vec2, rng: ptr) -> return vec2(spatial_id); } -fn reconstruct_world_position(pixel_id: vec2, depth: f32) -> vec3 { - let uv = (vec2(pixel_id) + 0.5) / view.main_pass_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.main_pass_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: LightSample, confidence_weight: f32, @@ -278,6 +236,7 @@ fn load_reservoir_b(pixel: vec2) -> Reservoir { struct ReservoirMergeResult { merged_reservoir: Reservoir, selected_sample_radiance: vec3, + wi: vec3, } fn merge_reservoirs( @@ -288,15 +247,16 @@ fn merge_reservoirs( diffuse_brdf: vec3, rng: ptr, ) -> ReservoirMergeResult { + let canonical_contribution = reservoir_contribution(canonical_reservoir, world_position, world_normal, diffuse_brdf); + let other_contribution = reservoir_contribution(other_reservoir, world_position, world_normal, diffuse_brdf); + 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 canonical_resampling_weight = canonical_mis_weight * (canonical_contribution.target_function * 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); + let other_resampling_weight = other_mis_weight * (other_contribution.target_function * other_reservoir.unbiased_contribution_weight); let weight_sum = canonical_resampling_weight + other_resampling_weight; @@ -306,24 +266,30 @@ fn merge_reservoirs( if rand_f(rng) < other_resampling_weight / 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); + let inverse_target_function = select(0.0, 1.0 / other_contribution.target_function, other_contribution.target_function > 0.0); combined_reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function; - return ReservoirMergeResult(combined_reservoir, other_target_function.rgb); + return ReservoirMergeResult(combined_reservoir, other_contribution.radiance, other_contribution.wi); } 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); + let inverse_target_function = select(0.0, 1.0 / canonical_contribution.target_function, canonical_contribution.target_function > 0.0); combined_reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function; - return ReservoirMergeResult(combined_reservoir, canonical_target_function.rgb); + return ReservoirMergeResult(combined_reservoir, canonical_contribution.radiance, canonical_contribution.wi); } } +struct ReservoirContribution { + radiance: vec3, + target_function: f32, + wi: vec3, +} + // TODO: Have input take ResolvedLightSample instead of reservoir.light_sample -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 = resolve_and_calculate_light_contribution(reservoir.sample, world_position, world_normal).radiance; - let target_function = luminance(light_contribution * diffuse_brdf); - return vec4(light_contribution, target_function); +fn reservoir_contribution(reservoir: Reservoir, world_position: vec3, world_normal: vec3, diffuse_brdf: vec3) -> ReservoirContribution { + if !reservoir_valid(reservoir) { return ReservoirContribution(vec3(0.0), 0.0, vec3(0.0)); } + let light_contribution = resolve_and_calculate_light_contribution(reservoir.sample, world_position, world_normal); + let target_function = luminance(light_contribution.radiance * diffuse_brdf); + return ReservoirContribution(light_contribution.radiance, target_function, light_contribution.wi); } diff --git a/crates/bevy_solari/src/realtime/restir_gi.wgsl b/crates/bevy_solari/src/realtime/restir_gi.wgsl index 66ececed93a16..1017993f08648 100644 --- a/crates/bevy_solari/src/realtime/restir_gi.wgsl +++ b/crates/bevy_solari/src/realtime/restir_gi.wgsl @@ -1,11 +1,12 @@ // 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::utils::{rand_f, sample_uniform_hemisphere, uniform_hemisphere_inverse_pdf, sample_disk, octahedral_decode} +#import bevy_pbr::utils::{rand_f, sample_uniform_hemisphere, uniform_hemisphere_inverse_pdf, sample_disk} #import bevy_render::maths::PI #import bevy_render::view::View +#import bevy_solari::brdf::evaluate_diffuse_brdf +#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar} #import bevy_solari::sampling::{sample_random_light, trace_point_visibility} #import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX} #import bevy_solari::world_cache::query_world_cache @@ -38,15 +39,11 @@ fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3) { 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 initial_reservoir = generate_initial_reservoir(world_position, world_normal, &rng); - let temporal = load_temporal_reservoir(global_id.xy, depth, world_position, world_normal); - let merge_result = merge_reservoirs(initial_reservoir, world_position, world_normal, diffuse_brdf, + let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip); + + let initial_reservoir = generate_initial_reservoir(surface.world_position, surface.world_normal, &rng); + let temporal = load_temporal_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal); + let merge_result = merge_reservoirs(initial_reservoir, surface.world_position, surface.world_normal, surface.material.base_color / PI, temporal.reservoir, temporal.world_position, temporal.world_normal, temporal.diffuse_brdf, &rng); gi_reservoirs_b[pixel_index] = merge_result.merged_reservoir; @@ -64,26 +61,24 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3) { 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 surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip); let input_reservoir = gi_reservoirs_b[pixel_index]; - let spatial = load_spatial_reservoir(global_id.xy, depth, world_position, world_normal, &rng); - let merge_result = merge_reservoirs(input_reservoir, world_position, world_normal, diffuse_brdf, + let spatial = load_spatial_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal, &rng); + let merge_result = merge_reservoirs(input_reservoir, surface.world_position, surface.world_normal, surface.material.base_color / PI, spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.diffuse_brdf, &rng); let combined_reservoir = merge_result.merged_reservoir; gi_reservoirs_a[pixel_index] = combined_reservoir; + let brdf = evaluate_diffuse_brdf(surface.material.base_color, surface.material.metallic); + var pixel_color = textureLoad(view_output, global_id.xy); - pixel_color += vec4(merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight * view.exposure, 0.0); + pixel_color += vec4(merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight * view.exposure * brdf, 0.0); textureStore(view_output, global_id.xy, pixel_color); #ifdef VISUALIZE_WORLD_CACHE - textureStore(view_output, global_id.xy, vec4(query_world_cache(world_position, world_normal, view.world_position) * view.exposure, 1.0)); + textureStore(view_output, global_id.xy, vec4(query_world_cache(surface.world_position, surface.world_normal, view.world_position) * view.exposure, 1.0)); #endif } @@ -138,12 +133,9 @@ fn load_temporal_reservoir(pixel_id: vec2, depth: f32, world_position: vec3 // Check if the pixel features have changed heavily between the current and previous frame 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)); - let temporal_base_color = pow(unpack4x8unorm(temporal_gpixel.r).rgb, vec3(2.2)); - let temporal_diffuse_brdf = temporal_base_color / PI; - if pixel_dissimilar(depth, world_position, temporal_world_position, world_normal, temporal_world_normal) { + let temporal_surface = gpixel_resolve(textureLoad(previous_gbuffer, temporal_pixel_id, 0), temporal_depth, temporal_pixel_id, view.main_pass_viewport.zw, previous_view.world_from_clip); + let temporal_diffuse_brdf = temporal_surface.material.base_color / PI; + if pixel_dissimilar(depth, world_position, temporal_surface.world_position, world_normal, temporal_surface.world_normal, view) { continue; } @@ -152,7 +144,7 @@ fn load_temporal_reservoir(pixel_id: vec2, depth: f32, world_position: vec3 temporal_reservoir.confidence_weight = min(temporal_reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP); - return NeighborInfo(temporal_reservoir, temporal_world_position, temporal_world_normal, temporal_diffuse_brdf); + return NeighborInfo(temporal_reservoir, temporal_surface.world_position, temporal_surface.world_normal, temporal_diffuse_brdf); } return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0)); @@ -171,13 +163,10 @@ fn load_spatial_reservoir(pixel_id: vec2, depth: f32, world_position: vec3< let spatial_pixel_id = get_neighbor_pixel_id(pixel_id, rng); let spatial_depth = textureLoad(depth_buffer, spatial_pixel_id, 0); - let spatial_gpixel = textureLoad(gbuffer, spatial_pixel_id, 0); - let spatial_world_position = reconstruct_world_position(spatial_pixel_id, spatial_depth); - let spatial_world_normal = octahedral_decode(unpack_24bit_normal(spatial_gpixel.a)); - let spatial_base_color = pow(unpack4x8unorm(spatial_gpixel.r).rgb, vec3(2.2)); - let spatial_diffuse_brdf = spatial_base_color / PI; - if pixel_dissimilar(depth, world_position, spatial_world_position, world_normal, spatial_world_normal) { - return NeighborInfo(empty_reservoir(), spatial_world_position, spatial_world_normal, spatial_diffuse_brdf); + let spatial_surface = gpixel_resolve(textureLoad(gbuffer, spatial_pixel_id, 0), spatial_depth, spatial_pixel_id, view.main_pass_viewport.zw, view.world_from_clip); + let spatial_diffuse_brdf = spatial_surface.material.base_color / PI; + if pixel_dissimilar(depth, world_position, spatial_surface.world_position, world_normal, spatial_surface.world_normal, view) { + return NeighborInfo(empty_reservoir(), spatial_surface.world_position, spatial_surface.world_normal, spatial_diffuse_brdf); } let spatial_pixel_index = spatial_pixel_id.x + spatial_pixel_id.y * u32(view.main_pass_viewport.z); @@ -185,7 +174,7 @@ fn load_spatial_reservoir(pixel_id: vec2, depth: f32, world_position: vec3< spatial_reservoir.radiance *= trace_point_visibility(world_position, spatial_reservoir.sample_point_world_position); - return NeighborInfo(spatial_reservoir, spatial_world_position, spatial_world_normal, spatial_diffuse_brdf); + return NeighborInfo(spatial_reservoir, spatial_surface.world_position, spatial_surface.world_normal, spatial_diffuse_brdf); } fn get_neighbor_pixel_id(center_pixel_id: vec2, rng: ptr) -> vec2 { @@ -225,40 +214,6 @@ fn isnan(x: f32) -> bool { return (bitcast(x) & 0x7fffffffu) > 0x7f800000u; } -fn reconstruct_world_position(pixel_id: vec2, depth: f32) -> vec3 { - let uv = (vec2(pixel_id) + 0.5) / view.main_pass_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.main_pass_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 -} - // Don't adjust the size of this struct without also adjusting GI_RESERVOIR_STRUCT_SIZE. struct Reservoir { sample_point_world_position: vec3, @@ -299,16 +254,14 @@ fn merge_reservoirs( // Radiances for resampling let canonical_sample_radiance = canonical_reservoir.radiance * - saturate(dot(normalize(canonical_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal)) * - canonical_diffuse_brdf; + saturate(dot(normalize(canonical_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal)); let other_sample_radiance = other_reservoir.radiance * - saturate(dot(normalize(other_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal)) * - canonical_diffuse_brdf; + saturate(dot(normalize(other_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal)); // Target functions for resampling and MIS - let canonical_target_function_canonical_sample = luminance(canonical_sample_radiance); - let canonical_target_function_other_sample = luminance(other_sample_radiance); + let canonical_target_function_canonical_sample = luminance(canonical_sample_radiance * canonical_diffuse_brdf); + let canonical_target_function_other_sample = luminance(other_sample_radiance * canonical_diffuse_brdf); // Extra target functions for MIS let other_target_function_canonical_sample = luminance( diff --git a/crates/bevy_solari/src/realtime/specular_trace.wgsl b/crates/bevy_solari/src/realtime/specular_trace.wgsl new file mode 100644 index 0000000000000..8699f60b13905 --- /dev/null +++ b/crates/bevy_solari/src/realtime/specular_trace.wgsl @@ -0,0 +1,111 @@ +#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace +#import bevy_render::maths::{orthonormalize, PI} +#import bevy_render::view::View +#import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf} +#import bevy_solari::gbuffer_utils::gpixel_resolve +#import bevy_solari::sampling::{sample_ggx_vndf, ggx_vndf_pdf} +#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX} +#import bevy_solari::world_cache::query_world_cache + +@group(1) @binding(0) var view_output: texture_storage_2d; +@group(1) @binding(5) var gi_reservoirs_a: array; +@group(1) @binding(7) var gbuffer: texture_2d; +@group(1) @binding(8) var depth_buffer: texture_depth_2d; +@group(1) @binding(12) var view: View; +struct PushConstants { frame_index: u32, reset: u32 } +var constants: PushConstants; + +@compute @workgroup_size(8, 8, 1) +fn specular_trace(@builtin(global_invocation_id) global_id: vec3) { + if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; } + + let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z); + var rng = pixel_index + constants.frame_index; + + let depth = textureLoad(depth_buffer, global_id.xy, 0); + if depth == 0.0 { + return; + } + let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip); + + let wo = normalize(view.world_position - surface.world_position); + + var radiance: vec3; + var wi: vec3; + if surface.material.roughness > 0.04 { + // Surface is very rough, reuse the ReSTIR GI reservoir + let gi_reservoir = gi_reservoirs_a[pixel_index]; + wi = normalize(gi_reservoir.sample_point_world_position - surface.world_position); + radiance = gi_reservoir.radiance * gi_reservoir.unbiased_contribution_weight; + } else { + // Surface is glossy or mirror-like, trace a new path + let TBN = orthonormalize(surface.world_normal); + let T = TBN[0]; + let B = TBN[1]; + let N = TBN[2]; + let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N)); + let wi_tangent = sample_ggx_vndf(wo_tangent, surface.material.roughness, &rng); + wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N; + let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, surface.material.roughness); + + radiance = trace_glossy_path(surface.world_position, wi, &rng) / pdf; + } + + let brdf = evaluate_specular_brdf(surface.world_normal, wo, wi, surface.material.base_color, surface.material.metallic, + surface.material.reflectance, surface.material.perceptual_roughness, surface.material.roughness); + let cos_theta = saturate(dot(wi, surface.world_normal)); + radiance *= brdf * cos_theta * view.exposure; + + var pixel_color = textureLoad(view_output, global_id.xy); + pixel_color += vec4(radiance, 0.0); + textureStore(view_output, global_id.xy, pixel_color); +} + +fn trace_glossy_path(initial_ray_origin: vec3, initial_wi: vec3, rng: ptr) -> vec3 { + var ray_origin = initial_ray_origin; + var wi = initial_wi; + + // Trace up to three bounces, getting the net throughput from them + var throughput = vec3(1.0); + for (var i = 0u; i < 3u; i += 1u) { + // Trace ray + let ray = trace_ray(ray_origin, wi, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE); + if ray.kind == RAY_QUERY_INTERSECTION_NONE { break; } + let ray_hit = resolve_ray_hit_full(ray); + + // Surface is very rough, terminate path in the world cache + if ray_hit.material.roughness > 0.04 || i == 2u { + let diffuse_brdf = ray_hit.material.base_color / PI; + return throughput * diffuse_brdf * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position); + } + + // Sample new ray direction from the GGX BRDF for next bounce + let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent); + let T = TBN[0]; + let B = TBN[1]; + let N = TBN[2]; + let wo = -wi; + let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N)); + let wi_tangent = sample_ggx_vndf(wo_tangent, ray_hit.material.roughness, rng); + wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N; + ray_origin = ray_hit.world_position; + + // Update throughput for next bounce + let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness); + let brdf = evaluate_brdf(N, wo, wi, ray_hit.material); + let cos_theta = dot(wi, N); + throughput *= (brdf * cos_theta) / pdf; + } + + return vec3(0.0); +} + +// Don't adjust the size of this struct without also adjusting GI_RESERVOIR_STRUCT_SIZE. +struct Reservoir { + sample_point_world_position: vec3, + weight_sum: f32, + radiance: vec3, + confidence_weight: f32, + sample_point_world_normal: vec3, + unbiased_contribution_weight: f32, +} diff --git a/crates/bevy_solari/src/scene/brdf.wgsl b/crates/bevy_solari/src/scene/brdf.wgsl index bc42203481928..dd78e83c7ff88 100644 --- a/crates/bevy_solari/src/scene/brdf.wgsl +++ b/crates/bevy_solari/src/scene/brdf.wgsl @@ -11,8 +11,8 @@ fn evaluate_brdf( wi: vec3, material: ResolvedMaterial, ) -> vec3 { - let diffuse_brdf = diffuse_brdf(material.base_color, material.metallic); - let specular_brdf = specular_brdf( + let diffuse_brdf = evaluate_diffuse_brdf(material.base_color, material.metallic); + let specular_brdf = evaluate_specular_brdf( world_normal, wo, wi, @@ -25,12 +25,12 @@ fn evaluate_brdf( return diffuse_brdf + specular_brdf; } -fn diffuse_brdf(base_color: vec3, metallic: f32) -> vec3 { +fn evaluate_diffuse_brdf(base_color: vec3, metallic: f32) -> vec3 { let diffuse_color = calculate_diffuse_color(base_color, metallic, 0.0, 0.0); return diffuse_color / PI; } -fn specular_brdf( +fn evaluate_specular_brdf( N: vec3, V: vec3, L: vec3,