1import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance;
2import bevy_pbr::render::utils::{rand_f, rand_range_u, sample_cosine_hemisphere};
3import bevy_render::maths::PI;
4import package::realtime::presample_light_tiles::unpack_resolved_light_sample;
5import package::scene::sampling::{calculate_resolved_light_contribution, trace_visibility, sample_environment_map_light};
6import package::scene::bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN};
7import package::realtime::world_cache_query::query_world_cache;
8import package::realtime::bindings::{
9    light_tile_resolved_samples,
10    view,
11    constants,
12    world_cache,
13};
14
15enable wgpu_ray_query;
16
17@compute @workgroup_size(64, 1, 1)
18fn sample_di(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) active_cell_id: vec3<u32>) {
19    if active_cell_id.x >= world_cache.active_cells_count { return; }
20
21    let cell_index = world_cache.active_cell_indices[active_cell_id.x];
22    let geometry_data = world_cache.geometry_data[cell_index];
23    var rng = cell_index + constants.frame_rng;
24
25    if rand_f(&rng) >= f32(constants.world_cache_cell_updates_soft_target) / f32(world_cache.active_cells_count) { return; }
26
27    let new_radiance = sample_random_light_ris(geometry_data.world_position, geometry_data.world_normal, workgroup_id.xy, &rng);
28
29    world_cache.active_cells_new_radiance[active_cell_id.x] = new_radiance;
30}
31
32@compute @workgroup_size(64, 1, 1)
33fn sample_gi(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(global_invocation_id) active_cell_id: vec3<u32>) {
34    if active_cell_id.x >= world_cache.active_cells_count { return; }
35
36    let cell_index = world_cache.active_cell_indices[active_cell_id.x];
37    let geometry_data = world_cache.geometry_data[cell_index];
38    var rng = cell_index + constants.frame_rng;
39
40    if rand_f(&rng) >= f32(constants.world_cache_cell_updates_soft_target) / f32(world_cache.active_cells_count) { return; }
41
42    let ray_direction = sample_cosine_hemisphere(geometry_data.world_normal, &rng);
43    let ray = trace_ray(geometry_data.world_position + (geometry_data.world_normal * RAY_T_MIN), ray_direction, RAY_T_MIN, constants.world_cache_max_gi_ray_distance, RAY_FLAG_NONE);
44    if ray.kind != RAY_QUERY_INTERSECTION_NONE {
45        let ray_hit = resolve_ray_hit_full(ray);
46        let cell_life = atomicLoad(&world_cache.life[cell_index]);
47        let radiance = query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, ray.t, cell_life, &rng);
48        world_cache.active_cells_new_radiance[active_cell_id.x] += ray_hit.material.base_color * radiance;
49    } else {
50        // The ray is cut short at world_cache_max_gi_ray_distance to save on traversal costs. 
51        // A miss here may actually have been a surface further away if we kept tracing, and not actually have been a miss.
52        // We could be leaking sky light due to this if world_cache_max_gi_ray_distance is too small.
53        world_cache.active_cells_new_radiance[active_cell_id.x] += PI * sample_environment_map_light(ray_direction);
54    }
55}
56
57@compute @workgroup_size(64, 1, 1)
58fn blend_new_samples(@builtin(global_invocation_id) active_cell_id: vec3<u32>) {
59    if active_cell_id.x >= world_cache.active_cells_count { return; }
60
61    let cell_index = world_cache.active_cell_indices[active_cell_id.x];
62    var rng = cell_index + constants.frame_rng;
63
64    if rand_f(&rng) >= f32(constants.world_cache_cell_updates_soft_target) / f32(world_cache.active_cells_count) { return; }
65
66    let old_radiance = world_cache.radiance[cell_index];
67    let new_radiance = world_cache.active_cells_new_radiance[active_cell_id.x];
68    let luminance_delta = world_cache.luminance_deltas[cell_index];
69
70    // https://bsky.app/profile/gboisse.bsky.social/post/3m5blga3ftk2a
71    let sample_count = min(old_radiance.a + 1.0, constants.world_cache_max_temporal_samples);
72    let alpha = abs(luminance_delta) / max(luminance(old_radiance.rgb), 0.001);
73    let max_sample_count = mix(constants.world_cache_max_temporal_samples, 1.0, pow(saturate(alpha), 1.0 / 8.0));
74    var blend_amount = 1.0 / min(sample_count, max_sample_count);
75    if bool(constants.reset) {
76        blend_amount = 1.0;
77    }
78
79    let blended_radiance = mix(old_radiance.rgb, new_radiance, blend_amount);
80    let blended_luminance_delta = select(mix(luminance_delta, luminance(blended_radiance) - luminance(old_radiance.rgb), 1.0 / 8.0), 0.0, bool(constants.reset));
81
82    world_cache.radiance[cell_index] = vec4(blended_radiance, sample_count);
83    world_cache.luminance_deltas[cell_index] = blended_luminance_delta;
84}
85
86fn sample_random_light_ris(world_position: vec3<f32>, world_normal: vec3<f32>, workgroup_id: vec2<u32>, rng: ptr<function, u32>) -> vec3<f32> {
87    var workgroup_rng = (workgroup_id.x * 5782582u) + workgroup_id.y;
88    let light_tile_start = rand_range_u(128u, &workgroup_rng) * 1024u;
89
90    var weight_sum = 0.0;
91    var selected_sample_radiance = vec3(0.0);
92    var selected_sample_target_function = 0.0;
93    var selected_sample_world_position = vec4(0.0);
94    let mis_weight = 1.0 / f32(constants.world_cache_direct_light_sample_count);
95    for (var i = 0u; i < constants.world_cache_direct_light_sample_count; i++) {
96        let tile_sample = light_tile_start + rand_range_u(1024u, rng);
97        let resolved_light_sample = unpack_resolved_light_sample(light_tile_resolved_samples[tile_sample], view.exposure);
98        let light_contribution = calculate_resolved_light_contribution(resolved_light_sample, world_position, world_normal);
99
100        let contribution = light_contribution.radiance * saturate(dot(light_contribution.wi, world_normal));
101        let target_function = luminance(contribution);
102        let resampling_weight = mis_weight * (target_function * light_contribution.inverse_pdf);
103
104        weight_sum += resampling_weight;
105
106        if rand_f(rng) * weight_sum < resampling_weight {
107            selected_sample_radiance = contribution;
108            selected_sample_target_function = target_function;
109            selected_sample_world_position = resolved_light_sample.world_position;
110        }
111    }
112
113    var unbiased_contribution_weight = 0.0;
114    if any(selected_sample_radiance != vec3(0.0)) {
115        let inverse_target_function = select(0.0, 1.0 / selected_sample_target_function, selected_sample_target_function > 0.0);
116        unbiased_contribution_weight = weight_sum * inverse_target_function;
117
118        unbiased_contribution_weight *= trace_visibility(world_position + (world_normal * RAY_T_MIN), selected_sample_world_position);
119    }
120
121    return selected_sample_radiance * unbiased_contribution_weight;
122}