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