1//! A compute shader that allocates `MeshUniform`s.
2//!
3//! This shader runs before mesh preprocessing in order to determine the
4//! positions of `MeshUniform`s. Unlike `MeshInputUniform`s, which are scattered
5//! throughout the buffer, `MeshUniform`s are indexed by instance ID, and so we
6//! must place instances of the same mesh together in the buffer. One dispatch
7//! call corresponds to one batch set (i.e. one multidraw operation), and one
8//! thread corresponds to one bin (a.k.a. draw, a.k.a. batch).
9//!
10//! Essentially, the goal of this shader is to perform a prefix sum, using the
11//! "scan-then-fan" approach. It has three phases:
12//!
13//! 1. *Local scan*: Perform a [Hillis-Steele scan] on each chunk of draws, where
14//! the size of each chunk (i.e. the number of draws) is equal to the workgroup
15//! size (256). Write the total size for this chunk to the fan buffer.
16//!
17//! 2. *Global scan*: Do a Hillis-Steele scan on the fan buffer. Now we know the
18//! running total for each chunk.
19//!
20//! 3. *Fan*: Copy the running total for each chunk to every element of that
21//! chunk.
22//!
23//! Note that, for batch sets (i.e. multidraw indirect calls) that have fewer
24//! than 256 batches in them, we only need step (1). This is the common case.
25//!
26//! [Hillis-Steele scan]: https://en.wikipedia.org/wiki/Prefix_sum#Algorithm_1:_Shorter_span,_more_parallel
27
28import bevy_render::occlusion_culling::mesh_preprocess_types::{BinMetadata, IndirectParametersMetadata};
29
30/// Information needed to allocate `MeshUniform`s.
31struct UniformAllocationMetadata {
32    /// The index of this batch set in the `IndirectBatchSet` array.
33    ///
34    /// We write this into the `indirect_parameters_metadata`.
35    batch_set_index: u32,
36
37    /// The number of bins (a.k.a. draws, a.k.a. batches) in this batch set.
38    bin_count: u32,
39
40    /// The index of the first set of indirect parameters for this batch set.
41    ///
42    /// This is also the index of the first `IndirectParametersMetadata`, as
43    /// that's a parallel array with the indirect parameters.
44    first_indirect_parameters_index: u32,
45
46    /// The index of the first `MeshUniform` slot for this batch set.
47    first_output_mesh_uniform_index: u32,
48
49    /// Padding.
50    pad: array<vec4<f32>, 15u>,
51};
52
53/// The number of threads in a workgroup.
54const WORKGROUP_SIZE: u32 = 256u;
55
56/// Information needed to allocate `MeshUniform`s.
57@group(0) @binding(0) var<uniform> allocate_uniforms_metadata: UniformAllocationMetadata;
58
59/// Information for each bin, including the indirect parameters offset and the
60/// instance count.
61@group(0) @binding(1) var<storage> bin_metadata: array<BinMetadata>;
62
63/// The array of indirect parameters metadata that we fill out, one for each
64/// batch.
65@group(0) @binding(2) var<storage, read_write> indirect_parameters_metadata:
66    array<IndirectParametersMetadata>;
67
68/// A temporary buffer that stores the mesh uniform index of the last instance
69/// plus one for each workgroup (i.e. for each 256-bin chunk).
70///
71/// This is accumulated in the second stage and written out in the third.
72@group(0) @binding(3) var<storage, read_write> fan_buffer: array<u32>;
73
74/// Scratch memory that stores the prefix sum for every element in our chunk.
75var<workgroup> output_offsets: array<u32, 256>;
76
77/// The first step of the prefix sum. This computes the prefix sum for each
78/// 256-element chunk.
79///
80/// Note that this will be the *only* step in the operation if the total number
81/// of bins in this batch set is 256 or fewer. Thus we must fill in the indirect
82/// parameters metadata for each batch here, as we can't guarantee that the
83/// following two steps will be run at all.
84@compute @workgroup_size(256, 1, 1)
85fn allocate_local_scan(
86    @builtin(local_invocation_id) local_id: vec3<u32>,
87    @builtin(workgroup_id) group_id: vec3<u32>,
88    @builtin(global_invocation_id) global_id: vec3<u32>
89) {
90    let bin_count = allocate_uniforms_metadata.bin_count;
91
92    let block_start = group_id.x * WORKGROUP_SIZE;
93    let block_end = min(block_start + WORKGROUP_SIZE, bin_count);
94
95    // If this is the first workgroup, take the first output index from the
96    // metadata into account. But if this is the second chunk or beyond, don't
97    // do that, as the second and third phases will add it in and we don't want
98    // to double-count it.
99    if (group_id.x == 0u) {
100        output_offsets[local_id.x] = allocate_uniforms_metadata.first_output_mesh_uniform_index;
101    } else {
102        output_offsets[local_id.x] = 0u;
103    }
104    workgroupBarrier();
105
106    // We're doing an inclusive sum, so put the instance count in the *next* bin.
107    if (global_id.x < block_end && local_id.x < WORKGROUP_SIZE - 1u) {
108        output_offsets[local_id.x + 1] = bin_metadata[global_id.x].instance_count;
109    }
110    workgroupBarrier();
111
112    // Prefix sum within our workgroup.
113    hillis_steele_scan(local_id.x);
114
115    // Now write the indirect parameters metadata for this batch. We fill in the
116    // `base_output_index` with the value of the prefix sum (which might be
117    // incomplete if this isn't the first chunk). We also populate a few
118    // bookkeeping fields for later rendering passes to use.
119    if (global_id.x < block_end) {
120        let indirect_parameters_offset =
121            allocate_uniforms_metadata.first_indirect_parameters_index +
122            bin_metadata[global_id.x].indirect_parameters_offset;
123        indirect_parameters_metadata[indirect_parameters_offset].base_output_index =
124            output_offsets[local_id.x];
125        indirect_parameters_metadata[indirect_parameters_offset].batch_set_index =
126            allocate_uniforms_metadata.batch_set_index;
127        // These parameters get filled in later. Initialize them to zero for now.
128        // This is required in the case of the early/late instance counts
129        // because the mesh preprocessing shader will atomically increment them.
130        indirect_parameters_metadata[indirect_parameters_offset].mesh_index = 0u;
131        indirect_parameters_metadata[indirect_parameters_offset].early_instance_count = 0u;
132        indirect_parameters_metadata[indirect_parameters_offset].late_instance_count = 0u;
133    }
134
135    // If this is the last element in the workgroup, put the total number of
136    // instances (plus the first output mesh uniform index if we're the first
137    // workgroup) in the fan buffer in preparation for the next phase.
138    if (local_id.x == WORKGROUP_SIZE - 1u) {
139        var chunk_total = output_offsets[WORKGROUP_SIZE - 1u];
140        if (global_id.x < block_end) {
141            chunk_total += bin_metadata[global_id.x].instance_count;
142        }
143        fan_buffer[group_id.x] = chunk_total;
144    }
145}
146
147/// The second step of the prefix sum.
148///
149/// This step takes the intermediate fan values computed in the previous step
150/// (i.e. the sum going out of each chunk) and performs one or more Hillis-Steele
151/// scans in order to compute the fan value going into each chunk.
152///
153/// This step is omitted if there are 256 or fewer total draws.
154@compute @workgroup_size(256, 1, 1)
155fn allocate_global_scan(@builtin(local_invocation_id) local_id: vec3<u32>) {
156    var sum = 0u;
157    let chunk_count = div_ceil(allocate_uniforms_metadata.bin_count, WORKGROUP_SIZE);
158
159    // Do a sequential loop over each block of 256 chunks. Because each
160    // iteration of this loop covers 64K meshes, the fact that it's sequential
161    // isn't going to be a problem in practice.
162    for (var block_start = 0u; block_start < chunk_count; block_start += WORKGROUP_SIZE) {
163        // Set up the Hillis-Steele scan.
164        let block_end = min(block_start + WORKGROUP_SIZE, chunk_count);
165        let global_id = block_start + local_id.x;
166        if (global_id < block_end) {
167            output_offsets[local_id.x] = fan_buffer[global_id];
168        }
169        workgroupBarrier();
170
171        // Perform the scan.
172        hillis_steele_scan(local_id.x);
173
174        // Write the value back.
175        // Note that we don't need a workgroup barrier here because
176        // `hillis_steele_scan` already did one.
177        if (global_id < block_end) {
178            fan_buffer[global_id] = sum + output_offsets[local_id.x];
179        }
180
181        // Save the sum coming out of this block for the next one.
182        sum += output_offsets[WORKGROUP_SIZE - 1u];
183        workgroupBarrier();
184    }
185}
186
187/// The third step of the prefix sum.
188///
189/// We take the summed fan value computed in the previous step and add it in to
190/// each value of each chunk beyond the first. We dispatch one fewer workgroup
191/// here than in step (1), because there's nothing to do for the first chunk.
192///
193/// This step is omitted if there are 256 or fewer total draws.
194@compute @workgroup_size(256, 1, 1)
195fn allocate_fan(
196    @builtin(workgroup_id) group_id: vec3<u32>,
197    @builtin(global_invocation_id) global_id: vec3<u32>
198) {
199    let id = global_id.x + WORKGROUP_SIZE;
200    let bin_count = allocate_uniforms_metadata.bin_count;
201    if (id >= bin_count) {
202        return;
203    }
204
205    let fan_value = fan_buffer[group_id.x];
206    let indirect_parameters_offset =
207        allocate_uniforms_metadata.first_indirect_parameters_index +
208        bin_metadata[id].indirect_parameters_offset;
209    indirect_parameters_metadata[indirect_parameters_offset].base_output_index += fan_value;
210}
211
212/// Calculates a running exclusive sum.
213/// https://en.wikipedia.org/wiki/Prefix_sum#Algorithm_1:_Shorter_span,_more_parallel
214fn hillis_steele_scan(local_id: u32) {
215    for (var offset = 1u; offset < WORKGROUP_SIZE; offset *= 2u) {
216        var term = 0u;
217        if (local_id >= offset) {
218            term = output_offsets[local_id - offset];
219        }
220        workgroupBarrier();
221        output_offsets[local_id] += term;
222        workgroupBarrier();
223    }
224}
225
226/// Divides unsigned integer a by b, rounding up.
227fn div_ceil(a: u32, b: u32) -> u32 {
228    return (a + b - 1u) / b;
229}