1enable wgpu_ray_query;
2
3#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
4#import bevy_pbr::utils::{rand_f, rand_range_u, sample_cosine_hemisphere}
5#import bevy_render::view::View
6#import bevy_solari::presample_light_tiles::{ResolvedLightSamplePacked, unpack_resolved_light_sample}
7#import bevy_solari::sampling::{calculate_resolved_light_contribution, trace_light_visibility}
8#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN}
9#import bevy_solari::world_cache::{
10    WORLD_CACHE_MAX_TEMPORAL_SAMPLES,
11    WORLD_CACHE_DIRECT_LIGHT_SAMPLE_COUNT,
12    WORLD_CACHE_MAX_GI_RAY_DISTANCE,
13    WORLD_CACHE_CELL_UPDATES_SOFT_CAP,
14    query_world_cache,
15}
16#import bevy_solari::realtime_bindings::{
17    light_tile_resolved_samples,
18    view,
19    constants,
20    world_cache_active_cells_count,
21    world_cache_active_cell_indices,
22    world_cache_life,
23    world_cache_geometry_data,
24    world_cache_radiance,
25    world_cache_luminance_deltas,
26    world_cache_active_cells_new_radiance,
27}
28
29@compute @workgroup_size(64, 1, 1)
30fn sample_di(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) active_cell_id: vec3<u32>) {
31    if active_cell_id.x >= world_cache_active_cells_count { return; }
32
33    let cell_index = world_cache_active_cell_indices[active_cell_id.x];
34    let geometry_data = world_cache_geometry_data[cell_index];
35    var rng = cell_index + constants.frame_index;
36
37    if rand_f(&rng) >= f32(WORLD_CACHE_CELL_UPDATES_SOFT_CAP) / f32(world_cache_active_cells_count) { return; }
38
39    let new_radiance = sample_random_light_ris(geometry_data.world_position, geometry_data.world_normal, workgroup_id.xy, &rng);
40
41    world_cache_active_cells_new_radiance[active_cell_id.x] = new_radiance;
42}
43
44@compute @workgroup_size(64, 1, 1)
45fn sample_gi(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) active_cell_id: vec3<u32>) {
46    if active_cell_id.x >= world_cache_active_cells_count { return; }
47
48    let cell_index = world_cache_active_cell_indices[active_cell_id.x];
49    let geometry_data = world_cache_geometry_data[cell_index];
50    var rng = cell_index + constants.frame_index;
51
52    if rand_f(&rng) >= f32(WORLD_CACHE_CELL_UPDATES_SOFT_CAP) / f32(world_cache_active_cells_count) { return; }
53
54    let ray_direction = sample_cosine_hemisphere(geometry_data.world_normal, &rng);
55    let ray = trace_ray(geometry_data.world_position + (geometry_data.world_normal * RAY_T_MIN), ray_direction, RAY_T_MIN, WORLD_CACHE_MAX_GI_RAY_DISTANCE, RAY_FLAG_NONE);
56    if ray.kind != RAY_QUERY_INTERSECTION_NONE {
57        let ray_hit = resolve_ray_hit_full(ray);
58        let cell_life = atomicLoad(&world_cache_life[cell_index]);
59        let radiance = query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, ray.t, cell_life, &rng);
60        world_cache_active_cells_new_radiance[active_cell_id.x] += ray_hit.material.base_color * radiance;
61    }
62}
63
64@compute @workgroup_size(64, 1, 1)
65fn blend_new_samples(@builtin(global_invocation_id) active_cell_id: vec3<u32>) {
66    if active_cell_id.x >= world_cache_active_cells_count { return; }
67
68    let cell_index = world_cache_active_cell_indices[active_cell_id.x];
69    var rng = cell_index + constants.frame_index;
70
71    if rand_f(&rng) >= f32(WORLD_CACHE_CELL_UPDATES_SOFT_CAP) / f32(world_cache_active_cells_count) { return; }
72
73    let old_radiance = world_cache_radiance[cell_index];
74    let new_radiance = world_cache_active_cells_new_radiance[active_cell_id.x];
75    let luminance_delta = world_cache_luminance_deltas[cell_index];
76
77    // https://bsky.app/profile/gboisse.bsky.social/post/3m5blga3ftk2a
78    let sample_count = min(old_radiance.a + 1.0, WORLD_CACHE_MAX_TEMPORAL_SAMPLES);
79    let alpha = abs(luminance_delta) / max(luminance(old_radiance.rgb), 0.001);
80    let max_sample_count = mix(WORLD_CACHE_MAX_TEMPORAL_SAMPLES, 1.0, pow(saturate(alpha), 1.0 / 8.0));
81    var blend_amount = 1.0 / min(sample_count, max_sample_count);
82    if bool(constants.reset) {
83        blend_amount = 1.0;
84    }
85
86    let blended_radiance = mix(old_radiance.rgb, new_radiance, blend_amount);
87    let blended_luminance_delta = select(mix(luminance_delta, luminance(blended_radiance) - luminance(old_radiance.rgb), 1.0 / 8.0), 0.0, bool(constants.reset));
88
89    world_cache_radiance[cell_index] = vec4(blended_radiance, sample_count);
90    world_cache_luminance_deltas[cell_index] = blended_luminance_delta;
91}
92
93fn sample_random_light_ris(world_position: vec3<f32>, world_normal: vec3<f32>, workgroup_id: vec2<u32>, rng: ptr<function, u32>) -> vec3<f32> {
94    var workgroup_rng = (workgroup_id.x * 5782582u) + workgroup_id.y;
95    let light_tile_start = rand_range_u(128u, &workgroup_rng) * 1024u;
96
97    var weight_sum = 0.0;
98    var selected_sample_radiance = vec3(0.0);
99    var selected_sample_target_function = 0.0;
100    var selected_sample_world_position = vec4(0.0);
101    let mis_weight = 1.0 / f32(WORLD_CACHE_DIRECT_LIGHT_SAMPLE_COUNT);
102    for (var i = 0u; i < WORLD_CACHE_DIRECT_LIGHT_SAMPLE_COUNT; i++) {
103        let tile_sample = light_tile_start + rand_range_u(1024u, rng);
104        let resolved_light_sample = unpack_resolved_light_sample(light_tile_resolved_samples[tile_sample], view.exposure);
105        let light_contribution = calculate_resolved_light_contribution(resolved_light_sample, world_position, world_normal);
106
107        let contribution = light_contribution.radiance * saturate(dot(light_contribution.wi, world_normal));
108        let target_function = luminance(contribution);
109        let resampling_weight = mis_weight * (target_function * light_contribution.inverse_pdf);
110
111        weight_sum += resampling_weight;
112
113        if rand_f(rng) < resampling_weight / weight_sum {
114            selected_sample_radiance = contribution;
115            selected_sample_target_function = target_function;
116            selected_sample_world_position = resolved_light_sample.world_position;
117        }
118    }
119
120    var unbiased_contribution_weight = 0.0;
121    if all(selected_sample_radiance != vec3(0.0)) {
122        let inverse_target_function = select(0.0, 1.0 / selected_sample_target_function, selected_sample_target_function > 0.0);
123        unbiased_contribution_weight = weight_sum * inverse_target_function;
124
125        unbiased_contribution_weight *= trace_light_visibility(world_position + (world_normal * RAY_T_MIN), selected_sample_world_position);
126    }
127
128    return selected_sample_radiance * unbiased_contribution_weight;
129}