1#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
2#import bevy_pbr::utils::{rand_f, rand_range_u, sample_cosine_hemisphere}
3#import bevy_render::view::View
4#import bevy_solari::presample_light_tiles::{ResolvedLightSamplePacked, unpack_resolved_light_sample}
5#import bevy_solari::sampling::{calculate_resolved_light_contribution, trace_light_visibility}
6#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
7#import bevy_solari::world_cache::{
8    WORLD_CACHE_MAX_TEMPORAL_SAMPLES,
9    query_world_cache,
10    world_cache_active_cells_count,
11    world_cache_active_cell_indices,
12    world_cache_geometry_data,
13    world_cache_radiance,
14    world_cache_active_cells_new_radiance,
15}
16
17@group(1) @binding(2) var<storage, read_write> light_tile_resolved_samples: array<ResolvedLightSamplePacked>;
18@group(1) @binding(12) var<uniform> view: View;
19struct PushConstants { frame_index: u32, reset: u32 }
20var<push_constant> constants: PushConstants;
21
22const DIRECT_LIGHT_SAMPLE_COUNT: u32 = 32u;
23
24@compute @workgroup_size(1024, 1, 1)
25fn sample_radiance(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) active_cell_id: vec3<u32>) {
26    if active_cell_id.x < world_cache_active_cells_count {
27        let cell_index = world_cache_active_cell_indices[active_cell_id.x];
28        let geometry_data = world_cache_geometry_data[cell_index];
29        var rng = cell_index + constants.frame_index;
30
31        // TODO: Initialize newly active cells with data from an adjacent LOD
32
33        var new_radiance = sample_random_light_ris(geometry_data.world_position, geometry_data.world_normal, workgroup_id.xy, &rng);
34
35#ifndef NO_MULTIBOUNCE
36        let ray_direction = sample_cosine_hemisphere(geometry_data.world_normal, &rng);
37        let ray_hit = trace_ray(geometry_data.world_position, ray_direction, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE);
38        if ray_hit.kind != RAY_QUERY_INTERSECTION_NONE {
39            let ray_hit = resolve_ray_hit_full(ray_hit);
40            new_radiance += ray_hit.material.base_color * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position);
41        }
42#endif
43
44        world_cache_active_cells_new_radiance[active_cell_id.x] = new_radiance;
45    }
46}
47
48@compute @workgroup_size(1024, 1, 1)
49fn blend_new_samples(@builtin(global_invocation_id) active_cell_id: vec3<u32>) {
50    if active_cell_id.x < world_cache_active_cells_count {
51        let cell_index = world_cache_active_cell_indices[active_cell_id.x];
52
53        let old_radiance = world_cache_radiance[cell_index];
54        let new_radiance = world_cache_active_cells_new_radiance[active_cell_id.x];
55        let sample_count = min(old_radiance.a + 1.0, WORLD_CACHE_MAX_TEMPORAL_SAMPLES);
56
57        let blended_radiance = mix(old_radiance.rgb, new_radiance, 1.0 / sample_count);
58
59        world_cache_radiance[cell_index] = vec4(blended_radiance, sample_count);
60    }
61}
62
63fn sample_random_light_ris(world_position: vec3<f32>, world_normal: vec3<f32>, workgroup_id: vec2<u32>, rng: ptr<function, u32>) -> vec3<f32> {
64    var workgroup_rng = (workgroup_id.x * 5782582u) + workgroup_id.y;
65    let light_tile_start = rand_range_u(128u, &workgroup_rng) * 1024u;
66
67    var weight_sum = 0.0;
68    var selected_sample_radiance = vec3(0.0);
69    var selected_sample_target_function = 0.0;
70    var selected_sample_world_position = vec4(0.0);
71    let mis_weight = 1.0 / f32(DIRECT_LIGHT_SAMPLE_COUNT);
72    for (var i = 0u; i < DIRECT_LIGHT_SAMPLE_COUNT; i++) {
73        let tile_sample = light_tile_start + rand_range_u(1024u, rng);
74        let resolved_light_sample = unpack_resolved_light_sample(light_tile_resolved_samples[tile_sample], view.exposure);
75        let light_contribution = calculate_resolved_light_contribution(resolved_light_sample, world_position, world_normal);
76
77        let target_function = luminance(light_contribution.radiance);
78        let resampling_weight = mis_weight * (target_function * light_contribution.inverse_pdf);
79
80        weight_sum += resampling_weight;
81
82        if rand_f(rng) < resampling_weight / weight_sum {
83            selected_sample_radiance = light_contribution.radiance;
84            selected_sample_target_function = target_function;
85            selected_sample_world_position = resolved_light_sample.world_position;
86        }
87    }
88
89    var unbiased_contribution_weight = 0.0;
90    if all(selected_sample_radiance != vec3(0.0)) {
91        let inverse_target_function = select(0.0, 1.0 / selected_sample_target_function, selected_sample_target_function > 0.0);
92        unbiased_contribution_weight = weight_sum * inverse_target_function;
93
94        unbiased_contribution_weight *= trace_light_visibility(world_position, selected_sample_world_position);
95    }
96
97    return selected_sample_radiance * unbiased_contribution_weight;
98}