1#import bevy_solari::world_cache::{WORLD_CACHE_EMPTY_CELL, world_cache_life, world_cache_checksums, world_cache_radiance, world_cache_a, world_cache_b, world_cache_active_cell_indices, world_cache_active_cells_count}
2
3@group(2) @binding(0) var<storage, read_write> world_cache_active_cells_dispatch: vec3<u32>;
4
5var<workgroup> w1: array<u32, 1024u>;
6var<workgroup> w2: array<u32, 1024u>;
7
8@compute @workgroup_size(1024, 1, 1)
9fn decay_world_cache(@builtin(global_invocation_id) global_id: vec3<u32>) {
10    var life = world_cache_life[global_id.x];
11    if life > 0u {
12        life -= 1u;
13        world_cache_life[global_id.x] = life;
14
15        if life == 0u {
16            world_cache_checksums[global_id.x] = WORLD_CACHE_EMPTY_CELL;
17            world_cache_radiance[global_id.x] = vec4(0.0);
18        }
19    }
20}
21
22@compute @workgroup_size(1024, 1, 1)
23fn compact_world_cache_single_block(
24    @builtin(global_invocation_id) cell_id: vec3<u32>,
25    @builtin(local_invocation_index) t: u32,
26) {
27    if t == 0u { w1[0u] = 0u; } else { w1[t] = u32(world_cache_life[cell_id.x - 1u] != 0u); }; workgroupBarrier();
28    if t < 1u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 1u]; } workgroupBarrier();
29    if t < 2u { w1[t] = w2[t]; } else { w1[t] = w2[t] + w2[t - 2u]; } workgroupBarrier();
30    if t < 4u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 4u]; } workgroupBarrier();
31    if t < 8u { w1[t] = w2[t]; } else { w1[t] = w2[t] + w2[t - 8u]; } workgroupBarrier();
32    if t < 16u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 16u]; } workgroupBarrier();
33    if t < 32u { w1[t] = w2[t]; } else { w1[t] = w2[t] + w2[t - 32u]; } workgroupBarrier();
34    if t < 64u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 64u]; } workgroupBarrier();
35    if t < 128u { w1[t] = w2[t]; } else { w1[t] = w2[t] + w2[t - 128u]; } workgroupBarrier();
36    if t < 256u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 256u]; } workgroupBarrier();
37    if t < 512u { world_cache_a[cell_id.x] = w2[t]; } else { world_cache_a[cell_id.x] = w2[t] + w2[t - 512u]; }
38}
39
40@compute @workgroup_size(1024, 1, 1)
41fn compact_world_cache_blocks(@builtin(local_invocation_index) t: u32) {
42    if t == 0u { w1[0u] = 0u; } else { w1[t] = world_cache_a[t * 1024u - 1u]; }; workgroupBarrier();
43    if t < 1u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 1u]; } workgroupBarrier();
44    if t < 2u { w1[t] = w2[t]; } else { w1[t] = w2[t] + w2[t - 2u]; } workgroupBarrier();
45    if t < 4u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 4u]; } workgroupBarrier();
46    if t < 8u { w1[t] = w2[t]; } else { w1[t] = w2[t] + w2[t - 8u]; } workgroupBarrier();
47    if t < 16u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 16u]; } workgroupBarrier();
48    if t < 32u { w1[t] = w2[t]; } else { w1[t] = w2[t] + w2[t - 32u]; } workgroupBarrier();
49    if t < 64u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 64u]; } workgroupBarrier();
50    if t < 128u { w1[t] = w2[t]; } else { w1[t] = w2[t] + w2[t - 128u]; } workgroupBarrier();
51    if t < 256u { w2[t] = w1[t]; } else { w2[t] = w1[t] + w1[t - 256u]; } workgroupBarrier();
52    if t < 512u { world_cache_b[t] = w2[t]; } else { world_cache_b[t] = w2[t] + w2[t - 512u]; }
53}
54
55@compute @workgroup_size(1024, 1, 1)
56fn compact_world_cache_write_active_cells(
57    @builtin(global_invocation_id) cell_id: vec3<u32>,
58    @builtin(workgroup_id) workgroup_id: vec3<u32>,
59    @builtin(local_invocation_index) thread_index: u32,
60) {
61    let compacted_index = world_cache_a[cell_id.x] + world_cache_b[workgroup_id.x];
62    if world_cache_life[cell_id.x] != 0u {
63        world_cache_active_cell_indices[compacted_index] = cell_id.x;
64    }
65
66    if thread_index == 1023u && workgroup_id.x == 1023u {
67        world_cache_active_cells_count = compacted_index + 1u; // TODO: This is 1 even when there are zero active entries in the cache
68        world_cache_active_cells_dispatch = vec3((world_cache_active_cells_count + 1023u) / 1024u, 1u, 1u);
69    }
70}