diff --git a/crates/bevy_solari/src/realtime/bindings.wesl b/crates/bevy_solari/src/realtime/bindings.wesl index dc10648054f20..a4457978d02da 100644 --- a/crates/bevy_solari/src/realtime/bindings.wesl +++ b/crates/bevy_solari/src/realtime/bindings.wesl @@ -7,17 +7,20 @@ enable wgpu_ray_query; @group(1) @binding(0) var view_output: texture_storage_2d; @group(1) @binding(1) var light_tile_samples: array; @group(1) @binding(2) var light_tile_resolved_samples: array; -@group(1) @binding(3) var reservoirs_a: array; -@group(1) @binding(4) var 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_storage_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; -@group(1) @binding(12) var world_cache: WorldCache; -@group(1) @binding(13) var constants: SolariLightingSettings; +@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_storage_2d; +@group(1) @binding(6) var view: View; +@group(1) @binding(7) var previous_view: PreviousViewUniforms; +@group(1) @binding(8) var world_cache: WorldCache; +@group(1) @binding(9) var constants: SolariLightingSettings; + +@if(RESTIR) { +@group(1) @binding(10) var previous_gbuffer: texture_2d; +@group(1) @binding(11) var previous_depth_buffer: texture_depth_2d; +@group(1) @binding(12) var reservoirs_a: array; +@group(1) @binding(13) var reservoirs_b: array; +} @if(DLSS_RR_GUIDE_BUFFERS) { @group(2) @binding(0) var diffuse_albedo: texture_storage_2d; diff --git a/crates/bevy_solari/src/realtime/initial_path.wesl b/crates/bevy_solari/src/realtime/initial_path.wesl index edd1f8de7bc59..cb755c778b2c1 100644 --- a/crates/bevy_solari/src/realtime/initial_path.wesl +++ b/crates/bevy_solari/src/realtime/initial_path.wesl @@ -20,9 +20,14 @@ const RECONNECTION_RELAX_DISTANCE = 1.0; const CACHE_TERMINATION_MIN_SOLID_ANGLE = PI; +// What tracing one path produced. +// +// `radiance` is shaded straight into the pixel. With ReSTIR the reservoir additionally carries the +// candidate chosen for reuse, and `radiance` holds the radiance that can not be reused. +// Without ReSTIR there is no reservoir and `radiance` is the entire estimate. struct InitialSamplingResult { - reservoir: Reservoir, - non_resampled_radiance: vec3, + @if(RESTIR) reservoir: Reservoir, + radiance: vec3, } // Path vertices use the following convention: x0 = camera, x1 = primary ray hit (the G-buffer @@ -32,44 +37,54 @@ struct PathState { normal: vec3, wo: vec3, material: ResolvedMaterial, - // Throughput past x1, excluding brdf*cos at x1 + // Throughput past x1. With ReSTIR the brdf*cos at x1 is factored out of it, see x1_brdf. throughput_past_first_hit: vec3, // Reconnection vertex x2, the first BRDF-sampled hit shared by every length >= 2 candidate - x2_position: vec3, - x2_normal: vec3, - // If false, candidates built on x2 are shaded directly into non_resampled_radiance instead of + @if(RESTIR) x2_position: vec3, + @if(RESTIR) x2_normal: vec3, + // If false, candidates built on x2 are shaded directly into radiance instead of // published to the reservoir - x2_reusable: bool, - // brdf*cos at x1 for the direction toward x2 - x1_brdf: vec3, + @if(RESTIR) x2_reusable: bool, + // brdf*cos at x1 for the direction toward x2, applied at shade time + @if(RESTIR) x1_brdf: vec3, + // Radiance shaded directly at this pixel + radiance: vec3, + @if(RESTIR) reservoir: Reservoir, + @if(RESTIR) weight_sum: f32, + @if(RESTIR) selected_target_function: f32, } -fn generate_initial_reservoir(world_position: vec3, world_normal: vec3, material: ResolvedMaterial, workgroup_id: vec2, pixel_id: vec2, rng: ptr) -> InitialSamplingResult { - var reservoir = empty_reservoir(); - reservoir.confidence_weight = 1.0; - - var non_resampled_radiance = vec3(0.0); - var weight_sum = 0.0; - var selected_target_function = 0.0; +fn new_path_state(world_position: vec3, world_normal: vec3, wo: vec3, material: ResolvedMaterial) -> PathState { + var path: PathState; + path.ray_origin = world_position + (world_normal * RAY_T_MIN); + path.normal = world_normal; + path.wo = wo; + path.material = material; + path.throughput_past_first_hit = vec3(1.0); + path.radiance = vec3(0.0); + @if(RESTIR) { + path.x1_brdf = vec3(0.0); + path.x2_position = vec3(0.0); + path.x2_normal = vec3(0.0); + path.x2_reusable = false; + path.reservoir = empty_reservoir(); + path.reservoir.confidence_weight = 1.0; + path.weight_sum = 0.0; + path.selected_target_function = 0.0; + } + return path; +} +fn generate_initial_path(world_position: vec3, world_normal: vec3, material: ResolvedMaterial, workgroup_id: vec2, pixel_id: vec2, rng: ptr) -> InitialSamplingResult { let wo = normalize(view.world_position - world_position); let primary_NdotV = max(dot(world_normal, wo), 0.0001); let primary_F_ab = F_AB(material.perceptual_roughness, primary_NdotV); + var path = new_path_state(world_position, world_normal, wo, material); + @if(DLSS_RR_GUIDE_BUFFERS) var psr = psr_init(world_normal, material); - var path: PathState; - path.ray_origin = world_position + (world_normal * RAY_T_MIN); - path.normal = world_normal; - path.wo = wo; - path.material = material; - path.throughput_past_first_hit = vec3(1.0); - path.x2_position = vec3(0.0); - path.x2_normal = vec3(0.0); - path.x2_reusable = false; - path.x1_brdf = vec3(0.0); - for (var bounce = 0u; bounce < constants.max_bounces; bounce++) { let NdotV = max(dot(path.normal, path.wo), 0.0001); let F_ab = F_AB(path.material.perceptual_roughness, NdotV); @@ -79,8 +94,7 @@ fn generate_initial_reservoir(world_position: vec3, world_normal: vec3 // BRDF-sampled emissive do the work. Pure dielectrics always run NEE. let p_nee = mix(1.0, path.material.perceptual_roughness, path.material.metallic); let di_samples = select(constants.secondary_di_samples, constants.primary_di_samples, bounce == 0u); - generate_nee_candidate(&reservoir, &weight_sum, &selected_target_function, &non_resampled_radiance, - path, F_ab, p_nee, di_samples, workgroup_id, bounce, rng); + generate_nee_candidate(&path, F_ab, p_nee, di_samples, workgroup_id, bounce, rng); // Sample the BRDF and trace the next ray let next_bounce = evaluate_and_sample_brdf(path.wo, path.normal, path.material, F_ab, rng); @@ -120,17 +134,25 @@ fn generate_initial_reservoir(world_position: vec3, world_normal: vec3 // Capture x2, the first BRDF-sampled hit if bounce == 0u { - path.x2_position = ray_hit.world_position; - path.x2_normal = ray_hit.world_normal; + @if(RESTIR) { + path.x2_position = ray_hit.world_position; + path.x2_normal = ray_hit.world_normal; + + path.x1_brdf = evaluate_brdf(wo, next_bounce.wi, world_normal, material, primary_F_ab); - path.x1_brdf = evaluate_brdf(wo, next_bounce.wi, world_normal, material, primary_F_ab); + path.x2_reusable = reconnection_reusable(ray.t, p_brdf, next_bounce.wi, next_bounce.diffuse_selected, ray_hit, world_position, material.perceptual_roughness, primary_NdotV); - path.x2_reusable = reconnection_reusable(ray.t, p_brdf, next_bounce.wi, next_bounce.diffuse_selected, ray_hit, world_position, material.perceptual_roughness, primary_NdotV); + // The primary brdf*cos is applied at shade time, so divide it out of next_bounce.throughput + // to leave 1/pdf (or 1/specular_weight for mirrors, avoiding the 1/INF = 0 that would kill + // mirror GI). + path.throughput_past_first_hit *= next_bounce.throughput / max(path.x1_brdf, vec3(0.0001)); + } - // The primary brdf*cos is applied at shade time, so divide it out of next_bounce.throughput - // to leave 1/pdf (or 1/specular_weight for mirrors, avoiding the 1/INF = 0 that would kill - // mirror GI). - path.throughput_past_first_hit *= next_bounce.throughput / max(path.x1_brdf, vec3(0.0001)); + @if(!RESTIR) { + // Nothing is reused, so there is no shade-time reconnection to factor the primary + // brdf*cos out for. + path.throughput_past_first_hit *= next_bounce.throughput; + } } else { // Later bounces keep the full brdf*cos/pdf for L_at_reconnection. path.throughput_past_first_hit *= next_bounce.throughput; @@ -138,12 +160,11 @@ fn generate_initial_reservoir(world_position: vec3, world_normal: vec3 // Resample emissive hits if any(ray_hit.material.emissive > vec3(0.0)) && dot(ray_hit.world_normal, -next_bounce.wi) > 0.0 { - generate_emissive_candidate(&reservoir, &weight_sum, &selected_target_function, &non_resampled_radiance, - path, ray_hit, next_bounce.wi, p_brdf, ray.t, p_nee, di_samples, bounce, rng); + generate_emissive_candidate(&path, ray_hit, next_bounce.wi, p_brdf, ray.t, p_nee, di_samples, bounce, rng); } // Try terminating into the world cache - if terminate_into_cache(&reservoir, &weight_sum, &selected_target_function, &non_resampled_radiance, path, ray_hit, ray.t, p_brdf, bounce, rng) { + if terminate_into_cache(&path, ray_hit, ray.t, p_brdf, bounce, rng) { break; } @@ -155,29 +176,32 @@ fn generate_initial_reservoir(world_position: vec3, world_normal: vec3 // Russian roulette for early termination if bounce > 0u { - // throughput_past_first_hit has the primary brdf*cos divided out (so it can be re-applied at shade - // time), which inflates it. Multiply x1_brdf back in to get the true energy-bounded path - // throughput, which is the correct quantity for the RR survival probability. - let full_throughput = path.throughput_past_first_hit * max(path.x1_brdf, vec3(0.0001)); + // With ReSTIR, throughput_past_first_hit has the primary brdf*cos divided out (so it can be + // re-applied at shade time), which inflates it. Multiply x1_brdf back in to get the true + // energy-bounded path throughput, which is the correct quantity for the RR survival + // probability. Without ReSTIR it already is that quantity. + var full_throughput = path.throughput_past_first_hit; + @if(RESTIR) + full_throughput *= max(path.x1_brdf, vec3(0.0001)); let rr = saturate(luminance(full_throughput)); if rand_f(rng) >= rr { break; } path.throughput_past_first_hit /= rr; } } - if selected_target_function > 0.0 { - reservoir.unbiased_contribution_weight = weight_sum / selected_target_function; + var result: InitialSamplingResult; + result.radiance = path.radiance; + @if(RESTIR) { + if path.selected_target_function > 0.0 { + path.reservoir.unbiased_contribution_weight = path.weight_sum / path.selected_target_function; + } + result.reservoir = path.reservoir; } - - return InitialSamplingResult(reservoir, non_resampled_radiance); + return result; } fn generate_nee_candidate( - reservoir: ptr, - weight_sum: ptr, - selected_target_function: ptr, - non_resampled_radiance: ptr>, - path: PathState, + path: ptr, F_ab: vec2, p_nee: f32, di_samples: u32, @@ -187,7 +211,7 @@ fn generate_nee_candidate( ) { if rand_f(rng) >= p_nee { return; } - let di = sample_light_ris(path.ray_origin, path.normal, path.wo, path.material, F_ab, di_samples, workgroup_id, bounce, rng); + let di = sample_light_ris((*path).ray_origin, (*path).normal, (*path).wo, (*path).material, F_ab, di_samples, workgroup_id, bounce, rng); let di_target_function = luminance(di.brdf_radiance); if di_target_function <= 0.0 { return; } @@ -196,38 +220,51 @@ fn generate_nee_candidate( var nee_mis_weight = 1.0; if di.brdf_rays_can_hit && di.inverse_solid_angle_pdf > 0.0 { let p_nee_strategy = f32(di_samples) * (1.0 / di.inverse_solid_angle_pdf) * p_nee; - let p_brdf_at_nee = brdf_pdf(path.wo, di.wi, path.normal, path.material, F_ab); + let p_brdf_at_nee = brdf_pdf((*path).wo, di.wi, (*path).normal, (*path).material, F_ab); nee_mis_weight = power_heuristic(p_nee_strategy, p_brdf_at_nee); } if bounce == 0u { - // Bounce 0: Candidate is the light sample, stored by reference and re-resolved each frame - // nee_mis_weight goes into the target function since it gets recomputed per-pixel during reuse - let target_function = di_target_function * nee_mis_weight; - let resampling_weight = target_function * di.unbiased_contribution_weight / p_nee; - - *weight_sum += resampling_weight; - if rand_f(rng) * (*weight_sum) < resampling_weight { - (*reservoir).light_sample = di.light_sample; - *selected_target_function = target_function; + @if(!RESTIR) { + (*path).radiance += di.brdf_radiance * di.unbiased_contribution_weight * nee_mis_weight / p_nee; + } + + @if(RESTIR) { + // Bounce 0: Candidate is the light sample, stored by reference and re-resolved each frame + // nee_mis_weight goes into the target function since it gets recomputed per-pixel during reuse + let target_function = di_target_function * nee_mis_weight; + let resampling_weight = target_function * di.unbiased_contribution_weight / p_nee; + + (*path).weight_sum += resampling_weight; + if rand_f(rng) * (*path).weight_sum < resampling_weight { + (*path).reservoir.light_sample = di.light_sample; + (*path).selected_target_function = target_function; + } } } else { - // Deeper bounces: Candidate is the reconnection radiance at x2 - let L_at_reconnection = path.throughput_past_first_hit * di.brdf_radiance * di.unbiased_contribution_weight * nee_mis_weight / p_nee; - if !path.x2_reusable { - // x1 -> x2 not reuse-safe: shade directly at this pixel instead of publishing. - *non_resampled_radiance += path.x1_brdf * L_at_reconnection; - } else { - let target_function = luminance(path.x1_brdf * L_at_reconnection); - let resampling_weight = target_function; + // Deeper bounces: the contribution is the reconnection radiance at x2 + let L_at_reconnection = (*path).throughput_past_first_hit * di.brdf_radiance * di.unbiased_contribution_weight * nee_mis_weight / p_nee; + + @if(!RESTIR) { + (*path).radiance += L_at_reconnection; + } - *weight_sum += resampling_weight; - if rand_f(rng) * (*weight_sum) < resampling_weight { - (*reservoir).light_sample = LightSample(NULL_LIGHT_ID, 0u); - (*reservoir).sample_point_world_position = path.x2_position; - (*reservoir).sample_point_world_normal = octahedral_encode(path.x2_normal); - (*reservoir).radiance = L_at_reconnection; - *selected_target_function = target_function; + @if(RESTIR) { + if !(*path).x2_reusable { + // x1 -> x2 not reuse-safe: shade directly at this pixel instead of publishing. + (*path).radiance += (*path).x1_brdf * L_at_reconnection; + } else { + let target_function = luminance((*path).x1_brdf * L_at_reconnection); + let resampling_weight = target_function; + + (*path).weight_sum += resampling_weight; + if rand_f(rng) * (*path).weight_sum < resampling_weight { + (*path).reservoir.light_sample = LightSample(NULL_LIGHT_ID, 0u); + (*path).reservoir.sample_point_world_position = (*path).x2_position; + (*path).reservoir.sample_point_world_normal = octahedral_encode((*path).x2_normal); + (*path).reservoir.radiance = L_at_reconnection; + (*path).selected_target_function = target_function; + } } } } @@ -288,11 +325,7 @@ fn sample_light_ris(ray_origin: vec3, normal: vec3, wo: vec3, mat } fn generate_emissive_candidate( - reservoir: ptr, - weight_sum: ptr, - selected_target_function: ptr, - non_resampled_radiance: ptr>, - path: PathState, + path: ptr, ray_hit: ResolvedRayHitFull, wi: vec3, p_brdf: f32, @@ -308,50 +341,49 @@ fn generate_emissive_candidate( let p_light = area_pdf * ray_t * ray_t / NdotV_hit; let emissive_mis_weight = power_heuristic(p_brdf, p_light * p_nee * f32(di_samples)); - if !path.x2_reusable { - // x1 -> x2 not reuse-safe (mirror/sharp lobe or failed gate): shade directly at this pixel - // instead of publishing, since a reuse shift would waste it or make a firefly. Mirror lobes - // always land here (p_brdf = INF, footprint 0), where emissive_mis_weight is 1. - *non_resampled_radiance += path.x1_brdf * path.throughput_past_first_hit * ray_hit.material.emissive * emissive_mis_weight; - return; + @if(!RESTIR) { + (*path).radiance += (*path).throughput_past_first_hit * ray_hit.material.emissive * emissive_mis_weight; } - if bounce == 0u { - // Bounce 0: Candidate is the emissive hit - let target_function = luminance(path.x1_brdf * ray_hit.material.emissive) * emissive_mis_weight; - let resampling_weight = luminance(path.x1_brdf * path.throughput_past_first_hit * ray_hit.material.emissive) * emissive_mis_weight; - - *weight_sum += resampling_weight; - if rand_f(rng) * (*weight_sum) < resampling_weight { - (*reservoir).light_sample = LightSample(NULL_LIGHT_ID, bitcast(area_pdf)); - (*reservoir).sample_point_world_position = path.x2_position; - (*reservoir).sample_point_world_normal = octahedral_encode(path.x2_normal); - (*reservoir).radiance = ray_hit.material.emissive; - *selected_target_function = target_function; - } - } else { - // Deeper bounces: Candidate is the reconnection radiance at x2 - let emissive_L_at_reconnection = path.throughput_past_first_hit * ray_hit.material.emissive * emissive_mis_weight; - let target_function = luminance(path.x1_brdf * emissive_L_at_reconnection); - let resampling_weight = target_function; - - *weight_sum += resampling_weight; - if rand_f(rng) * (*weight_sum) < resampling_weight { - (*reservoir).light_sample = LightSample(NULL_LIGHT_ID, 0u); - (*reservoir).sample_point_world_position = path.x2_position; - (*reservoir).sample_point_world_normal = octahedral_encode(path.x2_normal); - (*reservoir).radiance = emissive_L_at_reconnection; - *selected_target_function = target_function; + @if(RESTIR) { + if !(*path).x2_reusable { + // x1 -> x2 not reuse-safe (mirror/sharp lobe or failed gate): shade directly at this pixel + // instead of publishing, since a reuse shift would waste it or make a firefly. Mirror lobes + // always land here (p_brdf = INF, footprint 0), where emissive_mis_weight is 1. + (*path).radiance += (*path).x1_brdf * (*path).throughput_past_first_hit * ray_hit.material.emissive * emissive_mis_weight; + } else if bounce == 0u { + // Bounce 0: Candidate is the emissive hit + let target_function = luminance((*path).x1_brdf * ray_hit.material.emissive) * emissive_mis_weight; + let resampling_weight = luminance((*path).x1_brdf * (*path).throughput_past_first_hit * ray_hit.material.emissive) * emissive_mis_weight; + + (*path).weight_sum += resampling_weight; + if rand_f(rng) * (*path).weight_sum < resampling_weight { + (*path).reservoir.light_sample = LightSample(NULL_LIGHT_ID, bitcast(area_pdf)); + (*path).reservoir.sample_point_world_position = (*path).x2_position; + (*path).reservoir.sample_point_world_normal = octahedral_encode((*path).x2_normal); + (*path).reservoir.radiance = ray_hit.material.emissive; + (*path).selected_target_function = target_function; + } + } else { + // Deeper bounces: Candidate is the reconnection radiance at x2 + let emissive_L_at_reconnection = (*path).throughput_past_first_hit * ray_hit.material.emissive * emissive_mis_weight; + let target_function = luminance((*path).x1_brdf * emissive_L_at_reconnection); + let resampling_weight = target_function; + + (*path).weight_sum += resampling_weight; + if rand_f(rng) * (*path).weight_sum < resampling_weight { + (*path).reservoir.light_sample = LightSample(NULL_LIGHT_ID, 0u); + (*path).reservoir.sample_point_world_position = (*path).x2_position; + (*path).reservoir.sample_point_world_normal = octahedral_encode((*path).x2_normal); + (*path).reservoir.radiance = emissive_L_at_reconnection; + (*path).selected_target_function = target_function; + } } } } fn terminate_into_cache( - reservoir: ptr, - weight_sum: ptr, - selected_target_function: ptr, - non_resampled_radiance: ptr>, - path: PathState, + path: ptr, ray_hit: ResolvedRayHitFull, ray_t: f32, p_brdf: f32, @@ -375,21 +407,27 @@ fn terminate_into_cache( let cached_radiance = query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, ray_t, WORLD_CACHE_CELL_LIFETIME, rng); let cache_outgoing = (ray_hit.material.base_color / PI) * cached_radiance; - let cache_L_at_reconnection = path.throughput_past_first_hit * cache_outgoing; - if !path.x2_reusable { - *non_resampled_radiance += path.x1_brdf * cache_L_at_reconnection; - return true; + let cache_L_at_reconnection = (*path).throughput_past_first_hit * cache_outgoing; + + @if(!RESTIR) { + (*path).radiance += cache_L_at_reconnection; } - let target_function = luminance(path.x1_brdf * cache_L_at_reconnection); - let resampling_weight = target_function; - *weight_sum += resampling_weight; - if rand_f(rng) * (*weight_sum) < resampling_weight { - (*reservoir).light_sample = LightSample(NULL_LIGHT_ID, 0u); - (*reservoir).sample_point_world_position = path.x2_position; - (*reservoir).sample_point_world_normal = octahedral_encode(path.x2_normal); - (*reservoir).radiance = cache_L_at_reconnection; - *selected_target_function = target_function; + @if(RESTIR) { + if !(*path).x2_reusable { + (*path).radiance += (*path).x1_brdf * cache_L_at_reconnection; + } else { + let target_function = luminance((*path).x1_brdf * cache_L_at_reconnection); + let resampling_weight = target_function; + (*path).weight_sum += resampling_weight; + if rand_f(rng) * (*path).weight_sum < resampling_weight { + (*path).reservoir.light_sample = LightSample(NULL_LIGHT_ID, 0u); + (*path).reservoir.sample_point_world_position = (*path).x2_position; + (*path).reservoir.sample_point_world_normal = octahedral_encode((*path).x2_normal); + (*path).reservoir.radiance = cache_L_at_reconnection; + (*path).selected_target_function = target_function; + } + } } return true; @@ -398,6 +436,7 @@ fn terminate_into_cache( // ReSTIR PT Enhanced: Algorithmic Advances for Faster and More Robust ReSTIR Path Tracing // Section 4 (sorta) // https://research.nvidia.com/labs/rtr/publication/lin2026restirptenhanced/lin2026restirptenhanced.pdf +@if(RESTIR) fn reconnection_reusable(ray_t: f32, p_brdf: f32, wi: vec3, diffuse_selected: bool, ray_hit: ResolvedRayHitFull, world_position: vec3, x1_perceptual_roughness: f32, primary_NdotV: f32) -> bool { // ray_footprint = t^2 / (p_brdf * cos_x2) is the area a sample represents at x2. It goes to 0 for // mirror lobes (p_brdf = INF) and shrinks for sharp lobes or short segments. Compared against a diff --git a/crates/bevy_solari/src/realtime/mod.rs b/crates/bevy_solari/src/realtime/mod.rs index 271e38c54c499..67e588e5b58ab 100644 --- a/crates/bevy_solari/src/realtime/mod.rs +++ b/crates/bevy_solari/src/realtime/mod.rs @@ -3,7 +3,7 @@ mod node; mod prepare; use crate::SolariPlugins; -use bevy_app::{App, Plugin}; +use bevy_app::{App, Plugin, PostUpdate}; use bevy_asset::embedded_asset; use bevy_camera::Hdr; use bevy_core_pipeline::{ @@ -14,7 +14,14 @@ use bevy_core_pipeline::{ }, schedule::{Core3d, Core3dSystems}, }; -use bevy_ecs::{component::Component, reflect::ReflectComponent, schedule::IntoScheduleConfigs}; +use bevy_ecs::{ + component::Component, + entity::Entity, + query::Has, + reflect::ReflectComponent, + schedule::IntoScheduleConfigs, + system::{Commands, Query}, +}; use bevy_pbr::DefaultOpaqueRendererMethod; use bevy_reflect::{std_traits::ReflectDefault, Reflect}; use bevy_render::{ @@ -23,7 +30,9 @@ use bevy_render::{ use bevy_shader::load_shader_library; use extract::extract_solari_lighting; use node::{init_solari_lighting_pipelines, solari_lighting}; -use prepare::prepare_solari_lighting_resources; +use prepare::{ + prepare_solari_lighting_resources, setup_raytracing_scene_needs_previous_frame_data, +}; use tracing::warn; /// Raytraced direct and indirect lighting. @@ -39,6 +48,7 @@ impl Plugin for SolariLightingPlugin { load_shader_library!(app, "presample_light_tiles.wesl"); load_shader_library!(app, "initial_path.wesl"); embedded_asset!(app, "restir.wesl"); + embedded_asset!(app, "no_restir.wesl"); load_shader_library!(app, "world_cache_query.wesl"); embedded_asset!(app, "world_cache_compact.wesl"); embedded_asset!(app, "world_cache_update.wesl"); @@ -49,10 +59,11 @@ impl Plugin for SolariLightingPlugin { } fn finish(&self, app: &mut App) { - let render_app = app.sub_app_mut(RenderApp); - - let render_device = render_app.world().resource::(); - let features = render_device.features(); + let features = app + .sub_app(RenderApp) + .world() + .resource::() + .features(); if !features.contains(SolariPlugins::required_wgpu_features()) { warn!( "SolariLightingPlugin not loaded. GPU lacks support for required features: {:?}.", @@ -61,12 +72,18 @@ impl Plugin for SolariLightingPlugin { return; } - render_app + app.add_systems(PostUpdate, manage_prepass_double_buffers); + + app.sub_app_mut(RenderApp) .add_systems(RenderStartup, init_solari_lighting_pipelines) .add_systems(ExtractSchedule, extract_solari_lighting) .add_systems( Render, - prepare_solari_lighting_resources.in_set(RenderSystems::PrepareResources), + ( + prepare_solari_lighting_resources, + setup_raytracing_scene_needs_previous_frame_data, + ) + .in_set(RenderSystems::PrepareResources), ) .add_systems( Core3d, @@ -83,18 +100,27 @@ impl Plugin for SolariLightingPlugin { /// `Msaa::Off`. #[derive(Component, Reflect, Clone)] #[reflect(Component, Default, Clone)] -#[require( - Hdr, - DeferredPrepass, - DepthPrepass, - MotionVectorPrepass, - DeferredPrepassDoubleBuffer, - DepthPrepassDoubleBuffer -)] +#[require(Hdr, DeferredPrepass, DepthPrepass, MotionVectorPrepass)] pub struct SolariLighting { + /// [ReSTIR](https://en.wikipedia.org/wiki/Spatiotemporal_reservoir_resampling) is a technique to reuse path samples + /// between pixels and frames. This dramatically reduces noise, at the cost of a few extra rays per pixel. + /// + /// However, modern denoisers cope well with very noisy input. In many cases, turning this on + /// won't dramatically improve quality after denoising. + /// + /// If you want more fine shadow detail, or have scenes with more difficult lighting conditions, + /// turning this on may improve quality and stability, at the cost of a decent chunk of performance. + /// + /// Whether to enable this setting or not will be very scene dependent. + /// + /// Defaults to `false`. + pub restir: bool, + /// Maximum confidence weight (effective temporal history length) a pixel /// can accumulate during temporal resampling. /// + /// Has no effect when [`SolariLighting::restir`] is `false`. + /// /// Higher values are more stable but slower to react to lighting changes /// and will lead to increased artifacts. pub confidence_weight_cap: f32, @@ -115,8 +141,8 @@ pub struct SolariLighting { /// Maximum number of bounces traced when generating an initial path. /// - /// Higher values capture more indirect light for greater accuracy at the cost - /// of more rays traced per frame. Lower values are faster but lose + /// Higher values capture more detail in nested reflections and more indirect lighting, + /// at the cost of more rays traced per frame. Lower values are faster but lose /// multi-bounce lighting for specular paths. pub max_bounces: u32, @@ -183,6 +209,7 @@ pub struct SolariLighting { impl Default for SolariLighting { fn default() -> Self { Self { + restir: false, confidence_weight_cap: 8.0, primary_di_samples: 8, secondary_di_samples: 4, @@ -197,3 +224,33 @@ impl Default for SolariLighting { } } } + +/// Adds or removes the prepass double-buffer components according to [`SolariLighting::restir`]. +fn manage_prepass_double_buffers( + views: Query<( + Entity, + &SolariLighting, + Has, + Has, + )>, + mut commands: Commands, +) { + for (entity, solari_lighting, deferred_double_buffered, depth_double_buffered) in &views { + let mut entity = commands.entity(entity); + if solari_lighting.restir { + if !deferred_double_buffered { + entity.insert(DeferredPrepassDoubleBuffer); + } + if !depth_double_buffered { + entity.insert(DepthPrepassDoubleBuffer); + } + } else { + if deferred_double_buffered { + entity.remove::(); + } + if depth_double_buffered { + entity.remove::(); + } + } + } +} diff --git a/crates/bevy_solari/src/realtime/no_restir.wesl b/crates/bevy_solari/src/realtime/no_restir.wesl new file mode 100644 index 0000000000000..f15fcd02288e5 --- /dev/null +++ b/crates/bevy_solari/src/realtime/no_restir.wesl @@ -0,0 +1,30 @@ +import package::realtime::gbuffer_utils::gpixel_resolve; +import package::realtime::initial_path::generate_initial_path; +import package::realtime::bindings::{constants, depth_buffer, gbuffer, view, view_output}; +import package::scene::bindings::RAY_T_MAX; +import package::realtime::world_cache_query::{query_world_cache, WORLD_CACHE_CELL_LIFETIME}; + +enable wgpu_ray_query; + +@compute @workgroup_size(8, 8, 1) +fn initial_and_shade(@builtin(workgroup_id) workgroup_id: vec3, @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_rng; + + 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 path = generate_initial_path(surface.world_position, surface.world_normal, surface.material, workgroup_id.xy, global_id.xy, &rng); + + var pixel_color = path.radiance; + pixel_color += surface.material.emissive; + pixel_color *= view.exposure; + textureStore(view_output, global_id.xy, vec4(pixel_color, 1.0)); + + @if(VISUALIZE_WORLD_CACHE) + textureStore(view_output, global_id.xy, vec4(query_world_cache(surface.world_position, surface.world_normal, view.world_position, RAY_T_MAX, WORLD_CACHE_CELL_LIFETIME, &rng) * view.exposure, 1.0)); +} diff --git a/crates/bevy_solari/src/realtime/node.rs b/crates/bevy_solari/src/realtime/node.rs index c476c7c155426..14e35aebb5927 100644 --- a/crates/bevy_solari/src/realtime/node.rs +++ b/crates/bevy_solari/src/realtime/node.rs @@ -32,6 +32,7 @@ use bevy_utils::default; #[derive(Resource)] pub struct SolariLightingPipelines { bind_group_layout: BindGroupLayoutDescriptor, + bind_group_layout_restir: BindGroupLayoutDescriptor, bind_group_layout_world_cache_active_cells_dispatch: BindGroupLayoutDescriptor, #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] bind_group_layout_resolve_dlss_rr_textures: BindGroupLayoutDescriptor, @@ -43,14 +44,25 @@ pub struct SolariLightingPipelines { sample_gi_for_world_cache_pipeline: CachedComputePipelineId, blend_new_world_cache_samples_pipeline: CachedComputePipelineId, presample_light_tiles_pipeline: CachedComputePipelineId, - initial_and_temporal_pipeline: CachedComputePipelineId, - #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] - initial_and_temporal_with_psr_pipeline: CachedComputePipelineId, - spatial_and_shade_pipeline: CachedComputePipelineId, + restir: RestirPipelines, + no_restir: NoRestirPipelines, #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] resolve_dlss_rr_textures_pipeline: CachedComputePipelineId, } +struct RestirPipelines { + initial_and_temporal: CachedComputePipelineId, + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] + initial_and_temporal_with_psr: CachedComputePipelineId, + spatial_and_shade: CachedComputePipelineId, +} + +struct NoRestirPipelines { + initial_and_shade: CachedComputePipelineId, + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] + initial_and_shade_with_psr: CachedComputePipelineId, +} + #[cfg(any(not(feature = "dlss"), feature = "force_disable_dlss"))] type SolariLightingViewQuery = ( &'static SolariLightingResources, @@ -103,15 +115,37 @@ pub fn solari_lighting( return; }; + let restir = solari_lighting_resources.reservoirs.as_ref().zip( + view_prepass_textures + .previous_deferred_view() + .zip(view_prepass_textures.previous_depth_only_view()), + ); + #[cfg(any(not(feature = "dlss"), feature = "force_disable_dlss"))] - let initial_and_temporal_pipeline = pipelines.initial_and_temporal_pipeline; - #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] - let initial_and_temporal_pipeline = if view_dlss_rr_textures.is_some() { - pipelines.initial_and_temporal_with_psr_pipeline + let (initial_pipeline_id, spatial_pipeline_id) = if restir.is_some() { + ( + pipelines.restir.initial_and_temporal, + Some(pipelines.restir.spatial_and_shade), + ) } else { - pipelines.initial_and_temporal_pipeline + (pipelines.no_restir.initial_and_shade, None) }; + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] + let (initial_pipeline_id, spatial_pipeline_id) = + match (restir.is_some(), view_dlss_rr_textures.is_some()) { + (true, true) => ( + pipelines.restir.initial_and_temporal_with_psr, + Some(pipelines.restir.spatial_and_shade), + ), + (true, false) => ( + pipelines.restir.initial_and_temporal, + Some(pipelines.restir.spatial_and_shade), + ), + (false, true) => (pipelines.no_restir.initial_and_shade_with_psr, None), + (false, false) => (pipelines.no_restir.initial_and_shade, None), + }; + let ( Some(decay_world_cache_pipeline), Some(compact_world_cache_single_block_pipeline), @@ -121,14 +155,11 @@ pub fn solari_lighting( Some(sample_gi_for_world_cache_pipeline), Some(blend_new_world_cache_samples_pipeline), Some(presample_light_tiles_pipeline), - Some(initial_and_temporal_pipeline), - Some(spatial_and_shade_pipeline), + Some(initial_pipeline), Some(scene_bind_group), Some(gbuffer), Some(depth_buffer), Some(motion_vectors), - Some(previous_gbuffer), - Some(previous_depth_buffer), Some(view_uniforms_binding), Some(previous_view_uniforms_binding), ) = ( @@ -141,14 +172,11 @@ pub fn solari_lighting( pipeline_cache.get_compute_pipeline(pipelines.sample_gi_for_world_cache_pipeline), pipeline_cache.get_compute_pipeline(pipelines.blend_new_world_cache_samples_pipeline), pipeline_cache.get_compute_pipeline(pipelines.presample_light_tiles_pipeline), - pipeline_cache.get_compute_pipeline(initial_and_temporal_pipeline), - pipeline_cache.get_compute_pipeline(pipelines.spatial_and_shade_pipeline), + pipeline_cache.get_compute_pipeline(initial_pipeline_id), &scene_bindings.bind_group, view_prepass_textures.deferred_view(), view_prepass_textures.depth_only_view(), view_prepass_textures.motion_vectors_view(), - view_prepass_textures.previous_deferred_view(), - view_prepass_textures.previous_depth_only_view(), view_uniforms.uniforms.binding(), previous_view_uniforms.uniforms.binding(), ) @@ -156,6 +184,14 @@ pub fn solari_lighting( return; }; + let spatial_and_shade_pipeline = match spatial_pipeline_id { + Some(id) => match pipeline_cache.get_compute_pipeline(id) { + None => return, + pipeline => pipeline, + }, + None => None, + }; + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] let Some(resolve_dlss_rr_textures_pipeline) = pipeline_cache.get_compute_pipeline(pipelines.resolve_dlss_rr_textures_pipeline) @@ -173,19 +209,40 @@ pub fn solari_lighting( view_target_attachment.view, s.light_tile_samples.as_entire_binding(), s.light_tile_resolved_samples.as_entire_binding(), - s.reservoirs_a.as_entire_binding(), - s.reservoirs_b.as_entire_binding(), gbuffer, depth_buffer, motion_vectors, - previous_gbuffer, - previous_depth_buffer, - view_uniforms_binding, - previous_view_uniforms_binding, + view_uniforms_binding.clone(), + previous_view_uniforms_binding.clone(), s.world_cache.as_entire_binding(), s.constants.as_entire_binding(), )), ); + + let bind_group_restir = + restir.map(|(reservoirs, (previous_gbuffer, previous_depth_buffer))| { + render_device.create_bind_group( + "solari_lighting_bind_group_restir", + &pipeline_cache.get_bind_group_layout(&pipelines.bind_group_layout_restir), + &BindGroupEntries::sequential(( + view_target_attachment.view, + s.light_tile_samples.as_entire_binding(), + s.light_tile_resolved_samples.as_entire_binding(), + gbuffer, + depth_buffer, + motion_vectors, + view_uniforms_binding, + previous_view_uniforms_binding, + s.world_cache.as_entire_binding(), + s.constants.as_entire_binding(), + previous_gbuffer, + previous_depth_buffer, + reservoirs.a.as_entire_binding(), + reservoirs.b.as_entire_binding(), + )), + ) + }); + let bind_group_world_cache_active_cells_dispatch = render_device.create_bind_group( "solari_lighting_bind_group_world_cache_active_cells_dispatch", &pipeline_cache @@ -296,15 +353,28 @@ pub fn solari_lighting( let d = diagnostics.time_span(&mut pass, "solari_lighting/lighting"); + if let Some(bind_group_restir) = &bind_group_restir { + pass.set_bind_group( + 1, + bind_group_restir, + &[ + view_uniform_offset.offset, + previous_view_uniform_offset.offset, + ], + ); + } + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] if let Some(bind_group_resolve_dlss_rr_textures) = &bind_group_resolve_dlss_rr_textures { pass.set_bind_group(2, bind_group_resolve_dlss_rr_textures, &[]); } - pass.set_pipeline(initial_and_temporal_pipeline); + pass.set_pipeline(initial_pipeline); pass.dispatch_workgroups(dx, dy, 1); - pass.set_pipeline(spatial_and_shade_pipeline); - pass.dispatch_workgroups(dx, dy, 1); + if let Some(spatial_and_shade_pipeline) = spatial_and_shade_pipeline { + pass.set_pipeline(spatial_and_shade_pipeline); + pass.dispatch_workgroups(dx, dy, 1); + } d.end(&mut pass); @@ -336,17 +406,36 @@ pub fn init_solari_lighting_pipelines( texture_storage_2d(TextureFormat::Rgba16Float, StorageTextureAccess::ReadWrite), 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_storage_2d(TextureFormat::Rg16Float, StorageTextureAccess::ReadWrite), + uniform_buffer::(true), + uniform_buffer::(true), + storage_buffer_sized(false, None), + uniform_buffer_sized(false, None), + ), + ), + ); + + let bind_group_layout_restir = BindGroupLayoutDescriptor::new( + "solari_lighting_bind_group_layout_restir", + &BindGroupLayoutEntries::sequential( + ShaderStages::COMPUTE, + ( + texture_storage_2d(TextureFormat::Rgba16Float, StorageTextureAccess::ReadWrite), + storage_buffer_sized(false, None), + storage_buffer_sized(false, None), texture_2d(TextureSampleType::Uint), texture_depth_2d(), + texture_storage_2d(TextureFormat::Rg16Float, StorageTextureAccess::ReadWrite), uniform_buffer::(true), uniform_buffer::(true), storage_buffer_sized(false, None), uniform_buffer_sized(false, None), + texture_2d(TextureSampleType::Uint), + texture_depth_2d(), + storage_buffer_sized(false, None), + storage_buffer_sized(false, None), ), ), ); @@ -374,20 +463,43 @@ pub fn init_solari_lighting_pipelines( let create_pipeline = |label: &'static str, entry_point: &'static str, shader: Handle, - extra_bind_group_layout: Option<&BindGroupLayoutDescriptor>, + restir: bool, + extra_bind_group: ExtraBindGroup, extra_shader_defs: Vec| { - let mut layout = vec![ - scene_bindings.bind_group_layout.clone(), - bind_group_layout.clone(), - ]; - if let Some(extra_bind_group_layout) = extra_bind_group_layout { - layout.push(extra_bind_group_layout.clone()); + let group_1 = if restir { + &bind_group_layout_restir + } else { + &bind_group_layout + }; + let mut layout = vec![scene_bindings.bind_group_layout.clone(), group_1.clone()]; + match extra_bind_group { + ExtraBindGroup::None => {} + ExtraBindGroup::WorldCacheDispatch => { + layout.push(bind_group_layout_world_cache_active_cells_dispatch.clone()); + } + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] + ExtraBindGroup::DlssRrGuideBuffers => { + layout.push(bind_group_layout_resolve_dlss_rr_textures.clone()); + } } let mut shader_defs = vec![ShaderDefVal::UInt( "WORLD_CACHE_SIZE".into(), WORLD_CACHE_SIZE as u32, )]; + if restir { + shader_defs.push("RESTIR".into()); + } + match extra_bind_group { + ExtraBindGroup::None => {} + ExtraBindGroup::WorldCacheDispatch => { + shader_defs.push("WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER".into()); + } + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] + ExtraBindGroup::DlssRrGuideBuffers => { + shader_defs.push("DLSS_RR_GUIDE_BUFFERS".into()); + } + } shader_defs.extend_from_slice(&extra_shader_defs); pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor { @@ -402,6 +514,7 @@ pub fn init_solari_lighting_pipelines( commands.insert_resource(SolariLightingPipelines { bind_group_layout: bind_group_layout.clone(), + bind_group_layout_restir: bind_group_layout_restir.clone(), bind_group_layout_world_cache_active_cells_dispatch: bind_group_layout_world_cache_active_cells_dispatch.clone(), #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] @@ -411,87 +524,127 @@ pub fn init_solari_lighting_pipelines( "solari_lighting_decay_world_cache_pipeline", "decay_world_cache", load_embedded_asset!(asset_server.as_ref(), "world_cache_compact.wesl"), - Some(&bind_group_layout_world_cache_active_cells_dispatch), - vec!["WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER".into()], + false, + ExtraBindGroup::WorldCacheDispatch, + vec![], ), compact_world_cache_single_block_pipeline: create_pipeline( "solari_lighting_compact_world_cache_single_block_pipeline", "compact_world_cache_single_block", load_embedded_asset!(asset_server.as_ref(), "world_cache_compact.wesl"), - Some(&bind_group_layout_world_cache_active_cells_dispatch), - vec!["WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER".into()], + false, + ExtraBindGroup::WorldCacheDispatch, + vec![], ), compact_world_cache_blocks_pipeline: create_pipeline( "solari_lighting_compact_world_cache_blocks_pipeline", "compact_world_cache_blocks", load_embedded_asset!(asset_server.as_ref(), "world_cache_compact.wesl"), - Some(&bind_group_layout_world_cache_active_cells_dispatch), - vec!["WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER".into()], + false, + ExtraBindGroup::WorldCacheDispatch, + vec![], ), compact_world_cache_write_active_cells_pipeline: create_pipeline( "solari_lighting_compact_world_cache_write_active_cells_pipeline", "compact_world_cache_write_active_cells", load_embedded_asset!(asset_server.as_ref(), "world_cache_compact.wesl"), - Some(&bind_group_layout_world_cache_active_cells_dispatch), - vec!["WORLD_CACHE_NON_ATOMIC_LIFE_BUFFER".into()], + false, + ExtraBindGroup::WorldCacheDispatch, + vec![], ), sample_di_for_world_cache_pipeline: create_pipeline( "solari_lighting_sample_di_for_world_cache_pipeline", "sample_di", load_embedded_asset!(asset_server.as_ref(), "world_cache_update.wesl"), - None, + false, + ExtraBindGroup::None, vec![], ), sample_gi_for_world_cache_pipeline: create_pipeline( "solari_lighting_sample_gi_for_world_cache_pipeline", "sample_gi", load_embedded_asset!(asset_server.as_ref(), "world_cache_update.wesl"), - None, + false, + ExtraBindGroup::None, vec!["WORLD_CACHE_QUERY_ATOMIC_MAX_LIFETIME".into()], ), blend_new_world_cache_samples_pipeline: create_pipeline( "solari_lighting_blend_new_world_cache_samples_pipeline", "blend_new_samples", load_embedded_asset!(asset_server.as_ref(), "world_cache_update.wesl"), - None, + false, + ExtraBindGroup::None, vec![], ), presample_light_tiles_pipeline: create_pipeline( "solari_lighting_presample_light_tiles_pipeline", "presample_light_tiles", load_embedded_asset!(asset_server.as_ref(), "presample_light_tiles.wesl"), - None, - vec![], - ), - initial_and_temporal_pipeline: create_pipeline( - "solari_lighting_initial_and_temporal_pipeline", - "initial_and_temporal", - load_embedded_asset!(asset_server.as_ref(), "restir.wesl"), - None, + false, + ExtraBindGroup::None, vec![], ), - #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] - initial_and_temporal_with_psr_pipeline: create_pipeline( - "solari_lighting_initial_and_temporal_with_psr_pipeline", - "initial_and_temporal", - load_embedded_asset!(asset_server.as_ref(), "restir.wesl"), - Some(&bind_group_layout_resolve_dlss_rr_textures), - vec!["DLSS_RR_GUIDE_BUFFERS".into()], - ), - spatial_and_shade_pipeline: create_pipeline( - "solari_lighting_spatial_and_shade_pipeline", - "spatial_and_shade", - load_embedded_asset!(asset_server.as_ref(), "restir.wesl"), - None, - vec!["SPATIAL_MERGE".into()], - ), + restir: RestirPipelines { + initial_and_temporal: create_pipeline( + "solari_lighting_initial_and_temporal_pipeline", + "initial_and_temporal", + load_embedded_asset!(asset_server.as_ref(), "restir.wesl"), + true, + ExtraBindGroup::None, + vec![], + ), + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] + initial_and_temporal_with_psr: create_pipeline( + "solari_lighting_initial_and_temporal_with_psr_pipeline", + "initial_and_temporal", + load_embedded_asset!(asset_server.as_ref(), "restir.wesl"), + true, + ExtraBindGroup::DlssRrGuideBuffers, + vec![], + ), + spatial_and_shade: create_pipeline( + "solari_lighting_spatial_and_shade_pipeline", + "spatial_and_shade", + load_embedded_asset!(asset_server.as_ref(), "restir.wesl"), + true, + ExtraBindGroup::None, + vec!["SPATIAL_MERGE".into()], + ), + }, + no_restir: NoRestirPipelines { + initial_and_shade: create_pipeline( + "solari_lighting_initial_and_shade_pipeline", + "initial_and_shade", + load_embedded_asset!(asset_server.as_ref(), "no_restir.wesl"), + false, + ExtraBindGroup::None, + vec![], + ), + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] + initial_and_shade_with_psr: create_pipeline( + "solari_lighting_initial_and_shade_with_psr_pipeline", + "initial_and_shade", + load_embedded_asset!(asset_server.as_ref(), "no_restir.wesl"), + false, + ExtraBindGroup::DlssRrGuideBuffers, + vec![], + ), + }, #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] resolve_dlss_rr_textures_pipeline: create_pipeline( "solari_lighting_resolve_dlss_rr_textures_pipeline", "resolve_dlss_rr_textures", load_embedded_asset!(asset_server.as_ref(), "resolve_dlss_rr_textures.wesl"), - Some(&bind_group_layout_resolve_dlss_rr_textures), - vec!["DLSS_RR_GUIDE_BUFFERS".into()], + false, + ExtraBindGroup::DlssRrGuideBuffers, + vec![], ), }); } + +enum ExtraBindGroup { + None, + WorldCacheDispatch, + #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] + DlssRrGuideBuffers, +} diff --git a/crates/bevy_solari/src/realtime/prepare.rs b/crates/bevy_solari/src/realtime/prepare.rs index be7b3547c7a26..745b58cd7339d 100644 --- a/crates/bevy_solari/src/realtime/prepare.rs +++ b/crates/bevy_solari/src/realtime/prepare.rs @@ -1,4 +1,5 @@ use super::SolariLighting; +use crate::scene::RaytracingSceneNeedsPreviousFrameData; #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] use bevy_anti_alias::dlss::{ Dlss, DlssRayReconstructionFeature, ViewDlssRayReconstructionTextures, @@ -77,7 +78,7 @@ struct SolariLightingUniforms { } impl SolariLightingUniforms { - fn new(settings: &SolariLighting, frame_count: u32) -> Self { + fn new(settings: &SolariLighting, frame_count: u32, force_reset: bool) -> Self { Self { confidence_weight_cap: settings.confidence_weight_cap, primary_di_samples: settings.primary_di_samples, @@ -90,24 +91,41 @@ impl SolariLightingUniforms { world_cache_position_base_cell_size: settings.world_cache_position_base_cell_size, world_cache_position_lod_scale: settings.world_cache_position_lod_scale, frame_rng: frame_count.wrapping_mul(5782582), - reset: settings.reset as u32, + reset: (settings.reset || force_reset) as u32, } } } +pub fn setup_raytracing_scene_needs_previous_frame_data( + views: Query<&SolariLighting>, + needs_previous_frame_data: Option>, + mut commands: Commands, +) { + let restir_used = views.iter().any(|solari_lighting| solari_lighting.restir); + match (restir_used, needs_previous_frame_data.is_some()) { + (true, false) => commands.insert_resource(RaytracingSceneNeedsPreviousFrameData), + (false, true) => commands.remove_resource::(), + _ => {} + } +} + /// Internal rendering resources used for Solari lighting. #[derive(Component)] pub struct SolariLightingResources { pub constants: Buffer, pub light_tile_samples: Buffer, pub light_tile_resolved_samples: Buffer, - pub reservoirs_a: Buffer, - pub reservoirs_b: Buffer, + pub reservoirs: Option, pub world_cache: Buffer, pub world_cache_active_cells_dispatch: Buffer, pub view_size: UVec2, } +pub struct SolariReservoirBuffers { + pub a: Buffer, + pub b: Buffer, +} + pub fn prepare_solari_lighting_resources( #[cfg(any(not(feature = "dlss"), feature = "force_disable_dlss"))] query: Query<( Entity, @@ -150,11 +168,13 @@ pub fn prepare_solari_lighting_resources( view_size = *resolution_override; } - let uniforms = SolariLightingUniforms::new(solari_lighting, frame_count.0); + let reusable = solari_lighting_resources.filter(|r| { + r.view_size == view_size && r.reservoirs.is_some() == solari_lighting.restir + }); + let uniforms = + SolariLightingUniforms::new(solari_lighting, frame_count.0, reusable.is_none()); - if let Some(solari_lighting_resources) = solari_lighting_resources - && solari_lighting_resources.view_size == view_size - { + if let Some(solari_lighting_resources) = reusable { // The constants uniform can change every frame, so always upload it. render_queue.write_buffer( &solari_lighting_resources.constants, @@ -186,16 +206,20 @@ pub fn prepare_solari_lighting_resources( mapped_at_creation: false, }); - let reservoirs_buffer = |name| { - render_device.create_buffer(&BufferDescriptor { - label: Some(name), - size: (view_size.x * view_size.y) as u64 * RESERVOIR_STRUCT_SIZE, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }) - }; - let reservoirs_a = reservoirs_buffer("solari_lighting_reservoirs_a"); - let reservoirs_b = reservoirs_buffer("solari_lighting_reservoirs_b"); + let reservoirs = solari_lighting.restir.then(|| { + let reservoirs_buffer = |name| { + render_device.create_buffer(&BufferDescriptor { + label: Some(name), + size: (view_size.x * view_size.y) as u64 * RESERVOIR_STRUCT_SIZE, + usage: BufferUsages::STORAGE, + mapped_at_creation: false, + }) + }; + SolariReservoirBuffers { + a: reservoirs_buffer("solari_lighting_reservoirs_a"), + b: reservoirs_buffer("solari_lighting_reservoirs_b"), + } + }); let world_cache = render_device.create_buffer(&BufferDescriptor { label: Some("solari_lighting_world_cache"), @@ -215,8 +239,7 @@ pub fn prepare_solari_lighting_resources( constants, light_tile_samples, light_tile_resolved_samples, - reservoirs_a, - reservoirs_b, + reservoirs, world_cache, world_cache_active_cells_dispatch, view_size, diff --git a/crates/bevy_solari/src/realtime/restir.wesl b/crates/bevy_solari/src/realtime/restir.wesl index ee1b97780b27e..9d865a4aa8c78 100644 --- a/crates/bevy_solari/src/realtime/restir.wesl +++ b/crates/bevy_solari/src/realtime/restir.wesl @@ -5,7 +5,7 @@ import bevy_pbr::render::utils::{rand_f, rand_u, sample_disk}; import bevy_render::utils::octahedral_decode; import package::scene::brdf::{brdf_pdf, evaluate_brdf, F_AB}; import package::realtime::gbuffer_utils::{gpixel_resolve, permute_pixel, pixel_dissimilar}; -import package::realtime::initial_path::{generate_initial_reservoir, InitialSamplingResult}; +import package::realtime::initial_path::{generate_initial_path, InitialSamplingResult}; import package::realtime::bindings::{depth_buffer, empty_reservoir, gbuffer, motion_vectors, previous_depth_buffer, previous_gbuffer, previous_view, reservoirs_a, reservoirs_b, Reservoir, constants, view, view_output}; import package::scene::sampling::{balance_heuristic, calculate_resolved_light_contribution, isinf, isnan, LightSample, NULL_LIGHT_ID, power_heuristic, resolve_light_sample, ResolvedLightSample, trace_visibility, trace_visibility_previous_frame}; import package::scene::bindings::{light_sources, LIGHT_NOT_PRESENT_THIS_FRAME, previous_frame_light_id_translations, RAY_T_MAX, RAY_T_MIN, ResolvedMaterial}; @@ -29,8 +29,8 @@ fn initial_and_temporal(@builtin(workgroup_id) workgroup_id: vec3, @builtin } 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 = generate_initial_reservoir(surface.world_position, surface.world_normal, surface.material, workgroup_id.xy, global_id.xy, &rng); - textureStore(view_output, global_id.xy, vec4(initial.non_resampled_radiance, 0.0)); + let initial = generate_initial_path(surface.world_position, surface.world_normal, surface.material, workgroup_id.xy, global_id.xy, &rng); + textureStore(view_output, global_id.xy, vec4(initial.radiance, 0.0)); let temporal = load_temporal_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal); let previous_camera_homogeneous = previous_view.world_from_clip * (previous_view.clip_from_view * vec4(0.0, 0.0, 0.0, 1.0)); diff --git a/crates/bevy_solari/src/scene/binder.rs b/crates/bevy_solari/src/scene/binder.rs index ad84ff0475085..0d38389d0b023 100644 --- a/crates/bevy_solari/src/scene/binder.rs +++ b/crates/bevy_solari/src/scene/binder.rs @@ -28,6 +28,13 @@ const MAX_TEXTURE_COUNT: NonZeroU32 = NonZeroU32::new(5_000).unwrap(); const TEXTURE_MAP_NONE: u32 = u32::MAX; const LIGHT_NOT_PRESENT_THIS_FRAME: u32 = u32::MAX; +/// Insert this resource into the render world to make the raytracing scene maintain the previous frame's +/// TLAS and its light id translation table. +/// +/// This is useful for temporal techniques that need last frame's data. +#[derive(Resource, Default)] +pub struct RaytracingSceneNeedsPreviousFrameData; + #[derive(Resource)] pub struct RaytracingSceneBindings { pub bind_group: Option, @@ -49,6 +56,7 @@ pub fn prepare_raytracing_scene_bindings( Option<&PreviousGlobalTransform>, )>, directional_lights_query: Query<(Entity, &ExtractedDirectionalLight)>, + needs_previous_frame_data: Option>, mesh_allocator: Res, blas_manager: Res, material_assets: Res, @@ -218,10 +226,12 @@ pub fn prepare_raytracing_scene_bindings( (index_slice.range.len() / 3) as u32, )); - this_frame_entity_to_light_id.insert(entity, light_sources.get().len() as u32 - 1); - raytracing_scene_bindings - .previous_frame_light_entities - .push(entity); + if needs_previous_frame_data.is_some() { + this_frame_entity_to_light_id.insert(entity, light_sources.get().len() as u32 - 1); + raytracing_scene_bindings + .previous_frame_light_entities + .push(entity); + } } instance_id += 1; @@ -241,20 +251,24 @@ pub fn prepare_raytracing_scene_bindings( .get_mut() .push(GpuLightSource::new_directional_light(directional_light_id)); - this_frame_entity_to_light_id.insert(entity, light_sources.get().len() as u32 - 1); - raytracing_scene_bindings - .previous_frame_light_entities - .push(entity); + if needs_previous_frame_data.is_some() { + this_frame_entity_to_light_id.insert(entity, light_sources.get().len() as u32 - 1); + raytracing_scene_bindings + .previous_frame_light_entities + .push(entity); + } } - for previous_frame_light_entity in previous_frame_light_entities { - let current_frame_index = this_frame_entity_to_light_id - .get(&previous_frame_light_entity) - .copied() - .unwrap_or(LIGHT_NOT_PRESENT_THIS_FRAME); - previous_frame_light_id_translations - .get_mut() - .push(current_frame_index); + if needs_previous_frame_data.is_some() { + for previous_frame_light_entity in previous_frame_light_entities { + let current_frame_index = this_frame_entity_to_light_id + .get(&previous_frame_light_entity) + .copied() + .unwrap_or(LIGHT_NOT_PRESENT_THIS_FRAME); + previous_frame_light_id_translations + .get_mut() + .push(current_frame_index); + } } if light_sources.get().len() > u16::MAX as usize { @@ -313,7 +327,9 @@ pub fn prepare_raytracing_scene_bindings( )), )); - raytracing_scene_bindings.previous_frame_tlas = Some(tlas); + if needs_previous_frame_data.is_some() { + raytracing_scene_bindings.previous_frame_tlas = Some(tlas); + } } impl RaytracingSceneBindings { diff --git a/crates/bevy_solari/src/scene/mod.rs b/crates/bevy_solari/src/scene/mod.rs index 3b7d96f189f0b..ded973f9f3c7b 100644 --- a/crates/bevy_solari/src/scene/mod.rs +++ b/crates/bevy_solari/src/scene/mod.rs @@ -4,7 +4,7 @@ mod extract; mod types; use bevy_shader::load_shader_library; -pub use binder::RaytracingSceneBindings; +pub use binder::{RaytracingSceneBindings, RaytracingSceneNeedsPreviousFrameData}; pub use types::RaytracingMesh3d; use crate::SolariPlugins; diff --git a/examples/3d/solari.rs b/examples/3d/solari.rs index 2156be82c7c69..abc5f69f9794a 100644 --- a/examples/3d/solari.rs +++ b/examples/3d/solari.rs @@ -70,6 +70,8 @@ fn main() { #[cfg(all(feature = "dlss", not(feature = "force_disable_dlss")))] app.add_systems(Update, toggle_dlss_rr); + app.add_systems(Update, toggle_restir); + if args.many_lights != Some(true) { app.add_systems(Update, (pause_scene, toggle_lights, patrol_path)); } @@ -472,6 +474,15 @@ fn toggle_dlss_rr( } } +fn toggle_restir( + key_input: Res>, + mut solari_lighting: Single<&mut SolariLighting>, +) { + if key_input.just_pressed(KeyCode::KeyR) { + solari_lighting.restir = !solari_lighting.restir; + } +} + fn pause_scene(mut time: ResMut>, key_input: Res>) { if key_input.just_pressed(KeyCode::Space) { time.toggle(); @@ -556,6 +567,7 @@ struct ControlText; fn update_control_text( mut text: Single<&mut Text, With>, + solari_lighting: Single<&SolariLighting>, robot_light_material: Option>, materials: Res>, directional_light: Query>, @@ -609,6 +621,12 @@ fn update_control_text( #[cfg(any(not(feature = "dlss"), feature = "force_disable_dlss"))] text.0 .push_str("\nDenoising: App not compiled with DLSS support"); + + if solari_lighting.restir { + text.0.push_str("\n(R): Disable ReSTIR"); + } else { + text.0.push_str("\n(R): Enable ReSTIR"); + } } #[derive(Component)]