Bevy Shader Explorer
Search across packages, versions, files, declarations, and comments. Use multiple terms to narrow results (AND), and add @vertex, @fragment, or @compute to filter by shader stage.
Packages › bevy_solari › restir_gi.wgsl › source
← Back to declarations 
1// https://intro-to-restir.cwyman.org/presentations/2023ReSTIR_Course_Notes.pdf
2enable wgpu_ray_query;
4#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
5#import bevy_pbr::prepass_bindings::PreviousViewUniforms
6#import bevy_pbr::utils::{rand_f, sample_uniform_hemisphere, uniform_hemisphere_inverse_pdf, sample_disk}
7#import bevy_render::maths::PI
8#import bevy_render::view::View
9#import bevy_solari::brdf::evaluate_diffuse_brdf
10#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar, permute_pixel}
11#import bevy_solari::sampling::{sample_random_light, trace_point_visibility, balance_heuristic, isnan}
12#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
13#import bevy_solari::world_cache::{query_world_cache, WORLD_CACHE_CELL_LIFETIME}
14#import bevy_solari::realtime_bindings::{view_output, gi_reservoirs_a, gi_reservoirs_b, gbuffer, depth_buffer, motion_vectors, previous_gbuffer, previous_depth_buffer, view, previous_view, constants, Reservoir}
15#import bevy_solari::specular_gi::DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD
17const SPATIAL_REUSE_RADIUS_PIXELS = 30.0;
18const CONFIDENCE_WEIGHT_CAP = 8.0;
20@compute @workgroup_size(8, 8, 1)
21fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3<u32>) {
22    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
24    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
25    var rng = pixel_index + constants.frame_index;
27    let depth = textureLoad(depth_buffer, global_id.xy, 0);
28    if depth == 0.0 {
29        gi_reservoirs_b[pixel_index] = empty_reservoir();
30        return;
31    }
32    let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
33    if surface.material.metallic > 0.9999 && surface.material.roughness <= DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD {
34        gi_reservoirs_b[pixel_index] = empty_reservoir();
35        return;
36    }
38    let initial_reservoir = generate_initial_reservoir(surface.world_position, surface.world_normal, &rng);
39    let temporal = load_temporal_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal);
40    let merge_result = merge_reservoirs(initial_reservoir, surface.world_position, surface.world_normal, surface.material.base_color / PI,
41        temporal.reservoir, temporal.world_position, temporal.world_normal, temporal.diffuse_brdf, &rng);
43    gi_reservoirs_b[pixel_index] = merge_result.merged_reservoir;
46@compute @workgroup_size(8, 8, 1)
47fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
48    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
50    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
51    var rng = pixel_index + constants.frame_index;
53    let depth = textureLoad(depth_buffer, global_id.xy, 0);
54    if depth == 0.0 {
55        gi_reservoirs_a[pixel_index] = empty_reservoir();
56        return;
57    }
58    let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
59    if surface.material.metallic > 0.9999 && surface.material.roughness <= DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD {
60        gi_reservoirs_a[pixel_index] = empty_reservoir();
61        return;
62    }
64    let input_reservoir = gi_reservoirs_b[pixel_index];
65    let spatial = load_spatial_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal, &rng);
66    let merge_result = merge_reservoirs(input_reservoir, surface.world_position, surface.world_normal, surface.material.base_color / PI,
67        spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.diffuse_brdf, &rng);
68    var combined_reservoir = merge_result.merged_reservoir;
70    // More accuracy, less stability
71#ifndef BIASED_RESAMPLING
72    gi_reservoirs_a[pixel_index] = combined_reservoir;
73#endif
75    combined_reservoir.unbiased_contribution_weight *= trace_point_visibility(surface.world_position + (surface.world_normal * RAY_T_MIN), combined_reservoir.sample_point_world_position);
77    // More stability, less accuracy (shadows extend further out than they should)
78#ifdef BIASED_RESAMPLING
79    gi_reservoirs_a[pixel_index] = combined_reservoir;
80#endif
82    let wo = normalize(view.world_position - surface.world_position);
83    let brdf = evaluate_diffuse_brdf(wo, merge_result.wi, surface.world_normal, surface.material);
85    var pixel_color = textureLoad(view_output, global_id.xy);
86    pixel_color += vec4(merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight * view.exposure * brdf, 0.0);
87    textureStore(view_output, global_id.xy, pixel_color);
90fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> Reservoir {
91    var reservoir = empty_reservoir();
93    let ray_direction = sample_uniform_hemisphere(world_normal, rng);
94    let ray = trace_ray(world_position + (world_normal * RAY_T_MIN), ray_direction, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE);
96    if ray.kind == RAY_QUERY_INTERSECTION_NONE {
97        return reservoir;
98    }
100    let sample_point = resolve_ray_hit_full(ray);
102    if any(sample_point.material.emissive != vec3(0.0)) {
103        return reservoir;
104    }
106    reservoir.sample_point_world_position = sample_point.world_position;
107    reservoir.sample_point_world_normal = sample_point.world_normal;
108    reservoir.confidence_weight = 1.0;
110#ifdef NO_WORLD_CACHE
111    let direct_lighting = sample_random_light(sample_point.world_position, sample_point.world_normal, rng);
112    reservoir.radiance = direct_lighting.radiance * saturate(dot(direct_lighting.wi, sample_point.world_normal));
113    reservoir.unbiased_contribution_weight = direct_lighting.inverse_pdf * uniform_hemisphere_inverse_pdf();
114#else
115    reservoir.radiance = query_world_cache(sample_point.world_position, sample_point.geometric_world_normal, view.world_position, ray.t, WORLD_CACHE_CELL_LIFETIME, rng);
116    reservoir.unbiased_contribution_weight = uniform_hemisphere_inverse_pdf();
117#endif
119    let sample_point_diffuse_brdf = sample_point.material.base_color / PI;
120    reservoir.radiance *= sample_point_diffuse_brdf;
122    return reservoir;
125fn load_temporal_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> NeighborInfo {
126    let motion_vector = textureLoad(motion_vectors, pixel_id, 0).xy;
127    let temporal_pixel_id_float = round(vec2<f32>(pixel_id) - (motion_vector * view.main_pass_viewport.zw));
128    var point_temporal_pixel_id = vec2<u32>(temporal_pixel_id_float);
130    if bool(constants.reset) {
131        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
132    }
134    if any(temporal_pixel_id_float < vec2(0.0)) || any(temporal_pixel_id_float >= view.main_pass_viewport.zw) {
135        point_temporal_pixel_id = pixel_id;
136    }
138    let permuted_temporal_pixel_id = permute_pixel(point_temporal_pixel_id, constants.frame_index, view.main_pass_viewport.zw);
139    var temporal = load_temporal_reservoir_inner(permuted_temporal_pixel_id, depth, world_position, world_normal);
141    temporal.reservoir.confidence_weight = min(temporal.reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP);
143    return temporal;
146fn load_temporal_reservoir_inner(temporal_pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> NeighborInfo {
147    // Check if the pixel features have changed heavily between the current and previous frame
148    let temporal_depth = textureLoad(previous_depth_buffer, temporal_pixel_id, 0);
149    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);
150    let temporal_diffuse_brdf = temporal_surface.material.base_color / PI;
151    if pixel_dissimilar(depth, world_position, temporal_surface.world_position, world_normal, temporal_surface.world_normal, view) {
152        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
153    }
155    let temporal_pixel_index = temporal_pixel_id.x + temporal_pixel_id.y * u32(view.main_pass_viewport.z);
156    let temporal_reservoir = gi_reservoirs_a[temporal_pixel_index];
158    return NeighborInfo(temporal_reservoir, temporal_surface.world_position, temporal_surface.world_normal, temporal_diffuse_brdf);
161fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> NeighborInfo {
162    for (var i = 0u; i < 5u; i++) {
163        let spatial_pixel_id = get_neighbor_pixel_id(pixel_id, SPATIAL_REUSE_RADIUS_PIXELS, rng);
165        let spatial_depth = textureLoad(depth_buffer, spatial_pixel_id, 0);
166        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);
167        let spatial_diffuse_brdf = spatial_surface.material.base_color / PI;
168        if pixel_dissimilar(depth, world_position, spatial_surface.world_position, world_normal, spatial_surface.world_normal, view) {
169            continue;
170        }
172        let spatial_pixel_index = spatial_pixel_id.x + spatial_pixel_id.y * u32(view.main_pass_viewport.z);
173        let spatial_reservoir = gi_reservoirs_b[spatial_pixel_index];
174        return NeighborInfo(spatial_reservoir, spatial_surface.world_position, spatial_surface.world_normal, spatial_diffuse_brdf);
175    }
177    return NeighborInfo(empty_reservoir(), world_position, world_normal, vec3(0.0));
180fn get_neighbor_pixel_id(center_pixel_id: vec2<u32>, search_radius: f32, rng: ptr<function, u32>) -> vec2<u32> {
181    var spatial_id = vec2<f32>(center_pixel_id) + sample_disk(search_radius, rng);
182    spatial_id = clamp(spatial_id, vec2(0.0), view.main_pass_viewport.zw - 1.0);
183    return vec2<u32>(spatial_id);
186struct NeighborInfo {
187    reservoir: Reservoir,
188    world_position: vec3<f32>,
189    world_normal: vec3<f32>,
190    diffuse_brdf: vec3<f32>,
193fn jacobian(
194    new_world_position: vec3<f32>,
195    original_world_position: vec3<f32>,
196    sample_point_world_position: vec3<f32>,
197    sample_point_world_normal: vec3<f32>,
198) -> f32 {
199    let r = new_world_position - sample_point_world_position;
200    let q = original_world_position - sample_point_world_position;
201    let rl = length(r);
202    let ql = length(q);
203    let phi_r = saturate(dot(r / rl, sample_point_world_normal));
204    let phi_q = saturate(dot(q / ql, sample_point_world_normal));
205    let jacobian = (phi_r * ql * ql) / (phi_q * rl * rl);
206    return select(jacobian, 0.0, isinf(jacobian) || isnan(jacobian));
209fn isinf(x: f32) -> bool {
210    return (bitcast<u32>(x) & 0x7fffffffu) == 0x7f800000u;
213fn empty_reservoir() -> Reservoir {
214    return Reservoir(
215        vec3(0.0),
216        0.0,
217        vec3(0.0),
218        0.0,
219        vec3(0.0),
220        0.0,
221    );
224struct ReservoirMergeResult {
225    merged_reservoir: Reservoir,
226    selected_sample_radiance: vec3<f32>,
227    wi: vec3<f32>,
230fn merge_reservoirs(
231    canonical_reservoir: Reservoir,
232    canonical_world_position: vec3<f32>,
233    canonical_world_normal: vec3<f32>,
234    canonical_diffuse_brdf: vec3<f32>,
235    other_reservoir: Reservoir,
236    other_world_position: vec3<f32>,
237    other_world_normal: vec3<f32>,
238    other_diffuse_brdf: vec3<f32>,
239    rng: ptr<function, u32>,
240) -> ReservoirMergeResult {
241    // Radiances for resampling
242    let canonical_sample_wi = normalize(canonical_reservoir.sample_point_world_position - canonical_world_position);
243    let other_sample_wi = normalize(other_reservoir.sample_point_world_position - canonical_world_position);
245    // Target functions for resampling and MIS
246    let canonical_target_function_canonical_sample = luminance(
247        canonical_reservoir.radiance * saturate(dot(canonical_sample_wi, canonical_world_normal)) * canonical_diffuse_brdf
248    );
249    let canonical_target_function_other_sample = luminance(
250        other_reservoir.radiance * saturate(dot(other_sample_wi, canonical_world_normal)) * canonical_diffuse_brdf
251    );
253    // Extra target functions for MIS
254    let other_target_function_canonical_sample = luminance(
255        canonical_reservoir.radiance * saturate(dot(normalize(canonical_reservoir.sample_point_world_position - other_world_position), other_world_normal)) * other_diffuse_brdf
256    );
257    let other_target_function_other_sample = luminance(
258        other_reservoir.radiance * saturate(dot(normalize(other_reservoir.sample_point_world_position - other_world_position), other_world_normal)) * other_diffuse_brdf
259    );
261    // Jacobians for resampling and MIS
262    let canonical_target_function_other_sample_jacobian = jacobian(
263        canonical_world_position,
264        other_world_position,
265        other_reservoir.sample_point_world_position,
266        other_reservoir.sample_point_world_normal
267    );
268    let other_target_function_canonical_sample_jacobian = jacobian(
269        other_world_position,
270        canonical_world_position,
271        canonical_reservoir.sample_point_world_position,
272        canonical_reservoir.sample_point_world_normal
273    );
275    // Don't merge samples with huge jacobians, as it explodes the variance
276    if canonical_target_function_other_sample_jacobian > 1.2 || other_target_function_canonical_sample_jacobian > 1.2 {
277        return ReservoirMergeResult(canonical_reservoir, canonical_reservoir.radiance, canonical_sample_wi);
278    }
280    // Resampling weight for canonical sample
281    let canonical_sample_mis_weight = balance_heuristic(
282        canonical_reservoir.confidence_weight * canonical_target_function_canonical_sample,
283        other_reservoir.confidence_weight * other_target_function_canonical_sample * other_target_function_canonical_sample_jacobian,
284    );
285    let canonical_sample_resampling_weight = canonical_sample_mis_weight * canonical_target_function_canonical_sample * canonical_reservoir.unbiased_contribution_weight;
287    // Resampling weight for other sample
288    let other_sample_mis_weight = balance_heuristic(
289        other_reservoir.confidence_weight * other_target_function_other_sample,
290        canonical_reservoir.confidence_weight * canonical_target_function_other_sample * canonical_target_function_other_sample_jacobian,
291    );
292    let other_sample_resampling_weight = other_sample_mis_weight * canonical_target_function_other_sample * other_reservoir.unbiased_contribution_weight * canonical_target_function_other_sample_jacobian;
294    // Perform resampling
295    var combined_reservoir = empty_reservoir();
296    combined_reservoir.confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
297    combined_reservoir.weight_sum = canonical_sample_resampling_weight + other_sample_resampling_weight;
299    if rand_f(rng) < other_sample_resampling_weight / combined_reservoir.weight_sum {
300        combined_reservoir.sample_point_world_position = other_reservoir.sample_point_world_position;
301        combined_reservoir.sample_point_world_normal = other_reservoir.sample_point_world_normal;
302        combined_reservoir.radiance = other_reservoir.radiance;
304        let inverse_target_function = select(0.0, 1.0 / canonical_target_function_other_sample, canonical_target_function_other_sample > 0.0);
305        combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
307        return ReservoirMergeResult(combined_reservoir, other_reservoir.radiance, other_sample_wi);
308    } else {
309        combined_reservoir.sample_point_world_position = canonical_reservoir.sample_point_world_position;
310        combined_reservoir.sample_point_world_normal = canonical_reservoir.sample_point_world_normal;
311        combined_reservoir.radiance = canonical_reservoir.radiance;
313        let inverse_target_function = select(0.0, 1.0 / canonical_target_function_canonical_sample, canonical_target_function_canonical_sample > 0.0);
314        combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
316        return ReservoirMergeResult(combined_reservoir, canonical_reservoir.radiance, canonical_sample_wi);
317    }
Contents