1// https://cwyman.org/papers/hpg21_rearchitectingReSTIR.pdf
2
3enable wgpu_ray_query;
4
5#define_import_path bevy_solari::presample_light_tiles
6
7#import bevy_pbr::rgb9e5::{vec3_to_rgb9e5_, rgb9e5_to_vec3_}
8#import bevy_pbr::utils::{octahedral_encode, octahedral_decode}
9#import bevy_render::view::View
10#import bevy_solari::sampling::{generate_random_light_sample, LightSample, ResolvedLightSample}
11#import bevy_solari::realtime_bindings::{light_tile_samples, light_tile_resolved_samples, view, constants, ResolvedLightSamplePacked}
12
13@compute @workgroup_size(1024, 1, 1)
14fn presample_light_tiles(@builtin(workgroup_id) workgroup_id: vec3<u32>, @builtin(local_invocation_index) sample_index: u32) {
15    let tile_id = workgroup_id.x;
16    var rng = (tile_id * 5782582u) + sample_index + constants.frame_index;
17
18    let sample = generate_random_light_sample(&rng);
19
20    let i = (tile_id * 1024u) + sample_index;
21    light_tile_samples[i] = sample.light_sample;
22    light_tile_resolved_samples[i] = pack_resolved_light_sample(sample.resolved_light_sample);
23}
24
25fn pack_resolved_light_sample(sample: ResolvedLightSample) -> ResolvedLightSamplePacked {
26    return ResolvedLightSamplePacked(
27        sample.world_position.x,
28        sample.world_position.y,
29        sample.world_position.z,
30        pack2x16unorm(octahedral_encode(sample.world_normal)),
31        vec3_to_rgb9e5_(log2(sample.radiance * view.exposure + 1.0)),
32        sample.inverse_pdf * select(1.0, -1.0, sample.world_position.w == 0.0),
33    );
34}
35
36fn unpack_resolved_light_sample(packed: ResolvedLightSamplePacked, exposure: f32) -> ResolvedLightSample {
37    return ResolvedLightSample(
38        vec4(packed.world_position_x, packed.world_position_y, packed.world_position_z, select(1.0, 0.0, packed.inverse_pdf < 0.0)),
39        octahedral_decode(unpack2x16unorm(packed.world_normal)),
40        (exp2(rgb9e5_to_vec3_(packed.radiance)) - 1.0) / exposure,
41        abs(packed.inverse_pdf),
42    );
43}