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