1// https://intro-to-restir.cwyman.org/presentations/2023ReSTIR_Course_Notes.pdf
2
3#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
4#import bevy_pbr::prepass_bindings::PreviousViewUniforms
5#import bevy_pbr::utils::{rand_f, sample_uniform_hemisphere, uniform_hemisphere_inverse_pdf, sample_disk}
6#import bevy_render::maths::PI
7#import bevy_render::view::View
8#import bevy_solari::brdf::evaluate_diffuse_brdf
9#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar, permute_pixel}
10#import bevy_solari::sampling::{sample_random_light, trace_point_visibility, balance_heuristic}
11#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
12#import bevy_solari::world_cache::{query_world_cache, WORLD_CACHE_CELL_LIFETIME}
13
14@group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, read_write>;
15@group(1) @binding(5) var<storage, read_write> gi_reservoirs_a: array<Reservoir>;
16@group(1) @binding(6) var<storage, read_write> gi_reservoirs_b: array<Reservoir>;
17@group(1) @binding(7) var gbuffer: texture_2d<u32>;
18@group(1) @binding(8) var depth_buffer: texture_depth_2d;
19@group(1) @binding(9) var motion_vectors: texture_2d<f32>;
20@group(1) @binding(10) var previous_gbuffer: texture_2d<u32>;
21@group(1) @binding(11) var previous_depth_buffer: texture_depth_2d;
22@group(1) @binding(12) var<uniform> view: View;
23@group(1) @binding(13) var<uniform> previous_view: PreviousViewUniforms;
24struct PushConstants { frame_index: u32, reset: u32 }
25var<push_constant> constants: PushConstants;
26
27const SPATIAL_REUSE_RADIUS_PIXELS = 30.0;
28const CONFIDENCE_WEIGHT_CAP = 8.0;
29
30@compute @workgroup_size(8, 8, 1)
31fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3<u32>) {
32    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
33
34    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
35    var rng = pixel_index + constants.frame_index;
36
37    let depth = textureLoad(depth_buffer, global_id.xy, 0);
38    if depth == 0.0 {
39        gi_reservoirs_b[pixel_index] = empty_reservoir();
40        return;
41    }
42    let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
43
44    let initial_reservoir = generate_initial_reservoir(surface.world_position, surface.world_normal, &rng);
45    let temporal = load_temporal_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal);
46    let merge_result = merge_reservoirs(initial_reservoir, surface.world_position, surface.world_normal, surface.material.base_color / PI,
47        temporal.reservoir, temporal.world_position, temporal.world_normal, temporal.diffuse_brdf, &rng);
48
49    gi_reservoirs_b[pixel_index] = merge_result.merged_reservoir;
50}
51
52@compute @workgroup_size(8, 8, 1)
53fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
54    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
55
56    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
57    var rng = pixel_index + constants.frame_index;
58
59    let depth = textureLoad(depth_buffer, global_id.xy, 0);
60    if depth == 0.0 {
61        gi_reservoirs_a[pixel_index] = empty_reservoir();
62        return;
63    }
64    let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
65
66    let input_reservoir = gi_reservoirs_b[pixel_index];
67    let spatial = load_spatial_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal, &rng);
68    let merge_result = merge_reservoirs(input_reservoir, surface.world_position, surface.world_normal, surface.material.base_color / PI,
69        spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.diffuse_brdf, &rng);
70    var combined_reservoir = merge_result.merged_reservoir;
71
72    // More accuracy, less stability
73#ifndef BIASED_RESAMPLING
74    gi_reservoirs_a[pixel_index] = combined_reservoir;
75#endif
76
77    combined_reservoir.unbiased_contribution_weight *= trace_point_visibility(surface.world_position, combined_reservoir.sample_point_world_position);
78
79    // More stability, less accuracy (shadows extend further out than they should)
80#ifdef BIASED_RESAMPLING
81    gi_reservoirs_a[pixel_index] = combined_reservoir;
82#endif
83
84    let brdf = evaluate_diffuse_brdf(surface.material.base_color, surface.material.metallic);
85
86    var pixel_color = textureLoad(view_output, global_id.xy);
87    pixel_color += vec4(merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight * view.exposure * brdf, 0.0);
88    textureStore(view_output, global_id.xy, pixel_color);
89}
90
91fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> Reservoir {
92    var reservoir = empty_reservoir();
93
94    let ray_direction = sample_uniform_hemisphere(world_normal, rng);
95    let ray_hit = trace_ray(world_position, ray_direction, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE);
96
97    if ray_hit.kind == RAY_QUERY_INTERSECTION_NONE {
98        return reservoir;
99    }
100
101    let sample_point = resolve_ray_hit_full(ray_hit);
102
103    if all(sample_point.material.emissive != vec3(0.0)) {
104        return reservoir;
105    }
106
107    reservoir.sample_point_world_position = sample_point.world_position;
108    reservoir.sample_point_world_normal = sample_point.world_normal;
109    reservoir.confidence_weight = 1.0;
110
111#ifdef NO_WORLD_CACHE
112    let direct_lighting = sample_random_light(sample_point.world_position, sample_point.world_normal, rng);
113    reservoir.radiance = direct_lighting.radiance;
114    reservoir.unbiased_contribution_weight = direct_lighting.inverse_pdf * uniform_hemisphere_inverse_pdf();
115#else
116    reservoir.radiance = query_world_cache(sample_point.world_position, sample_point.geometric_world_normal, view.world_position, WORLD_CACHE_CELL_LIFETIME, rng);
117    reservoir.unbiased_contribution_weight = uniform_hemisphere_inverse_pdf();
118#endif
119
120    let sample_point_diffuse_brdf = sample_point.material.base_color / PI;
121    reservoir.radiance *= sample_point_diffuse_brdf;
122
123    return reservoir;
124}
125
126fn load_temporal_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> NeighborInfo {
127    let motion_vector = textureLoad(motion_vectors, pixel_id, 0).xy;
128    let temporal_pixel_id_float = round(vec2<f32>(pixel_id) - (motion_vector * view.main_pass_viewport.zw));
129
130    // Check if the current pixel was off screen during the previous frame (current pixel is newly visible),
131    // or if all temporal history should assumed to be invalid
132    if any(temporal_pixel_id_float < vec2(0.0)) || any(temporal_pixel_id_float >= view.main_pass_viewport.zw) || bool(constants.reset) {
133        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
134    }
135
136    let permuted_temporal_pixel_id = permute_pixel(vec2<u32>(temporal_pixel_id_float), constants.frame_index, view.main_pass_viewport.zw);
137    var temporal = load_temporal_reservoir_inner(permuted_temporal_pixel_id, depth, world_position, world_normal);
138
139    // If permuted reprojection failed (tends to happen on object edges), try point reprojection
140    if all(temporal.reservoir.radiance == vec3(0.0)) {
141        temporal = load_temporal_reservoir_inner(vec2<u32>(temporal_pixel_id_float), depth, world_position, world_normal);
142    }
143
144    temporal.reservoir.confidence_weight = min(temporal.reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP);
145
146    return temporal;
147}
148
149fn load_temporal_reservoir_inner(temporal_pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> NeighborInfo {
150    // Check if the pixel features have changed heavily between the current and previous frame
151    let temporal_depth = textureLoad(previous_depth_buffer, temporal_pixel_id, 0);
152    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);
153    let temporal_diffuse_brdf = temporal_surface.material.base_color / PI;
154    if pixel_dissimilar(depth, world_position, temporal_surface.world_position, world_normal, temporal_surface.world_normal, view) {
155        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
156    }
157
158    let temporal_pixel_index = temporal_pixel_id.x + temporal_pixel_id.y * u32(view.main_pass_viewport.z);
159    let temporal_reservoir = gi_reservoirs_a[temporal_pixel_index];
160
161    return NeighborInfo(temporal_reservoir, temporal_surface.world_position, temporal_surface.world_normal, temporal_diffuse_brdf);
162}
163
164fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> NeighborInfo {
165    for (var i = 0u; i < 5u; i++) {
166        let spatial_pixel_id = get_neighbor_pixel_id(pixel_id, rng);
167
168        let spatial_depth = textureLoad(depth_buffer, spatial_pixel_id, 0);
169        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);
170        let spatial_diffuse_brdf = spatial_surface.material.base_color / PI;
171        if pixel_dissimilar(depth, world_position, spatial_surface.world_position, world_normal, spatial_surface.world_normal, view) {
172            continue;
173        }
174
175        let spatial_pixel_index = spatial_pixel_id.x + spatial_pixel_id.y * u32(view.main_pass_viewport.z);
176        let spatial_reservoir = gi_reservoirs_b[spatial_pixel_index];
177        return NeighborInfo(spatial_reservoir, spatial_surface.world_position, spatial_surface.world_normal, spatial_diffuse_brdf);
178    }
179
180    return NeighborInfo(empty_reservoir(), world_position, world_normal, vec3(0.0));
181}
182
183fn get_neighbor_pixel_id(center_pixel_id: vec2<u32>, rng: ptr<function, u32>) -> vec2<u32> {
184    var spatial_id = vec2<f32>(center_pixel_id) + sample_disk(SPATIAL_REUSE_RADIUS_PIXELS, rng);
185    spatial_id = clamp(spatial_id, vec2(0.0), view.main_pass_viewport.zw - 1.0);
186    return vec2<u32>(spatial_id);
187}
188
189struct NeighborInfo {
190    reservoir: Reservoir,
191    world_position: vec3<f32>,
192    world_normal: vec3<f32>,
193    diffuse_brdf: vec3<f32>,
194}
195
196fn jacobian(
197    new_world_position: vec3<f32>,
198    original_world_position: vec3<f32>,
199    sample_point_world_position: vec3<f32>,
200    sample_point_world_normal: vec3<f32>,
201) -> f32 {
202    let r = new_world_position - sample_point_world_position;
203    let q = original_world_position - sample_point_world_position;
204    let rl = length(r);
205    let ql = length(q);
206    let phi_r = saturate(dot(r / rl, sample_point_world_normal));
207    let phi_q = saturate(dot(q / ql, sample_point_world_normal));
208    let jacobian = (phi_r * ql * ql) / (phi_q * rl * rl);
209    return select(jacobian, 0.0, isinf(jacobian) || isnan(jacobian));
210}
211
212fn isinf(x: f32) -> bool {
213    return (bitcast<u32>(x) & 0x7fffffffu) == 0x7f800000u;
214}
215
216fn isnan(x: f32) -> bool {
217    return (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
218}
219
220// Don't adjust the size of this struct without also adjusting GI_RESERVOIR_STRUCT_SIZE.
221struct Reservoir {
222    sample_point_world_position: vec3<f32>,
223    weight_sum: f32,
224    radiance: vec3<f32>,
225    confidence_weight: f32,
226    sample_point_world_normal: vec3<f32>,
227    unbiased_contribution_weight: f32,
228}
229
230fn empty_reservoir() -> Reservoir {
231    return Reservoir(
232        vec3(0.0),
233        0.0,
234        vec3(0.0),
235        0.0,
236        vec3(0.0),
237        0.0,
238    );
239}
240
241struct ReservoirMergeResult {
242    merged_reservoir: Reservoir,
243    selected_sample_radiance: vec3<f32>,
244}
245
246fn merge_reservoirs(
247    canonical_reservoir: Reservoir,
248    canonical_world_position: vec3<f32>,
249    canonical_world_normal: vec3<f32>,
250    canonical_diffuse_brdf: vec3<f32>,
251    other_reservoir: Reservoir,
252    other_world_position: vec3<f32>,
253    other_world_normal: vec3<f32>,
254    other_diffuse_brdf: vec3<f32>,
255    rng: ptr<function, u32>,
256) -> ReservoirMergeResult {
257    // Radiances for resampling
258    let canonical_sample_radiance = canonical_reservoir.radiance * saturate(dot(normalize(canonical_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal));
259    let other_sample_radiance = other_reservoir.radiance * saturate(dot(normalize(other_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal));
260
261    // Target functions for resampling and MIS
262    let canonical_target_function_canonical_sample = luminance(canonical_sample_radiance * canonical_diffuse_brdf);
263    let canonical_target_function_other_sample = luminance(other_sample_radiance * canonical_diffuse_brdf);
264
265    // Extra target functions for MIS
266    let other_target_function_canonical_sample = luminance(
267        canonical_reservoir.radiance * saturate(dot(normalize(canonical_reservoir.sample_point_world_position - other_world_position), other_world_normal)) * other_diffuse_brdf
268    );
269    let other_target_function_other_sample = luminance(
270        other_reservoir.radiance * saturate(dot(normalize(other_reservoir.sample_point_world_position - other_world_position), other_world_normal)) * other_diffuse_brdf
271    );
272
273    // Jacobians for resampling and MIS
274    let canonical_target_function_other_sample_jacobian = jacobian(
275        canonical_world_position,
276        other_world_position,
277        other_reservoir.sample_point_world_position,
278        other_reservoir.sample_point_world_normal
279    );
280    let other_target_function_canonical_sample_jacobian = jacobian(
281        other_world_position,
282        canonical_world_position,
283        canonical_reservoir.sample_point_world_position,
284        canonical_reservoir.sample_point_world_normal
285    );
286
287    // Don't merge samples with huge jacobians, as it explodes the variance
288    if canonical_target_function_other_sample_jacobian > 1.2 {
289        return ReservoirMergeResult(canonical_reservoir, canonical_sample_radiance);
290    }
291
292    // Resampling weight for canonical sample
293    let canonical_sample_mis_weight = balance_heuristic(
294        canonical_reservoir.confidence_weight * canonical_target_function_canonical_sample,
295        other_reservoir.confidence_weight * other_target_function_canonical_sample * other_target_function_canonical_sample_jacobian,
296    );
297    let canonical_sample_resampling_weight = canonical_sample_mis_weight * canonical_target_function_canonical_sample * canonical_reservoir.unbiased_contribution_weight;
298
299    // Resampling weight for other sample
300    let other_sample_mis_weight = balance_heuristic(
301        other_reservoir.confidence_weight * other_target_function_other_sample,
302        canonical_reservoir.confidence_weight * canonical_target_function_other_sample * canonical_target_function_other_sample_jacobian,
303    );
304    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;
305
306    // Perform resampling
307    var combined_reservoir = empty_reservoir();
308    combined_reservoir.confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
309    combined_reservoir.weight_sum = canonical_sample_resampling_weight + other_sample_resampling_weight;
310
311    if rand_f(rng) < other_sample_resampling_weight / combined_reservoir.weight_sum {
312        combined_reservoir.sample_point_world_position = other_reservoir.sample_point_world_position;
313        combined_reservoir.sample_point_world_normal = other_reservoir.sample_point_world_normal;
314        combined_reservoir.radiance = other_reservoir.radiance;
315
316        let inverse_target_function = select(0.0, 1.0 / canonical_target_function_other_sample, canonical_target_function_other_sample > 0.0);
317        combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
318
319        return ReservoirMergeResult(combined_reservoir, other_sample_radiance);
320    } else {
321        combined_reservoir.sample_point_world_position = canonical_reservoir.sample_point_world_position;
322        combined_reservoir.sample_point_world_normal = canonical_reservoir.sample_point_world_normal;
323        combined_reservoir.radiance = canonical_reservoir.radiance;
324
325        let inverse_target_function = select(0.0, 1.0 / canonical_target_function_canonical_sample, canonical_target_function_canonical_sample > 0.0);
326        combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
327
328        return ReservoirMergeResult(combined_reservoir, canonical_sample_radiance);
329    }
330}