1//! https://intro-to-restir.cwyman.org/presentations/2023ReSTIR_Course_Notes.pdf
2
3import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance;
4import bevy_pbr::render::utils::{rand_f, rand_u, sample_disk};
5import bevy_render::utils::octahedral_decode;
6import package::scene::brdf::{brdf_pdf, evaluate_brdf, F_AB};
7import package::realtime::gbuffer_utils::{gpixel_resolve, permute_pixel, pixel_dissimilar};
8import package::realtime::initial_path::{generate_initial_path, InitialSamplingResult};
9import package::realtime::bindings::{depth_buffer, empty_reservoir, gbuffer, load_motion_vector, previous_depth_buffer, previous_gbuffer, previous_view, reservoirs_a, reservoirs_b, Reservoir, constants, view, view_output};
10import package::scene::sampling::{balance_heuristic, calculate_resolved_light_contribution, isinf, isnan, LightSample, NULL_LIGHT_ID, power_heuristic, resolve_light_sample, ResolvedLightSample, trace_visibility, trace_visibility_previous_frame};
11import package::scene::bindings::{light_sources, LIGHT_NOT_PRESENT_THIS_FRAME, previous_frame_light_id_translations, RAY_T_MAX, RAY_T_MIN, ResolvedMaterial};
12import package::realtime::world_cache_query::{query_world_cache, WORLD_CACHE_CELL_LIFETIME};
13
14enable wgpu_ray_query;
15
16const SPATIAL_REUSE_RADIUS_PIXELS = 30.0;
17
18@compute @workgroup_size(8, 8, 1)
19fn initial_and_temporal(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) global_id: vec3<u32>) {
20    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
21
22    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
23    var rng = pixel_index + constants.frame_rng;
24
25    let depth = textureLoad(depth_buffer, global_id.xy, 0);
26    if depth == 0.0 {
27        reservoirs_b[pixel_index] = empty_reservoir();
28        return;
29    }
30    let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
31
32    let initial = generate_initial_path(surface.world_position, surface.world_normal, surface.material, workgroup_id.xy, global_id.xy, &rng);
33    textureStore(view_output, global_id.xy, vec4(initial.radiance, 0.0));
34
35    let temporal = load_temporal_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal);
36    let previous_camera_homogeneous = previous_view.world_from_clip * (previous_view.clip_from_view * vec4(0.0, 0.0, 0.0, 1.0));
37    let previous_camera_world_position = previous_camera_homogeneous.xyz / previous_camera_homogeneous.w;
38    let merge_result = merge_reservoirs(initial.reservoir, surface.world_position, surface.world_normal, surface.material,
39        temporal.reservoir, temporal.world_position, temporal.world_normal, temporal.material, previous_camera_world_position, &rng);
40
41    reservoirs_b[pixel_index] = merge_result.merged_reservoir;
42}
43
44@compute @workgroup_size(8, 8, 1)
45fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
46    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
47
48    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
49    var rng = pixel_index + constants.frame_rng + 0x6A09E667u;
50
51    let depth = textureLoad(depth_buffer, global_id.xy, 0);
52    if depth == 0.0 {
53        reservoirs_a[pixel_index] = empty_reservoir();
54        return;
55    }
56    let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
57
58    let input_reservoir = reservoirs_b[pixel_index];
59    let spatial = load_spatial_reservoir(global_id.xy, depth, surface.world_position, surface.world_normal, &rng);
60    let merge_result = merge_reservoirs(input_reservoir, surface.world_position, surface.world_normal, surface.material,
61        spatial.reservoir, spatial.world_position, spatial.world_normal, spatial.material, view.world_position, &rng);
62
63    reservoirs_a[pixel_index] = merge_result.merged_reservoir;
64
65    var pixel_color = merge_result.selected_sample_brdf_radiance * merge_result.merged_reservoir.unbiased_contribution_weight;
66    pixel_color += surface.material.emissive;
67    pixel_color += textureLoad(view_output, global_id.xy).rgb;
68    pixel_color *= view.exposure;
69    textureStore(view_output, global_id.xy, vec4(pixel_color, 1.0));
70
71    @if(VISUALIZE_WORLD_CACHE)
72    textureStore(view_output, global_id.xy, vec4(query_world_cache(surface.world_position, surface.world_normal, view.world_position, RAY_T_MAX, WORLD_CACHE_CELL_LIFETIME, &rng) * view.exposure, 1.0));
73}
74
75fn load_temporal_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>) -> NeighborInfo {
76    if bool(constants.reset) {
77        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), empty_material());
78    }
79
80    let motion_vector = load_motion_vector(pixel_id);
81    let temporal_pixel_id_float = round(vec2<f32>(pixel_id) - (motion_vector * view.main_pass_viewport.zw));
82
83    // If reprojection lands off-screen, fall back to this pixel's own previous reservoir rather than
84    // dropping history. The dissimilarity check below still validates the surface, and a same-pixel
85    // guess that passes it beats restarting from a confidence-1 reservoir at the screen edge.
86    var point_temporal_pixel_id = pixel_id;
87    if all(temporal_pixel_id_float >= vec2(0.0)) && all(temporal_pixel_id_float < view.main_pass_viewport.zw) {
88        point_temporal_pixel_id = vec2<u32>(temporal_pixel_id_float);
89    }
90
91    var permute_rng = constants.frame_rng;
92    let permuted_temporal_pixel_id = permute_pixel(point_temporal_pixel_id, rand_u(&permute_rng), view.main_pass_viewport.zw);
93
94    // Check if the pixel features have changed heavily between the current and previous frame
95    let temporal_depth = textureLoad(previous_depth_buffer, permuted_temporal_pixel_id, 0);
96    let temporal_surface = gpixel_resolve(textureLoad(previous_gbuffer, permuted_temporal_pixel_id, 0), temporal_depth, permuted_temporal_pixel_id, view.main_pass_viewport.zw, previous_view.world_from_clip);
97    if pixel_dissimilar(depth, world_position, temporal_surface.world_position, world_normal, temporal_surface.world_normal, view) {
98        return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), empty_material());
99    }
100
101    let temporal_pixel_index = permuted_temporal_pixel_id.x + permuted_temporal_pixel_id.y * u32(view.main_pass_viewport.z);
102    var temporal = NeighborInfo(reservoirs_a[temporal_pixel_index], temporal_surface.world_position, temporal_surface.world_normal, temporal_surface.material);
103
104    // Check if the light selected in the previous frame no longer exists in the current frame (e.g. entity despawned)
105    if temporal.reservoir.light_sample.light_id != NULL_LIGHT_ID {
106        let previous_light_id = temporal.reservoir.light_sample.light_id >> 16u;
107        let triangle_id = temporal.reservoir.light_sample.light_id & 0xFFFFu;
108        let light_id = previous_frame_light_id_translations[previous_light_id];
109        if light_id == LIGHT_NOT_PRESENT_THIS_FRAME {
110            return NeighborInfo(empty_reservoir(), vec3(0.0), vec3(0.0), empty_material());
111        }
112        temporal.reservoir.light_sample.light_id = (light_id << 16u) | triangle_id;
113    }
114
115    temporal.reservoir.confidence_weight = min(temporal.reservoir.confidence_weight, constants.confidence_weight_cap);
116
117    return temporal;
118}
119
120fn load_spatial_reservoir(pixel_id: vec2<u32>, depth: f32, world_position: vec3<f32>, world_normal: vec3<f32>, rng: ptr<function, u32>) -> NeighborInfo {
121    for (var i = 0u; i < 5u; i++) {
122        let spatial_pixel_id = get_neighbor_pixel_id(pixel_id, SPATIAL_REUSE_RADIUS_PIXELS, rng);
123
124        if all(spatial_pixel_id == pixel_id) {
125            continue;
126        }
127
128        let spatial_depth = textureLoad(depth_buffer, spatial_pixel_id, 0);
129        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);
130        if pixel_dissimilar(depth, world_position, spatial_surface.world_position, world_normal, spatial_surface.world_normal, view) {
131            continue;
132        }
133
134        let spatial_pixel_index = spatial_pixel_id.x + spatial_pixel_id.y * u32(view.main_pass_viewport.z);
135        return NeighborInfo(reservoirs_b[spatial_pixel_index], spatial_surface.world_position, spatial_surface.world_normal, spatial_surface.material);
136    }
137
138    return NeighborInfo(empty_reservoir(), world_position, world_normal, empty_material());
139}
140
141fn get_neighbor_pixel_id(center_pixel_id: vec2<u32>, search_radius: f32, rng: ptr<function, u32>) -> vec2<u32> {
142    var spatial_id = vec2<f32>(center_pixel_id) + sample_disk(search_radius, rng);
143    spatial_id = clamp(spatial_id, vec2(0.0), view.main_pass_viewport.zw - 1.0);
144    return vec2<u32>(spatial_id);
145}
146
147struct NeighborInfo {
148    reservoir: Reservoir,
149    world_position: vec3<f32>,
150    world_normal: vec3<f32>,
151    material: ResolvedMaterial,
152}
153
154fn empty_material() -> ResolvedMaterial {
155    return ResolvedMaterial(vec3(0.0), vec3(0.0), 0.0, 0.0, 0.0, 0.0);
156}
157
158fn jacobian(
159    new_world_position: vec3<f32>,
160    original_world_position: vec3<f32>,
161    sample_point_world_position: vec3<f32>,
162    sample_point_world_normal: vec3<f32>,
163) -> f32 {
164    let r = new_world_position - sample_point_world_position;
165    let q = original_world_position - sample_point_world_position;
166    let rl = length(r);
167    let ql = length(q);
168    let phi_r = saturate(dot(r / rl, sample_point_world_normal));
169    let phi_q = saturate(dot(q / ql, sample_point_world_normal));
170    let jacobian = (phi_r * ql * ql) / (phi_q * rl * rl);
171    return select(jacobian, 0.0, isinf(jacobian) || isnan(jacobian));
172}
173
174struct ReservoirMergeResult {
175    merged_reservoir: Reservoir,
176    selected_sample_brdf_radiance: vec3<f32>,
177}
178
179fn merge_reservoirs(
180    canonical_reservoir: Reservoir,
181    canonical_world_position: vec3<f32>,
182    canonical_world_normal: vec3<f32>,
183    canonical_material: ResolvedMaterial,
184    other_reservoir: Reservoir,
185    other_world_position: vec3<f32>,
186    other_world_normal: vec3<f32>,
187    other_material: ResolvedMaterial,
188    other_view_position: vec3<f32>,
189    rng: ptr<function, u32>,
190) -> ReservoirMergeResult {
191    var canonical_resolved: ResolvedLightSample;
192    if canonical_reservoir.light_sample.light_id != NULL_LIGHT_ID {
193        canonical_resolved = resolve_light_sample(canonical_reservoir.light_sample, light_sources[canonical_reservoir.light_sample.light_id >> 16u]);
194    }
195
196    let canonical_wo = normalize(view.world_position - canonical_world_position);
197    let canonical_NdotV = max(dot(canonical_world_normal, canonical_wo), 0.0001);
198    let canonical_F_ab = F_AB(canonical_material.perceptual_roughness, canonical_NdotV);
199    let canonical_sample_at_canonical = reservoir_contribution(canonical_reservoir, canonical_resolved, canonical_world_position, canonical_world_normal, canonical_wo, canonical_material, canonical_F_ab);
200
201    // Skip resampling empty reservoirs
202    if other_reservoir.confidence_weight == 0.0 {
203        return ReservoirMergeResult(canonical_reservoir, canonical_sample_at_canonical.brdf_radiance);
204    }
205
206    var other_resolved: ResolvedLightSample;
207    if other_reservoir.light_sample.light_id != NULL_LIGHT_ID {
208        other_resolved = resolve_light_sample(other_reservoir.light_sample, light_sources[other_reservoir.light_sample.light_id >> 16u]);
209    }
210    let other_wo = normalize(other_view_position - other_world_position);
211    let other_NdotV = max(dot(other_world_normal, other_wo), 0.0001);
212    let other_F_ab = F_AB(other_material.perceptual_roughness, other_NdotV);
213
214    // Contributions for resampling and MIS
215    var other_sample_at_canonical = reservoir_contribution(other_reservoir, other_resolved, canonical_world_position, canonical_world_normal, canonical_wo, canonical_material, canonical_F_ab);
216    var canonical_sample_at_other = reservoir_contribution(canonical_reservoir, canonical_resolved, other_world_position, other_world_normal, other_wo, other_material, other_F_ab);
217    let other_sample_at_other = reservoir_contribution(other_reservoir, other_resolved, other_world_position, other_world_normal, other_wo, other_material, other_F_ab);
218
219    // Jacobians for resampling and MIS. Light samples don't need a reprojection jacobian,
220    // since calculate_resolved_light_contribution already accounts for the shading point's geometry.
221    var other_sample_at_canonical_jacobian = 1.0;
222    if other_reservoir.light_sample.light_id == NULL_LIGHT_ID {
223        other_sample_at_canonical_jacobian = jacobian(
224            canonical_world_position,
225            other_world_position,
226            other_reservoir.sample_point_world_position,
227            octahedral_decode(other_reservoir.sample_point_world_normal)
228        );
229    }
230    var canonical_sample_at_other_jacobian = 1.0;
231    if canonical_reservoir.light_sample.light_id == NULL_LIGHT_ID {
232        canonical_sample_at_other_jacobian = jacobian(
233            other_world_position,
234            canonical_world_position,
235            canonical_reservoir.sample_point_world_position,
236            octahedral_decode(canonical_reservoir.sample_point_world_normal)
237        );
238    }
239
240    // Don't merge samples with huge jacobians, as it explodes the variance
241    if other_sample_at_canonical_jacobian < 0.125 || other_sample_at_canonical_jacobian > 8.0 {
242        other_sample_at_canonical_jacobian = 0.0;
243    }
244    if canonical_sample_at_other_jacobian < 0.125 || canonical_sample_at_other_jacobian > 8.0 {
245        canonical_sample_at_other_jacobian = 0.0;
246    }
247
248    // Visibility for the cross-domain targets
249    if other_sample_at_canonical.target_function > 0.0 && other_sample_at_canonical_jacobian > 0.0 {
250        let visibility = trace_visibility(canonical_world_position + canonical_world_normal * RAY_T_MIN, other_sample_at_canonical.sample_world_position);
251        other_sample_at_canonical.target_function *= visibility;
252    }
253    if canonical_sample_at_other.target_function > 0.0 && canonical_sample_at_other_jacobian > 0.0 {
254        @if(SPATIAL_MERGE)
255        let visibility = trace_visibility(other_world_position + other_world_normal * RAY_T_MIN, canonical_sample_at_other.sample_world_position);
256        @else
257        let visibility = trace_visibility_previous_frame(other_world_position + other_world_normal * RAY_T_MIN, canonical_sample_at_other.sample_world_position);
258        canonical_sample_at_other.target_function *= visibility;
259    }
260
261    // Defensive balance heuristic MIS (for spatial reuse only)
262    let total_confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
263    var defensive_t_c = 0.0;
264    @if(SPATIAL_MERGE)
265    defensive_t_c = select(1.0, canonical_reservoir.confidence_weight / total_confidence_weight, total_confidence_weight > 0.0);
266
267    // Resampling weight for canonical sample
268    let canonical_balance_mis_weight = balance_heuristic(
269        canonical_reservoir.confidence_weight * canonical_sample_at_canonical.target_function,
270        other_reservoir.confidence_weight * canonical_sample_at_other.target_function * canonical_sample_at_other_jacobian,
271    );
272    let canonical_sample_mis_weight = mix(canonical_balance_mis_weight, 1.0, defensive_t_c);
273    let canonical_sample_resampling_weight = canonical_sample_mis_weight * canonical_sample_at_canonical.target_function * canonical_reservoir.unbiased_contribution_weight;
274
275    // Resampling weight for other sample
276    let other_balance_mis_weight = balance_heuristic(
277        other_reservoir.confidence_weight * other_sample_at_other.target_function,
278        canonical_reservoir.confidence_weight * other_sample_at_canonical.target_function * other_sample_at_canonical_jacobian,
279    );
280    let other_sample_mis_weight = mix(other_balance_mis_weight, 0.0, defensive_t_c);
281    let other_sample_resampling_weight = other_sample_mis_weight * other_sample_at_canonical.target_function * other_reservoir.unbiased_contribution_weight * other_sample_at_canonical_jacobian;
282
283    // Perform resampling
284    var combined_reservoir = empty_reservoir();
285    combined_reservoir.confidence_weight = canonical_reservoir.confidence_weight + other_reservoir.confidence_weight;
286    let weight_sum = canonical_sample_resampling_weight + other_sample_resampling_weight;
287
288    if rand_f(rng) * weight_sum < other_sample_resampling_weight {
289        combined_reservoir.sample_point_world_position = other_reservoir.sample_point_world_position;
290        combined_reservoir.sample_point_world_normal = other_reservoir.sample_point_world_normal;
291        combined_reservoir.radiance = other_reservoir.radiance;
292        combined_reservoir.light_sample = other_reservoir.light_sample;
293
294        let inverse_target_function = select(0.0, 1.0 / other_sample_at_canonical.target_function, other_sample_at_canonical.target_function > 0.0);
295        combined_reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function;
296
297        return ReservoirMergeResult(combined_reservoir, other_sample_at_canonical.brdf_radiance);
298    } else {
299        combined_reservoir.sample_point_world_position = canonical_reservoir.sample_point_world_position;
300        combined_reservoir.sample_point_world_normal = canonical_reservoir.sample_point_world_normal;
301        combined_reservoir.radiance = canonical_reservoir.radiance;
302        combined_reservoir.light_sample = canonical_reservoir.light_sample;
303
304        let inverse_target_function = select(0.0, 1.0 / canonical_sample_at_canonical.target_function, canonical_sample_at_canonical.target_function > 0.0);
305        combined_reservoir.unbiased_contribution_weight = weight_sum * inverse_target_function;
306
307        return ReservoirMergeResult(combined_reservoir, canonical_sample_at_canonical.brdf_radiance);
308    }
309}
310
311struct ReservoirContribution {
312    brdf_radiance: vec3<f32>,
313    target_function: f32,
314    sample_world_position: vec4<f32>,
315}
316
317fn reservoir_contribution(reservoir: Reservoir, resolved: ResolvedLightSample, world_position: vec3<f32>, world_normal: vec3<f32>, wo: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> ReservoirContribution {
318    if reservoir.light_sample.light_id != NULL_LIGHT_ID {
319        let light_contribution = calculate_resolved_light_contribution(resolved, world_position, world_normal);
320
321        // MIS weight against the bounce-0 BRDF-emissive strategy, recomputed from this surface's
322        // brdf and material rather than baked into the unbiased contribution weight at generation. Mirrors the bounce-0
323        // nee_mis_weight in generate_nee_candidate and generate_emissive_candidate, which puts the same factor in the target.
324        var nee_mis_weight = 1.0;
325        if light_contribution.brdf_rays_can_hit && light_contribution.inverse_solid_angle_pdf > 0.0 {
326            let light_count = arrayLength(&light_sources);
327            let inverse_solid_angle_pdf = light_contribution.inverse_solid_angle_pdf * f32(light_count);
328            let p_nee = mix(1.0, material.perceptual_roughness, material.metallic);
329            let p_nee_strategy = f32(constants.primary_di_samples) * (1.0 / inverse_solid_angle_pdf) * p_nee;
330            let p_brdf_at_nee = brdf_pdf(wo, light_contribution.wi, world_normal, material, F_ab);
331            nee_mis_weight = power_heuristic(p_nee_strategy, p_brdf_at_nee);
332        }
333
334        let brdf_radiance = light_contribution.radiance * evaluate_brdf(wo, light_contribution.wi, world_normal, material, F_ab) * nee_mis_weight;
335        return ReservoirContribution(brdf_radiance, luminance(brdf_radiance), resolved.world_position);
336    } else if any(reservoir.radiance != vec3(0.0)) {
337        let delta = reservoir.sample_point_world_position - (world_position + world_normal * RAY_T_MIN);
338        let sample_distance = length(delta);
339        let wi = delta / sample_distance;
340        var brdf_radiance = reservoir.radiance * evaluate_brdf(wo, wi, world_normal, material, F_ab);
341
342        // Bounce-0 BRDF-emissive sample (directly-visible light). The seed field carries the light
343        // triangle's bitcast area pdf, and the stored radiance is the raw emission. Rebuild the MIS
344        // weight against this surface's NEE strategy, the dual of nee_mis_weight above and a mirror
345        // of the emissive candidate in generate_initial_path.
346        if reservoir.light_sample.seed != 0u {
347            let area_pdf = bitcast<f32>(reservoir.light_sample.seed);
348            let light_normal = octahedral_decode(reservoir.sample_point_world_normal);
349            let cos_theta_light = dot(-wi, light_normal);
350            if cos_theta_light <= 0.0 {
351                brdf_radiance = vec3(0.0);
352            } else {
353                let p_light = area_pdf * sample_distance * sample_distance / cos_theta_light;
354                let p_nee = mix(1.0, material.perceptual_roughness, material.metallic);
355                let p_brdf = brdf_pdf(wo, wi, world_normal, material, F_ab);
356                brdf_radiance *= power_heuristic(p_brdf, p_light * p_nee * f32(constants.primary_di_samples));
357            }
358        }
359
360        return ReservoirContribution(brdf_radiance, luminance(brdf_radiance), vec4(reservoir.sample_point_world_position, 1.0));
361    } else {
362        return ReservoirContribution(vec3(0.0), 0.0, vec4(reservoir.sample_point_world_position, 1.0));
363    }
364}