1#import bevy_hanabi::vfx_common::{
2    ChildInfo, ChildInfoBuffer, SimParams, Spawner,
3    EM_OFFSET_ALIVE_COUNT, EM_OFFSET_MAX_UPDATE, EM_OFFSET_CAPACITY,
4    EM_OFFSET_MAX_SPAWN, EM_OFFSET_INDIRECT_DISPATCH_INDEX, DRAW_INDEXED_INDIRECT_STRIDE,
5    EM_OFFSET_INDIRECT_WRITE_INDEX, DISPATCH_INDIRECT_STRIDE, EFFECT_METADATA_STRIDE
6}
7
8@group(0) @binding(0) var<uniform> sim_params : SimParams;
9
10// Tightly packed array of EffectMetadata[], accessed as u32 array.
11@group(1) @binding(0) var<storage, read_write> effect_metadata_buffer : array<u32>;
12// Tightly packed array of DispatchIndirectArgs[], accessed as u32 array.
13@group(1) @binding(1) var<storage, read_write> dispatch_indirect_buffer : array<u32>;
14// Tightly packed array of DrawIndexedIndirectArgs[], accessed as u32 array. This can contain
15// some DrawIndirectArgs[] instead, but in that case the stride is adjusted so all rows have
16// the same size. Since we access the instance_count, which is at the same position in both,
17// we ignore their size difference (the non-indexed one is padded).
18@group(1) @binding(2) var<storage, read_write> draw_indirect_buffer : array<u32>;
19
20@group(2) @binding(0) var<storage, read_write> spawner_buffer : array<Spawner>;
21
22#ifdef HAS_GPU_SPAWN_EVENTS
23@group(3) @binding(0) var<storage, read_write> child_info_buffer : ChildInfoBuffer;
24#endif
25
26/// Calculate the indirect workgroups counts based on the number of particles alive.
27@compute @workgroup_size(64)
28fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
29    let thread_index = global_invocation_id.x;
30
31#ifdef HAS_GPU_SPAWN_EVENTS
32    // Clear any GPU event. The indexing is safe because there are always less child effects
33    // than there are effects in total, so 'index' will always cover the entire child info array.
34    if (thread_index < arrayLength(&child_info_buffer.rows)) {
35        child_info_buffer.rows[thread_index].event_count = 0;
36    }
37#endif
38
39    // Cap at maximum number of effects
40    let effect_index = thread_index;
41    if (effect_index >= sim_params.num_effects) {
42        return;
43    }
44
45    let effect_metadata_index = spawner_buffer[effect_index].effect_metadata_index;
46    let em_base = EFFECT_METADATA_STRIDE * effect_metadata_index;
47
48    // Clear the rendering instance count, which will be upgraded by the update pass
49    // with the particles actually alive at the end of their update (after aged).
50    let draw_indirect_index = spawner_buffer[effect_index].draw_indirect_index;
51    let dri_base = DRAW_INDEXED_INDIRECT_STRIDE * draw_indirect_index;
52    draw_indirect_buffer[dri_base + 1u] = 0u;
53
54    let capacity = effect_metadata_buffer[em_base + EM_OFFSET_CAPACITY];
55    let alive_count = effect_metadata_buffer[em_base + EM_OFFSET_ALIVE_COUNT];
56    let dead_count = capacity - alive_count;
57
58    // Update max_update from current value of alive_count, so that the
59    // update pass coming next can cap its threads to this value, while also
60    // atomically modifying alive_count itself for next frame.
61    effect_metadata_buffer[em_base + EM_OFFSET_MAX_UPDATE] = alive_count;
62
63    // Copy the number of dead particles to a constant location, so that the
64    // init pass on next frame can atomically modify alive_count in parallel
65    // yet still read its initial value at the beginning of the init pass,
66    // and limit the number of particles spawned to the number of dead
67    // particles to recycle.
68    effect_metadata_buffer[em_base + EM_OFFSET_MAX_SPAWN] = dead_count;
69
70    // Calculate the number of workgroups (thread groups) to dispatch for the update
71    // pass, which is the number of alive particles rounded up to 64 (workgroup_size).
72    let indirect_dispatch_index = effect_metadata_buffer[em_base + EM_OFFSET_INDIRECT_DISPATCH_INDEX];
73    let di_base = DISPATCH_INDIRECT_STRIDE * indirect_dispatch_index;
74    dispatch_indirect_buffer[di_base] = (alive_count + 63u) >> 6u;
75    dispatch_indirect_buffer[di_base + 1u] = 1u;
76    dispatch_indirect_buffer[di_base + 2u] = 1u;
77
78    // Swap ping/pong buffers. The update pass always writes into ping, and both the update
79    // pass and the render pass always read from pong.
80    let ping = effect_metadata_buffer[em_base + EM_OFFSET_INDIRECT_WRITE_INDEX];
81    let pong = 1u - ping;
82    effect_metadata_buffer[em_base + EM_OFFSET_INDIRECT_WRITE_INDEX] = pong;
83
84    // Copy the new pong into the spawner buffer, which will be used during rendering
85    // to determine where to read particle indices.
86    spawner_buffer[effect_index].render_indirect_read_index = pong;
87}