1// https://intro-to-restir.cwyman.org/presentations/2023ReSTIR_Course_Notes.pdf
2// https://d1qx31qr3h6wln.cloudfront.net/publications/ReSTIR%20GI.pdf
3
4enable wgpu_ray_query;
5
6#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
7#import bevy_pbr::prepass_bindings::PreviousViewUniforms
8#import bevy_pbr::utils::{rand_f, rand_range_u, sample_disk}
9#import bevy_render::maths::PI
10#import bevy_render::view::View
11#import bevy_solari::brdf::{evaluate_diffuse_brdf, evaluate_specular_brdf}
12#import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar, permute_pixel}
13#import bevy_solari::presample_light_tiles::unpack_resolved_light_sample
14#import bevy_solari::sampling::{LightSample, ResolvedLightSample, NULL_LIGHT_ID, calculate_resolved_light_contribution, resolve_and_calculate_light_contribution, resolve_light_sample, trace_light_visibility, balance_heuristic}
15#import bevy_solari::scene_bindings::{light_sources, previous_frame_light_id_translations, LIGHT_NOT_PRESENT_THIS_FRAME, RAY_T_MIN}
16#import bevy_solari::specular_gi::SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD
17#import bevy_solari::realtime_bindings::{view_output, light_tile_samples, light_tile_resolved_samples, di_reservoirs_a, di_reservoirs_b, gbuffer, depth_buffer, motion_vectors, previous_gbuffer, previous_depth_buffer, view, previous_view, constants, ResolvedLightSamplePacked}
18
19const INITIAL_SAMPLES = 8u;
20const SPATIAL_REUSE_RADIUS_PIXELS = 30.0;
21const CONFIDENCE_WEIGHT_CAP = 20.0;
22
23@compute @workgroup_size(8, 8, 1)
24fn initial_and_temporal(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) global_id: vec3<u32>) {
25    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
26
27    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
28    var rng = pixel_index + constants.frame_index;
29
30    let depth = textureLoad(depth_buffer, global_id.xy, 0);
31    if depth == 0.0 {
32        store_reservoir_b(global_id.xy, empty_reservoir());
33        return;
34    }
35    let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
36
37    let diffuse_brdf = surface.material.base_color / PI;
38    let initial_reservoir = generate_initial_reservoir(surface.world_position, surface.world_normal, diffuse_brdf, workgroup_id.xy, &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, diffuse_brdf,
41        temporal.reservoir, temporal.world_position, temporal.world_normal, temporal.diffuse_brdf, &rng);
42
43    store_reservoir_b(global_id.xy, merge_result.merged_reservoir);
44}
45
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; }
49
50    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
51    var rng = pixel_index + constants.frame_index;
52
53    let depth = textureLoad(depth_buffer, global_id.xy, 0);
54    if depth == 0.0 {
55        store_reservoir_a(global_id.xy, 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
60    let diffuse_brdf = surface.material.base_color / PI;
61    let input_reservoir = load_reservoir_b(global_id.xy);
62    let spatial = load_spatial_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal, &rng);
63    let merge_result = merge_reservoirs(input_reservoir, surface.world_position, surface.world_normal, diffuse_brdf,
64        spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.diffuse_brdf, &rng);
65    var combined_reservoir = merge_result.merged_reservoir;
66
67    // More accuracy, less stability
68#ifndef BIASED_RESAMPLING
69    store_reservoir_a(global_id.xy, combined_reservoir);
70#endif
71
72    if reservoir_valid(combined_reservoir) {
73        combined_reservoir.unbiased_contribution_weight *= trace_light_visibility(surface.world_position + (surface.world_normal * RAY_T_MIN), merge_result.selected_light_world_position);
74    }
75
76    // More stability, less accuracy (shadows extend further out than they should)
77#ifdef BIASED_RESAMPLING
78    store_reservoir_a(global_id.xy, combined_reservoir);
79#endif
80
81    let wo = normalize(view.world_position - surface.world_position);
82    var brdf = evaluate_diffuse_brdf(wo, merge_result.wi, surface.world_normal, surface.material);
83    // Only consider the specular lobe if the surface is not smooth, else leave it for the specular GI pass to handle
84    if surface.material.roughness > SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD {
85        brdf += evaluate_specular_brdf(wo, merge_result.wi, surface.world_normal, surface.material);
86    }
87
88    var pixel_color = merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight;
89    pixel_color *= brdf;
90    pixel_color += surface.material.emissive;
91    pixel_color *= view.exposure;
92    textureStore(view_output, global_id.xy, vec4(pixel_color, 1.0));
93}
94
95fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>, diffuse_brdf: vec3<f32>, workgroup_id: vec2<u32>, rng: ptr<function, u32>) -> Reservoir {
96    var workgroup_rng = (workgroup_id.x * 5782582u) + workgroup_id.y;
97    let light_tile_start = rand_range_u(128u, &workgroup_rng) * 1024u;
98
99    var reservoir = empty_reservoir();
100    var weight_sum = 0.0;
101    let mis_weight = 1.0 / f32(INITIAL_SAMPLES);
102
103    var reservoir_target_function = 0.0;
104    var light_sample_world_position = vec4(0.0);
105    var selected_tile_sample = 0u;
106    for (var i = 0u; i < INITIAL_SAMPLES; i++) {
107        let tile_sample = light_tile_start + rand_range_u(1024u, rng);
108        let resolved_light_sample = unpack_resolved_light_sample(light_tile_resolved_samples[tile_sample], view.exposure);
109        let light_contribution = calculate_resolved_light_contribution(resolved_light_sample, world_position, world_normal);
110
111        let target_function = luminance(light_contribution.radiance * diffuse_brdf * saturate(dot(light_contribution.wi, world_normal)));
112        let resampling_weight = mis_weight * (target_function * light_contribution.inverse_pdf);
113
114        weight_sum += resampling_weight;
115
116        if rand_f(rng) < resampling_weight / weight_sum {
117            reservoir_target_function = target_function;
118            light_sample_world_position = resolved_light_sample.world_position;
119            selected_tile_sample = tile_sample;
120        }
121    }
122
123    if reservoir_target_function != 0.0 {
124        reservoir.sample = light_tile_samples[selected_tile_sample];
125    }
126
127    if reservoir_valid(reservoir) {
128        let inverse_target_function = select(0.0, 1.0 / reservoir_target_function, reservoir_target_function > 0.0);
129        reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function;
130
131        reservoir.unbiased_contribution_weight *= trace_light_visibility(world_position + (world_normal * RAY_T_MIN), light_sample_world_position);
132    }
133
134    reservoir.confidence_weight = 1.0;
135    return reservoir;
136}
137
138fn load_temporal_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> NeighborInfo {
139    let motion_vector = textureLoad(motion_vectors, pixel_id, 0).xy;
140    let temporal_pixel_id_float = round(vec2<f32>(pixel_id) - (motion_vector * view.main_pass_viewport.zw));
141    var point_temporal_pixel_id = vec2<u32>(temporal_pixel_id_float);
142
143    if bool(constants.reset) {
144        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
145    }
146
147    if any(temporal_pixel_id_float < vec2(0.0)) || any(temporal_pixel_id_float >= view.main_pass_viewport.zw) {
148        point_temporal_pixel_id = pixel_id;
149    }
150
151    let permuted_temporal_pixel_id = permute_pixel(point_temporal_pixel_id, constants.frame_index, view.main_pass_viewport.zw);
152    var temporal = load_temporal_reservoir_inner(permuted_temporal_pixel_id, depth, world_position, world_normal);
153
154    // Check if the light selected in the previous frame no longer exists in the current frame (e.g. entity despawned)
155    let previous_light_id = temporal.reservoir.sample.light_id >> 16u;
156    let triangle_id = temporal.reservoir.sample.light_id & 0xFFFFu;
157    let light_id = previous_frame_light_id_translations[previous_light_id];
158    if light_id == LIGHT_NOT_PRESENT_THIS_FRAME {
159        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
160    }
161    temporal.reservoir.sample.light_id = (light_id << 16u) | triangle_id;
162
163    temporal.reservoir.confidence_weight = min(temporal.reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP);
164
165    return temporal;
166}
167
168fn load_temporal_reservoir_inner(temporal_pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> NeighborInfo {
169    // Check if the pixel features have changed heavily between the current and previous frame
170    let temporal_depth = textureLoad(previous_depth_buffer, temporal_pixel_id, 0);
171    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);
172    let temporal_diffuse_brdf = temporal_surface.material.base_color / PI;
173    if pixel_dissimilar(depth, world_position, temporal_surface.world_position, world_normal, temporal_surface.world_normal, view) {
174        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
175    }
176
177    let temporal_reservoir = load_reservoir_a(temporal_pixel_id);
178    return NeighborInfo(temporal_reservoir, temporal_surface.world_position, temporal_surface.world_normal, temporal_diffuse_brdf);
179}
180
181fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> NeighborInfo {
182    for (var i = 0u; i < 5u; i++) {
183        let spatial_pixel_id = get_neighbor_pixel_id(pixel_id, SPATIAL_REUSE_RADIUS_PIXELS, rng);
184
185        let spatial_depth = textureLoad(depth_buffer, spatial_pixel_id, 0);
186        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);
187        let spatial_diffuse_brdf = spatial_surface.material.base_color / PI;
188        if pixel_dissimilar(depth, world_position, spatial_surface.world_position, world_normal, spatial_surface.world_normal, view) {
189            continue;
190        }
191
192        let spatial_reservoir = load_reservoir_b(spatial_pixel_id);
193        return NeighborInfo(spatial_reservoir, spatial_surface.world_position, spatial_surface.world_normal, spatial_diffuse_brdf);
194    }
195
196    return NeighborInfo(empty_reservoir(), world_position, world_normal, vec3(0.0));
197}
198
199fn get_neighbor_pixel_id(center_pixel_id: vec2<u32>, search_radius: f32, rng: ptr<function, u32>) -> vec2<u32> {
200    var spatial_id = vec2<f32>(center_pixel_id) + sample_disk(search_radius, rng);
201    spatial_id = clamp(spatial_id, vec2(0.0), view.main_pass_viewport.zw - 1.0);
202    return vec2<u32>(spatial_id);
203}
204
205struct NeighborInfo {
206    reservoir: Reservoir,
207    world_position: vec3<f32>,
208    world_normal: vec3<f32>,
209    diffuse_brdf: vec3<f32>,
210}
211
212struct Reservoir {
213    sample: LightSample,
214    confidence_weight: f32,
215    unbiased_contribution_weight: f32,
216}
217
218fn empty_reservoir() -> Reservoir {
219    return Reservoir(
220        LightSample(NULL_LIGHT_ID, 0u),
221        0.0,
222        0.0,
223    );
224}
225
226fn reservoir_valid(reservoir: Reservoir) -> bool {
227    return reservoir.sample.light_id != NULL_LIGHT_ID;
228}
229
230fn pack_reservoir(reservoir: Reservoir) -> vec4<u32> {
231    let weights = bitcast<vec2<u32>>(vec2<f32>(reservoir.confidence_weight, reservoir.unbiased_contribution_weight));
232    return vec4<u32>(reservoir.sample.light_id, reservoir.sample.seed, weights);
233}
234
235fn store_reservoir_a(pixel: vec2<u32>, reservoir: Reservoir) {
236    textureStore(di_reservoirs_a, pixel, pack_reservoir(reservoir));
237}
238
239fn store_reservoir_b(pixel: vec2<u32>, reservoir: Reservoir) {
240    textureStore(di_reservoirs_b, pixel, pack_reservoir(reservoir));
241}
242
243fn unpack_reservoir(packed: vec4<u32>) -> Reservoir {
244    let weights = bitcast<vec2<f32>>(packed.zw);
245    return Reservoir(LightSample(packed.x, packed.y), weights.x, weights.y);
246}
247
248fn load_reservoir_a(pixel: vec2<u32>) -> Reservoir {
249    return unpack_reservoir(textureLoad(di_reservoirs_a, pixel));
250}
251
252fn load_reservoir_b(pixel: vec2<u32>) -> Reservoir {
253    return unpack_reservoir(textureLoad(di_reservoirs_b, pixel));
254}
255
256struct ReservoirMergeResult {
257    merged_reservoir: Reservoir,
258    selected_sample_radiance: vec3<f32>,
259    wi: vec3<f32>,
260    selected_light_world_position: vec4<f32>,
261}
262
263fn merge_reservoirs(
264    canonical_reservoir: Reservoir,
265    canonical_world_position: vec3<f32>,
266    canonical_world_normal: vec3<f32>,
267    canonical_diffuse_brdf: vec3<f32>,
268    other_reservoir: Reservoir,
269    other_world_position: vec3<f32>,
270    other_world_normal: vec3<f32>,
271    other_diffuse_brdf: vec3<f32>,
272    rng: ptr<function, u32>,
273) -> ReservoirMergeResult {
274    // Resolve each light sample once, then evaluate at both positions
275    var canonical_resolved: ResolvedLightSample;
276    if reservoir_valid(canonical_reservoir) {
277        canonical_resolved = resolve_light_sample(canonical_reservoir.sample, light_sources[canonical_reservoir.sample.light_id >> 16u]);
278    }
279    var other_resolved: ResolvedLightSample;
280    if reservoir_valid(other_reservoir) {
281        other_resolved = resolve_light_sample(other_reservoir.sample, light_sources[other_reservoir.sample.light_id >> 16u]);
282    }
283
284    // Contributions for resampling (evaluate resolved samples at canonical position)
285    let canonical_contribution_canonical_sample = resolved_reservoir_contribution(canonical_reservoir, canonical_resolved, canonical_world_position, canonical_world_normal, canonical_diffuse_brdf);
286    let canonical_contribution_other_sample = resolved_reservoir_contribution(other_reservoir, other_resolved, canonical_world_position, canonical_world_normal, canonical_diffuse_brdf);
287
288    // Extra contributions for MIS (evaluate resolved samples at other position)
289    let other_contribution_canonical_sample = resolved_reservoir_contribution(canonical_reservoir, canonical_resolved, other_world_position, other_world_normal, other_diffuse_brdf);
290    let other_contribution_other_sample = resolved_reservoir_contribution(other_reservoir, other_resolved, other_world_position, other_world_normal, other_diffuse_brdf);
291
292    // Resampling weight for canonical sample
293    let canonical_sample_mis_weight = balance_heuristic(
294        canonical_reservoir.confidence_weight * canonical_contribution_canonical_sample.target_function,
295        other_reservoir.confidence_weight * other_contribution_canonical_sample.target_function,
296    );
297    let canonical_sample_resampling_weight = canonical_sample_mis_weight * canonical_contribution_canonical_sample.target_function * 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_contribution_other_sample.target_function,
302        canonical_reservoir.confidence_weight * canonical_contribution_other_sample.target_function,
303    );
304    let other_sample_resampling_weight = other_sample_mis_weight * canonical_contribution_other_sample.target_function * other_reservoir.unbiased_contribution_weight;
305
306    // Perform resampling
307    var combined_reservoir = empty_reservoir();
308    combined_reservoir.confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
309    let weight_sum = canonical_sample_resampling_weight + other_sample_resampling_weight;
310
311    if rand_f(rng) < other_sample_resampling_weight / weight_sum {
312        combined_reservoir.sample = other_reservoir.sample;
313
314        let inverse_target_function = select(0.0, 1.0 / canonical_contribution_other_sample.target_function, canonical_contribution_other_sample.target_function > 0.0);
315        combined_reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function;
316
317        return ReservoirMergeResult(combined_reservoir, canonical_contribution_other_sample.radiance, canonical_contribution_other_sample.wi, other_resolved.world_position);
318    } else {
319        combined_reservoir.sample = canonical_reservoir.sample;
320
321        let inverse_target_function = select(0.0, 1.0 / canonical_contribution_canonical_sample.target_function, canonical_contribution_canonical_sample.target_function > 0.0);
322        combined_reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function;
323
324        return ReservoirMergeResult(combined_reservoir, canonical_contribution_canonical_sample.radiance, canonical_contribution_canonical_sample.wi, canonical_resolved.world_position);
325    }
326}
327
328struct ReservoirContribution {
329    radiance: vec3<f32>,
330    target_function: f32,
331    wi: vec3<f32>,
332}
333
334fn resolved_reservoir_contribution(reservoir: Reservoir, resolved: ResolvedLightSample, world_position: vec3<f32>, world_normal: vec3<f32>, diffuse_brdf: vec3<f32>) -> ReservoirContribution {
335    if !reservoir_valid(reservoir) { return ReservoirContribution(vec3(0.0), 0.0, vec3(0.0)); }
336    let light_contribution = calculate_resolved_light_contribution(resolved, world_position, world_normal);
337    let target_function = luminance(light_contribution.radiance * diffuse_brdf * saturate(dot(light_contribution.wi, world_normal)));
338    return ReservoirContribution(light_contribution.radiance, target_function, light_contribution.wi);
339}