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::pbr_deferred_types::unpack_24bit_normal
5#import bevy_pbr::prepass_bindings::PreviousViewUniforms
6#import bevy_pbr::utils::{rand_f, sample_uniform_hemisphere, uniform_hemisphere_inverse_pdf, sample_disk, octahedral_decode}
7#import bevy_render::maths::PI
8#import bevy_render::view::View
9#import bevy_solari::sampling::{sample_random_light, trace_point_visibility}
10#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
11#import bevy_solari::world_cache::query_world_cache
12
13@group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, read_write>;
14@group(1) @binding(5) var<storage, read_write> gi_reservoirs_a: array<Reservoir>;
15@group(1) @binding(6) var<storage, read_write> gi_reservoirs_b: array<Reservoir>;
16@group(1) @binding(7) var gbuffer: texture_2d<u32>;
17@group(1) @binding(8) var depth_buffer: texture_depth_2d;
18@group(1) @binding(9) var motion_vectors: texture_2d<f32>;
19@group(1) @binding(10) var previous_gbuffer: texture_2d<u32>;
20@group(1) @binding(11) var previous_depth_buffer: texture_depth_2d;
21@group(1) @binding(12) var<uniform> view: View;
22@group(1) @binding(13) var<uniform> previous_view: PreviousViewUniforms;
23struct PushConstants { frame_index: u32, reset: u32 }
24var<push_constant> constants: PushConstants;
25
26const SPATIAL_REUSE_RADIUS_PIXELS = 30.0;
27const CONFIDENCE_WEIGHT_CAP = 8.0;
28
29@compute @workgroup_size(8, 8, 1)
30fn initial_and_temporal(@builtin(global_invocation_id) global_id: vec3<u32>) {
31    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
32
33    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
34    var rng = pixel_index + constants.frame_index;
35
36    let depth = textureLoad(depth_buffer, global_id.xy, 0);
37    if depth == 0.0 {
38        gi_reservoirs_b[pixel_index] = empty_reservoir();
39        return;
40    }
41    let gpixel = textureLoad(gbuffer, global_id.xy, 0);
42    let world_position = reconstruct_world_position(global_id.xy, depth);
43    let world_normal = octahedral_decode(unpack_24bit_normal(gpixel.a));
44    let base_color = pow(unpack4x8unorm(gpixel.r).rgb, vec3(2.2));
45    let diffuse_brdf = base_color / PI;
46
47    let initial_reservoir = generate_initial_reservoir(world_position, world_normal, &rng);
48    let temporal = load_temporal_reservoir(global_id.xy, depth, world_position, world_normal);
49    let merge_result = merge_reservoirs(initial_reservoir, world_position, world_normal, diffuse_brdf,
50        temporal.reservoir, temporal.world_position, temporal.world_normal, temporal.diffuse_brdf, &rng);
51
52    gi_reservoirs_b[pixel_index] = merge_result.merged_reservoir;
53}
54
55@compute @workgroup_size(8, 8, 1)
56fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
57    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
58
59    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
60    var rng = pixel_index + constants.frame_index;
61
62    let depth = textureLoad(depth_buffer, global_id.xy, 0);
63    if depth == 0.0 {
64        gi_reservoirs_a[pixel_index] = empty_reservoir();
65        return;
66    }
67    let gpixel = textureLoad(gbuffer, global_id.xy, 0);
68    let world_position = reconstruct_world_position(global_id.xy, depth);
69    let world_normal = octahedral_decode(unpack_24bit_normal(gpixel.a));
70    let base_color = pow(unpack4x8unorm(gpixel.r).rgb, vec3(2.2));
71    let diffuse_brdf = base_color / PI;
72
73    let input_reservoir = gi_reservoirs_b[pixel_index];
74    let spatial = load_spatial_reservoir(global_id.xy, depth, world_position, world_normal, &rng);
75    let merge_result = merge_reservoirs(input_reservoir, world_position, world_normal, diffuse_brdf,
76        spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.diffuse_brdf, &rng);
77    let combined_reservoir = merge_result.merged_reservoir;
78
79    gi_reservoirs_a[pixel_index] = combined_reservoir;
80
81    var pixel_color = textureLoad(view_output, global_id.xy);
82    pixel_color += vec4(merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight * view.exposure, 0.0);
83    textureStore(view_output, global_id.xy, pixel_color);
84
85#ifdef VISUALIZE_WORLD_CACHE
86    textureStore(view_output, global_id.xy, vec4(query_world_cache(world_position, world_normal, view.world_position) * view.exposure, 1.0));
87#endif
88}
89
90fn generate_initial_reservoir(world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> Reservoir {
91    var reservoir = empty_reservoir();
92
93    let ray_direction = sample_uniform_hemisphere(world_normal, rng);
94    let ray_hit = trace_ray(world_position, ray_direction, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE);
95
96    if ray_hit.kind == RAY_QUERY_INTERSECTION_NONE {
97        return reservoir;
98    }
99
100    let sample_point = resolve_ray_hit_full(ray_hit);
101
102    if all(sample_point.material.emissive != vec3(0.0)) {
103        return reservoir;
104    }
105
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;
109
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;
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);
116    reservoir.unbiased_contribution_weight = uniform_hemisphere_inverse_pdf();
117#endif
118
119    let sample_point_diffuse_brdf = sample_point.material.base_color / PI;
120    reservoir.radiance *= sample_point_diffuse_brdf;
121
122    return reservoir;
123}
124
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
129    // Check if the current pixel was off screen during the previous frame (current pixel is newly visible),
130    // or if all temporal history should assumed to be invalid
131    if any(temporal_pixel_id_float < vec2(0.0)) || any(temporal_pixel_id_float >= view.main_pass_viewport.zw) || bool(constants.reset) {
132        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
133    }
134
135    let temporal_pixel_id_base = vec2<u32>(round(temporal_pixel_id_float));
136    for (var i = 0u; i < 4u; i++) {
137        let temporal_pixel_id = permute_pixel(temporal_pixel_id_base, i);
138
139        // Check if the pixel features have changed heavily between the current and previous frame
140        let temporal_depth = textureLoad(previous_depth_buffer, temporal_pixel_id, 0);
141        let temporal_gpixel = textureLoad(previous_gbuffer, temporal_pixel_id, 0);
142        let temporal_world_position = reconstruct_previous_world_position(temporal_pixel_id, temporal_depth);
143        let temporal_world_normal = octahedral_decode(unpack_24bit_normal(temporal_gpixel.a));
144        let temporal_base_color = pow(unpack4x8unorm(temporal_gpixel.r).rgb, vec3(2.2));
145        let temporal_diffuse_brdf = temporal_base_color / PI;
146        if pixel_dissimilar(depth, world_position, temporal_world_position, world_normal, temporal_world_normal) {
147            continue;
148        }
149
150        let temporal_pixel_index = temporal_pixel_id.x + temporal_pixel_id.y * u32(view.main_pass_viewport.z);
151        var temporal_reservoir = gi_reservoirs_a[temporal_pixel_index];
152
153        temporal_reservoir.confidence_weight = min(temporal_reservoir.confidence_weight, CONFIDENCE_WEIGHT_CAP);
154
155        return NeighborInfo(temporal_reservoir, temporal_world_position, temporal_world_normal, temporal_diffuse_brdf);
156    }
157
158    return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), vec3(0.0));
159}
160
161fn permute_pixel(pixel_id: vec2<u32>, i: u32) -> vec2<u32> {
162    let r = constants.frame_index + i;
163    let offset = vec2(r & 3u, (r >> 2u) & 3u);
164    var shifted_pixel_id = pixel_id + offset;
165    shifted_pixel_id ^= vec2(3u);
166    shifted_pixel_id -= offset;
167    return min(shifted_pixel_id, vec2<u32>(view.main_pass_viewport.zw - 1.0));
168}
169
170fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> NeighborInfo {
171    let spatial_pixel_id = get_neighbor_pixel_id(pixel_id, rng);
172
173    let spatial_depth = textureLoad(depth_buffer, spatial_pixel_id, 0);
174    let spatial_gpixel = textureLoad(gbuffer, spatial_pixel_id, 0);
175    let spatial_world_position = reconstruct_world_position(spatial_pixel_id, spatial_depth);
176    let spatial_world_normal = octahedral_decode(unpack_24bit_normal(spatial_gpixel.a));
177    let spatial_base_color = pow(unpack4x8unorm(spatial_gpixel.r).rgb, vec3(2.2));
178    let spatial_diffuse_brdf = spatial_base_color / PI;
179    if pixel_dissimilar(depth, world_position, spatial_world_position, world_normal, spatial_world_normal) {
180        return NeighborInfo(empty_reservoir(), spatial_world_position, spatial_world_normal, spatial_diffuse_brdf);
181    }
182
183    let spatial_pixel_index = spatial_pixel_id.x + spatial_pixel_id.y * u32(view.main_pass_viewport.z);
184    var spatial_reservoir = gi_reservoirs_b[spatial_pixel_index];
185
186    spatial_reservoir.radiance *= trace_point_visibility(world_position, spatial_reservoir.sample_point_world_position);
187
188    return NeighborInfo(spatial_reservoir, spatial_world_position, spatial_world_normal, spatial_diffuse_brdf);
189}
190
191fn get_neighbor_pixel_id(center_pixel_id: vec2<u32>, rng: ptr<function, u32>) -> vec2<u32> {
192    var spatial_id = vec2<f32>(center_pixel_id) + sample_disk(SPATIAL_REUSE_RADIUS_PIXELS, rng);
193    spatial_id = clamp(spatial_id, vec2(0.0), view.main_pass_viewport.zw - 1.0);
194    return vec2<u32>(spatial_id);
195}
196
197struct NeighborInfo {
198    reservoir: Reservoir,
199    world_position: vec3<f32>,
200    world_normal: vec3<f32>,
201    diffuse_brdf: vec3<f32>,
202}
203
204fn jacobian(
205    new_world_position: vec3<f32>,
206    original_world_position: vec3<f32>,
207    sample_point_world_position: vec3<f32>,
208    sample_point_world_normal: vec3<f32>,
209) -> f32 {
210    let r = new_world_position - sample_point_world_position;
211    let q = original_world_position - sample_point_world_position;
212    let rl = length(r);
213    let ql = length(q);
214    let phi_r = saturate(dot(r / rl, sample_point_world_normal));
215    let phi_q = saturate(dot(q / ql, sample_point_world_normal));
216    let jacobian = (phi_r * ql * ql) / (phi_q * rl * rl);
217    return select(jacobian, 0.0, isinf(jacobian) || isnan(jacobian));
218}
219
220fn isinf(x: f32) -> bool {
221    return (bitcast<u32>(x) & 0x7fffffffu) == 0x7f800000u;
222}
223
224fn isnan(x: f32) -> bool {
225    return (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
226}
227
228fn reconstruct_world_position(pixel_id: vec2<u32>, depth: f32) -> vec3<f32> {
229    let uv = (vec2<f32>(pixel_id) + 0.5) / view.main_pass_viewport.zw;
230    let xy_ndc = (uv - vec2(0.5)) * vec2(2.0, -2.0);
231    let world_pos = view.world_from_clip * vec4(xy_ndc, depth, 1.0);
232    return world_pos.xyz / world_pos.w;
233}
234
235fn reconstruct_previous_world_position(pixel_id: vec2<u32>, depth: f32) -> vec3<f32> {
236    let uv = (vec2<f32>(pixel_id) + 0.5) / view.main_pass_viewport.zw;
237    let xy_ndc = (uv - vec2(0.5)) * vec2(2.0, -2.0);
238    let world_pos = previous_view.world_from_clip * vec4(xy_ndc, depth, 1.0);
239    return world_pos.xyz / world_pos.w;
240}
241
242// Reject if tangent plane difference difference more than 0.3% or angle between normals more than 25 degrees
243fn pixel_dissimilar(depth: f32, world_position: vec3<f32>, other_world_position: vec3<f32>, normal: vec3<f32>, other_normal: vec3<f32>) -> bool {
244    // https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s22699-fast-denoising-with-self-stabilizing-recurrent-blurs.pdf#page=45
245    let tangent_plane_distance = abs(dot(normal, other_world_position - world_position));
246    let view_z = -depth_ndc_to_view_z(depth);
247
248    return tangent_plane_distance / view_z > 0.003 || dot(normal, other_normal) < 0.906;
249}
250
251fn depth_ndc_to_view_z(ndc_depth: f32) -> f32 {
252#ifdef VIEW_PROJECTION_PERSPECTIVE
253    return -view.clip_from_view[3][2]() / ndc_depth;
254#else ifdef VIEW_PROJECTION_ORTHOGRAPHIC
255    return -(view.clip_from_view[3][2] - ndc_depth) / view.clip_from_view[2][2];
256#else
257    let view_pos = view.view_from_clip * vec4(0.0, 0.0, ndc_depth, 1.0);
258    return view_pos.z / view_pos.w;
259#endif
260}
261
262// Don't adjust the size of this struct without also adjusting GI_RESERVOIR_STRUCT_SIZE.
263struct Reservoir {
264    sample_point_world_position: vec3<f32>,
265    weight_sum: f32,
266    radiance: vec3<f32>,
267    confidence_weight: f32,
268    sample_point_world_normal: vec3<f32>,
269    unbiased_contribution_weight: f32,
270}
271
272fn empty_reservoir() -> Reservoir {
273    return Reservoir(
274        vec3(0.0),
275        0.0,
276        vec3(0.0),
277        0.0,
278        vec3(0.0),
279        0.0,
280    );
281}
282
283struct ReservoirMergeResult {
284    merged_reservoir: Reservoir,
285    selected_sample_radiance: vec3<f32>,
286}
287
288fn merge_reservoirs(
289    canonical_reservoir: Reservoir,
290    canonical_world_position: vec3<f32>,
291    canonical_world_normal: vec3<f32>,
292    canonical_diffuse_brdf: vec3<f32>,
293    other_reservoir: Reservoir,
294    other_world_position: vec3<f32>,
295    other_world_normal: vec3<f32>,
296    other_diffuse_brdf: vec3<f32>,
297    rng: ptr<function, u32>,
298) -> ReservoirMergeResult {
299    // Radiances for resampling
300    let canonical_sample_radiance =
301        canonical_reservoir.radiance *
302        saturate(dot(normalize(canonical_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal)) *
303        canonical_diffuse_brdf;
304    let other_sample_radiance =
305        other_reservoir.radiance *
306        saturate(dot(normalize(other_reservoir.sample_point_world_position - canonical_world_position), canonical_world_normal)) *
307        canonical_diffuse_brdf;
308
309    // Target functions for resampling and MIS
310    let canonical_target_function_canonical_sample = luminance(canonical_sample_radiance);
311    let canonical_target_function_other_sample = luminance(other_sample_radiance);
312
313    // Extra target functions for MIS
314    let other_target_function_canonical_sample = luminance(
315        canonical_reservoir.radiance *
316        saturate(dot(normalize(canonical_reservoir.sample_point_world_position - other_world_position), other_world_normal)) *
317        other_diffuse_brdf
318    );
319    let other_target_function_other_sample = luminance(
320        other_reservoir.radiance *
321        saturate(dot(normalize(other_reservoir.sample_point_world_position - other_world_position), other_world_normal)) *
322        other_diffuse_brdf
323    );
324
325    // Jacobians for resampling and MIS
326    let canonical_target_function_other_sample_jacobian = jacobian(
327        canonical_world_position,
328        other_world_position,
329        other_reservoir.sample_point_world_position,
330        other_reservoir.sample_point_world_normal
331    );
332    let other_target_function_canonical_sample_jacobian = jacobian(
333        other_world_position,
334        canonical_world_position,
335        canonical_reservoir.sample_point_world_position,
336        canonical_reservoir.sample_point_world_normal
337    );
338
339    // Don't merge samples with huge jacobians, as it explodes the variance
340    if canonical_target_function_other_sample_jacobian > 2.0 {
341        return ReservoirMergeResult(canonical_reservoir, canonical_sample_radiance);
342    }
343
344    // Resampling weight for canonical sample
345    let canonical_sample_mis_weight = balance_heuristic(
346        canonical_reservoir.confidence_weight * canonical_target_function_canonical_sample,
347        other_reservoir.confidence_weight * other_target_function_canonical_sample * other_target_function_canonical_sample_jacobian,
348    );
349    let canonical_sample_resampling_weight = canonical_sample_mis_weight *
350        canonical_target_function_canonical_sample *
351        canonical_reservoir.unbiased_contribution_weight;
352
353    // Resampling weight for other sample
354    let other_sample_mis_weight = balance_heuristic(
355        other_reservoir.confidence_weight * other_target_function_other_sample,
356        canonical_reservoir.confidence_weight * canonical_target_function_other_sample * canonical_target_function_other_sample_jacobian,
357    );
358    let other_sample_resampling_weight = other_sample_mis_weight *
359        canonical_target_function_other_sample *
360        other_reservoir.unbiased_contribution_weight *
361        canonical_target_function_other_sample_jacobian;
362
363    // Perform resampling
364    var combined_reservoir = empty_reservoir();
365    combined_reservoir.confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
366    combined_reservoir.weight_sum = canonical_sample_resampling_weight + other_sample_resampling_weight;
367
368    if rand_f(rng) < other_sample_resampling_weight / combined_reservoir.weight_sum {
369        combined_reservoir.sample_point_world_position = other_reservoir.sample_point_world_position;
370        combined_reservoir.sample_point_world_normal = other_reservoir.sample_point_world_normal;
371        combined_reservoir.radiance = other_reservoir.radiance;
372
373        let inverse_target_function = select(0.0, 1.0 / canonical_target_function_other_sample, canonical_target_function_other_sample > 0.0);
374        combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
375
376        return ReservoirMergeResult(combined_reservoir, other_sample_radiance);
377    } else {
378        combined_reservoir.sample_point_world_position = canonical_reservoir.sample_point_world_position;
379        combined_reservoir.sample_point_world_normal = canonical_reservoir.sample_point_world_normal;
380        combined_reservoir.radiance = canonical_reservoir.radiance;
381
382        let inverse_target_function = select(0.0, 1.0 / canonical_target_function_canonical_sample, canonical_target_function_canonical_sample > 0.0);
383        combined_reservoir.unbiased_contribution_weight = combined_reservoir.weight_sum * inverse_target_function;
384
385        return ReservoirMergeResult(combined_reservoir, canonical_sample_radiance);
386    }
387}
388
389fn balance_heuristic(x: f32, y: f32) -> f32 {
390    let sum = x + y;
391    if sum == 0.0 {
392        return 0.0;
393    }
394    return x / sum;
395}