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::pbr_deferred_types::unpack_24bit_normal
6#import bevy_pbr::prepass_bindings::PreviousViewUniforms
7#import bevy_pbr::rgb9e5::rgb9e5_to_vec3_
8#import bevy_pbr::utils::{rand_f, rand_range_u, octahedral_decode, sample_disk}
9#import bevy_render::maths::PI
10#import bevy_render::view::View
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}
13#import bevy_solari::scene_bindings::{light_sources, previous_frame_light_id_translations, LIGHT_NOT_PRESENT_THIS_FRAME}
14
15@group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, read_write>;
16@group(1) @binding(1) var<storage, read_write> light_tile_samples: array<LightSample>;
17@group(1) @binding(2) var<storage, read_write> light_tile_resolved_samples: array<ResolvedLightSamplePacked>;
18@group(1) @binding(3) var di_reservoirs_a: texture_storage_2d<rgba32uint, read_write>;
19@group(1) @binding(4) var di_reservoirs_b: texture_storage_2d<rgba32uint, read_write>;
20@group(1) @binding(7) var gbuffer: texture_2d<u32>;
21@group(1) @binding(8) var depth_buffer: texture_depth_2d;
22@group(1) @binding(9) var motion_vectors: texture_2d<f32>;
23@group(1) @binding(10) var previous_gbuffer: texture_2d<u32>;
24@group(1) @binding(11) var previous_depth_buffer: texture_depth_2d;
25@group(1) @binding(12) var<uniform> view: View;
26@group(1) @binding(13) var<uniform> previous_view: PreviousViewUniforms;
27struct PushConstants { frame_index: u32, reset: u32 }
28var<push_constant> constants: PushConstants;
29
30const INITIAL_SAMPLES = 32u;
31const SPATIAL_REUSE_RADIUS_PIXELS = 30.0;
32const CONFIDENCE_WEIGHT_CAP = 20.0;
33
34const NULL_RESERVOIR_SAMPLE = 0xFFFFFFFFu;
35
36@compute @workgroup_size(8, 8, 1)
37fn initial_and_temporal(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) global_id: vec3<u32>) {
38    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
39
40    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
41    var rng = pixel_index + constants.frame_index;
42
43    let depth = textureLoad(depth_buffer, global_id.xy, 0);
44    if depth == 0.0 {
45        store_reservoir_b(global_id.xy, empty_reservoir());
46        return;
47    }
48    let gpixel = textureLoad(gbuffer, global_id.xy, 0);
49    let world_position = reconstruct_world_position(global_id.xy, depth);
50    let world_normal = octahedral_decode(unpack_24bit_normal(gpixel.a));
51    let base_color = pow(unpack4x8unorm(gpixel.r).rgb, vec3(2.2));
52    let diffuse_brdf = base_color / PI;
53
54    let initial_reservoir = generate_initial_reservoir(world_position, world_normal, diffuse_brdf, workgroup_id.xy, &rng);
55    let temporal_reservoir = load_temporal_reservoir(global_id.xy, depth, world_position, world_normal);
56    let merge_result = merge_reservoirs(initial_reservoir, temporal_reservoir, world_position, world_normal, diffuse_brdf, &rng);
57
58    store_reservoir_b(global_id.xy, merge_result.merged_reservoir);
59}
60
61@compute @workgroup_size(8, 8, 1)
62fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
63    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
64
65    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
66    var rng = pixel_index + constants.frame_index;
67
68    let depth = textureLoad(depth_buffer, global_id.xy, 0);
69    if depth == 0.0 {
70        store_reservoir_a(global_id.xy, empty_reservoir());
71        textureStore(view_output, global_id.xy, vec4(vec3(0.0), 1.0));
72        return;
73    }
74    let gpixel = textureLoad(gbuffer, global_id.xy, 0);
75    let world_position = reconstruct_world_position(global_id.xy, depth);
76    let world_normal = octahedral_decode(unpack_24bit_normal(gpixel.a));
77    let base_color = pow(unpack4x8unorm(gpixel.r).rgb, vec3(2.2));
78    let diffuse_brdf = base_color / PI;
79    let emissive = rgb9e5_to_vec3_(gpixel.g);
80
81    let input_reservoir = load_reservoir_b(global_id.xy);
82    let spatial_reservoir = load_spatial_reservoir(global_id.xy, depth, world_position, world_normal, &rng);
83    let merge_result = merge_reservoirs(input_reservoir, spatial_reservoir, world_position, world_normal, diffuse_brdf, &rng);
84    let combined_reservoir = merge_result.merged_reservoir;
85
86    store_reservoir_a(global_id.xy, combined_reservoir);
87
88    var pixel_color = merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight;
89    pixel_color *= view.exposure;
90    pixel_color *= diffuse_brdf;
91    pixel_color += emissive;
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);
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, 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>) -> Reservoir {
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    let temporal_pixel_id = vec2<u32>(temporal_pixel_id_float);
142
143    // Check if the current pixel was off screen during the previous frame (current pixel is newly visible),
144    // or if all temporal history should assumed to be invalid
145    if any(temporal_pixel_id_float < vec2(0.0)) || any(temporal_pixel_id_float >= view.main_pass_viewport.zw) || bool(constants.reset) {
146        return empty_reservoir();
147    }
148
149    // Check if the pixel features have changed heavily between the current and previous frame
150    let temporal_depth = textureLoad(previous_depth_buffer, temporal_pixel_id, 0);
151    let temporal_gpixel = textureLoad(previous_gbuffer, temporal_pixel_id, 0);
152    let temporal_world_position = reconstruct_previous_world_position(temporal_pixel_id, temporal_depth);
153    let temporal_world_normal = octahedral_decode(unpack_24bit_normal(temporal_gpixel.a));
154    if pixel_dissimilar(depth, world_position, temporal_world_position, world_normal, temporal_world_normal) {
155        return empty_reservoir();
156    }
157
158    var temporal_reservoir = load_reservoir_a(temporal_pixel_id);
159
160    // Check if the light selected in the previous frame no longer exists in the current frame (e.g. entity despawned)
161    let previous_light_id = temporal_reservoir.sample.light_id >> 16u;
162    let triangle_id = temporal_reservoir.sample.light_id & 0xFFFFu;
163    let light_id = previous_frame_light_id_translations[previous_light_id];
164    if light_id == LIGHT_NOT_PRESENT_THIS_FRAME {
165        return empty_reservoir();
166    }
167    temporal_reservoir.sample.light_id = (light_id << 16u) | triangle_id;
168
169    temporal_reservoir.confidence_weight = min(temporal_reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP);
170
171    return temporal_reservoir;
172}
173
174fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> Reservoir {
175    let spatial_pixel_id = get_neighbor_pixel_id(pixel_id, rng);
176
177    let spatial_depth = textureLoad(depth_buffer, spatial_pixel_id, 0);
178    let spatial_gpixel = textureLoad(gbuffer, spatial_pixel_id, 0);
179    let spatial_world_position = reconstruct_world_position(spatial_pixel_id, spatial_depth);
180    let spatial_world_normal = octahedral_decode(unpack_24bit_normal(spatial_gpixel.a));
181    if pixel_dissimilar(depth, world_position, spatial_world_position, world_normal, spatial_world_normal) {
182        return empty_reservoir();
183    }
184
185    var spatial_reservoir = load_reservoir_b(spatial_pixel_id);
186
187    if reservoir_valid(spatial_reservoir) {
188        let resolved_light_sample = resolve_light_sample(spatial_reservoir.sample, light_sources[spatial_reservoir.sample.light_id >> 16u]);
189        spatial_reservoir.unbiased_contribution_weight *= trace_light_visibility(world_position, resolved_light_sample.world_position);
190    }
191
192    return spatial_reservoir;
193}
194
195fn get_neighbor_pixel_id(center_pixel_id: vec2<u32>, rng: ptr<function, u32>) -> vec2<u32> {
196    var spatial_id = vec2<f32>(center_pixel_id) + sample_disk(SPATIAL_REUSE_RADIUS_PIXELS, rng);
197    spatial_id = clamp(spatial_id, vec2(0.0), view.main_pass_viewport.zw - 1.0);
198    return vec2<u32>(spatial_id);
199}
200
201fn reconstruct_world_position(pixel_id: vec2<u32>, depth: f32) -> vec3<f32> {
202    let uv = (vec2<f32>(pixel_id) + 0.5) / view.main_pass_viewport.zw;
203    let xy_ndc = (uv - vec2(0.5)) * vec2(2.0, -2.0);
204    let world_pos = view.world_from_clip * vec4(xy_ndc, depth, 1.0);
205    return world_pos.xyz / world_pos.w;
206}
207
208fn reconstruct_previous_world_position(pixel_id: vec2<u32>, depth: f32) -> vec3<f32> {
209    let uv = (vec2<f32>(pixel_id) + 0.5) / view.main_pass_viewport.zw;
210    let xy_ndc = (uv - vec2(0.5)) * vec2(2.0, -2.0);
211    let world_pos = previous_view.world_from_clip * vec4(xy_ndc, depth, 1.0);
212    return world_pos.xyz / world_pos.w;
213}
214
215// Reject if tangent plane difference difference more than 0.3% or angle between normals more than 25 degrees
216fn pixel_dissimilar(depth: f32, world_position: vec3<f32>, other_world_position: vec3<f32>, normal: vec3<f32>, other_normal: vec3<f32>) -> bool {
217    // https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s22699-fast-denoising-with-self-stabilizing-recurrent-blurs.pdf#page=45
218    let tangent_plane_distance = abs(dot(normal, other_world_position - world_position));
219    let view_z = -depth_ndc_to_view_z(depth);
220
221    return tangent_plane_distance / view_z > 0.003 || dot(normal, other_normal) < 0.906;
222}
223
224fn depth_ndc_to_view_z(ndc_depth: f32) -> f32 {
225#ifdef VIEW_PROJECTION_PERSPECTIVE
226    return -view.clip_from_view[3][2]() / ndc_depth;
227#else ifdef VIEW_PROJECTION_ORTHOGRAPHIC
228    return -(view.clip_from_view[3][2] - ndc_depth) / view.clip_from_view[2][2];
229#else
230    let view_pos = view.view_from_clip * vec4(0.0, 0.0, ndc_depth, 1.0);
231    return view_pos.z / view_pos.w;
232#endif
233}
234
235struct Reservoir {
236    sample: LightSample,
237    confidence_weight: f32,
238    unbiased_contribution_weight: f32,
239}
240
241fn empty_reservoir() -> Reservoir {
242    return Reservoir(
243        LightSample(NULL_RESERVOIR_SAMPLE, 0u),
244        0.0,
245        0.0,
246    );
247}
248
249fn reservoir_valid(reservoir: Reservoir) -> bool {
250    return reservoir.sample.light_id != NULL_RESERVOIR_SAMPLE;
251}
252
253fn pack_reservoir(reservoir: Reservoir) -> vec4<u32> {
254    let weights = bitcast<vec2<u32>>(vec2<f32>(reservoir.confidence_weight, reservoir.unbiased_contribution_weight));
255    return vec4<u32>(reservoir.sample.light_id, reservoir.sample.seed, weights);
256}
257
258fn store_reservoir_a(pixel: vec2<u32>, reservoir: Reservoir) {
259    textureStore(di_reservoirs_a, pixel, pack_reservoir(reservoir));
260}
261
262fn store_reservoir_b(pixel: vec2<u32>, reservoir: Reservoir) {
263    textureStore(di_reservoirs_b, pixel, pack_reservoir(reservoir));
264}
265
266fn unpack_reservoir(packed: vec4<u32>) -> Reservoir {
267    let weights = bitcast<vec2<f32>>(packed.zw);
268    return Reservoir(LightSample(packed.x, packed.y), weights.x, weights.y);
269}
270
271fn load_reservoir_a(pixel: vec2<u32>) -> Reservoir {
272    return unpack_reservoir(textureLoad(di_reservoirs_a, pixel));
273}
274
275fn load_reservoir_b(pixel: vec2<u32>) -> Reservoir {
276    return unpack_reservoir(textureLoad(di_reservoirs_b, pixel));
277}
278
279struct ReservoirMergeResult {
280    merged_reservoir: Reservoir,
281    selected_sample_radiance: vec3<f32>,
282}
283
284fn merge_reservoirs(
285    canonical_reservoir: Reservoir,
286    other_reservoir: Reservoir,
287    world_position: vec3<f32>,
288    world_normal: vec3<f32>,
289    diffuse_brdf: vec3<f32>,
290    rng: ptr<function, u32>,
291) -> ReservoirMergeResult {
292    let mis_weight_denominator = 1.0 / (canonical_reservoir.confidence_weight + other_reservoir.confidence_weight);
293
294    let canonical_mis_weight = canonical_reservoir.confidence_weight * mis_weight_denominator;
295    let canonical_target_function = reservoir_target_function(canonical_reservoir, world_position, world_normal, diffuse_brdf);
296    let canonical_resampling_weight = canonical_mis_weight * (canonical_target_function.a * canonical_reservoir.unbiased_contribution_weight);
297
298    let other_mis_weight = other_reservoir.confidence_weight * mis_weight_denominator;
299    let other_target_function = reservoir_target_function(other_reservoir, world_position, world_normal, diffuse_brdf);
300    let other_resampling_weight = other_mis_weight * (other_target_function.a * other_reservoir.unbiased_contribution_weight);
301
302    let weight_sum = canonical_resampling_weight + other_resampling_weight;
303
304    var combined_reservoir = empty_reservoir();
305    combined_reservoir.confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
306
307    if rand_f(rng) < other_resampling_weight / weight_sum {
308        combined_reservoir.sample = other_reservoir.sample;
309
310        let inverse_target_function = select(0.0, 1.0 / other_target_function.a, other_target_function.a > 0.0);
311        combined_reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function;
312
313        return ReservoirMergeResult(combined_reservoir, other_target_function.rgb);
314    } else {
315        combined_reservoir.sample = canonical_reservoir.sample;
316
317        let inverse_target_function = select(0.0, 1.0 / canonical_target_function.a, canonical_target_function.a > 0.0);
318        combined_reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function;
319
320        return ReservoirMergeResult(combined_reservoir, canonical_target_function.rgb);
321    }
322}
323
324// TODO: Have input take ResolvedLightSample instead of reservoir.light_sample
325fn reservoir_target_function(reservoir: Reservoir, world_position: vec3<f32>, world_normal: vec3<f32>, diffuse_brdf: vec3<f32>) -> vec4<f32> {
326    if !reservoir_valid(reservoir) { return vec4(0.0); }
327    let light_contribution = resolve_and_calculate_light_contribution(reservoir.sample, world_position, world_normal).radiance;
328    let target_function = luminance(light_contribution * diffuse_brdf);
329    return vec4(light_contribution, target_function);
330}