1//! Builds GPU indirect draw parameters from metadata.
2//!
3//! This only runs when indirect drawing is enabled. It takes the output of
4//! `mesh_preprocess.wesl` and creates indirect parameters for the GPU.
5//!
6//! This shader runs separately for indexed and non-indexed meshes. Unlike
7//! `mesh_preprocess.wesl`, which runs one instance per mesh *instance*, one
8//! instance of this shader corresponds to a single *batch* which could contain
9//! arbitrarily many instances of a single mesh.
10
11import bevy_render::occlusion_culling::mesh_preprocess_types::{
12    IndirectBatchSet,
13    IndirectParametersIndexed,
14    IndirectParametersNonIndexed,
15    IndirectParametersMetadata,
16    MeshInput
17};
18
19/// Specifies the batches that this shader invocation is to process.
20struct IndirectParametersBuildJob {
21    /// The first batch index that this shader invocation should process
22    /// (inclusive).
23    first_batch_index: u32,
24    /// The last batch index that this shader invocation should process
25    /// (exclusive).
26    last_batch_index: u32,
27}
28
29/// The data for each mesh that the CPU supplied to the GPU.
30@group(0) @binding(0) var<storage> current_input: array<MeshInput>;
31
32/// Data that we use to generate the indirect parameters.
33///
34/// The `mesh_preprocess.wesl` shader emits these.
35@group(0) @binding(1) var<storage> indirect_parameters_metadata:
36    array<IndirectParametersMetadata>;
37
38/// Information about each batch set.
39///
40/// A *batch set* is a set of meshes that might be multi-drawn together.
41@group(0) @binding(3) var<storage, read_write> indirect_batch_sets: array<IndirectBatchSet>;
42
43/// Specifies the batches that this shader invocation is to process.
44@group(0) @binding(4) var<uniform> indirect_parameters_build_job: IndirectParametersBuildJob;
45
46@if(INDEXED)
47/// The buffer of indirect draw parameters that we generate, and that the GPU
48/// reads to issue the draws.
49///
50/// This buffer is for indexed meshes.
51@group(0) @binding(5) var<storage, read_write> indirect_parameters:
52    array<IndirectParametersIndexed>;
53@else   // INDEXED
54/// The buffer of indirect draw parameters that we generate, and that the GPU
55/// reads to issue the draws.
56///
57/// This buffer is for non-indexed meshes.
58@group(0) @binding(5) var<storage, read_write> indirect_parameters:
59    array<IndirectParametersNonIndexed>;
60
61@compute
62@workgroup_size(64)
63fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
64    // Figure out our instance index (i.e. batch index). If this thread doesn't
65    // correspond to a valid index in the range.
66    let instance_index = global_invocation_id.x + indirect_parameters_build_job.first_batch_index;
67    if (instance_index >= indirect_parameters_build_job.last_batch_index) {
68        return;
69    }
70
71    // Unpack the metadata for this batch.
72    let base_output_index = indirect_parameters_metadata[instance_index].base_output_index;
73    let batch_set_index = indirect_parameters_metadata[instance_index].batch_set_index;
74    let mesh_index = indirect_parameters_metadata[instance_index].mesh_index;
75
76    // If we aren't using `multi_draw_indirect_count`, we have a 1:1 fixed
77    // assignment of batches to slots in the indirect parameters buffer, so we
78    // can just use the instance index as the index of our indirect parameters.
79    let early_instance_count =
80        indirect_parameters_metadata[instance_index].early_instance_count;
81    let late_instance_count = indirect_parameters_metadata[instance_index].late_instance_count;
82
83    // If in the early phase, we draw only the early meshes. If in the late
84    // phase, we draw only the late meshes. If in the main phase, draw all the
85    // meshes.
86@if(EARLY_PHASE)
87    let instance_count = early_instance_count;
88@elif(LATE_PHASE)
89    let instance_count = late_instance_count;
90@else   // LATE_PHASE
91    let instance_count = early_instance_count + late_instance_count;
92
93    var indirect_parameters_index = instance_index;
94
95    // If the current hardware and driver support `multi_draw_indirect_count`,
96    // dynamically reserve an index for the indirect parameters we're to
97    // generate.
98@if(MULTI_DRAW_INDIRECT_COUNT_SUPPORTED)
99    // If this batch belongs to a batch set, then allocate space for the
100    // indirect commands in that batch set.
101    if (batch_set_index != 0xffffffffu) {
102        // Bail out now if there are no instances. Note that we can only bail if
103        // we're in a batch set. That's because only batch sets are drawn using
104        // `multi_draw_indirect_count`. If we aren't using
105        // `multi_draw_indirect_count`, then we need to continue in order to
106        // zero out the instance count; otherwise, it'll have garbage data in
107        // it.
108        if (instance_count == 0u) {
109            return;
110        }
111
112        let indirect_parameters_base =
113            indirect_batch_sets[batch_set_index].indirect_parameters_base;
114        let indirect_parameters_offset =
115            atomicAdd(&indirect_batch_sets[batch_set_index].indirect_parameters_count, 1u);
116
117        indirect_parameters_index = indirect_parameters_base + indirect_parameters_offset;
118    }
119
120    // Build up the indirect parameters. The structures for indexed and
121    // non-indexed meshes are slightly different.
122
123    indirect_parameters[indirect_parameters_index].instance_count = instance_count;
124
125@if(LATE_PHASE)
126    // The late mesh instances are stored after the early mesh instances, so we
127    // offset the output index by the number of early mesh instances.
128    indirect_parameters[indirect_parameters_index].first_instance =
129        base_output_index + early_instance_count;
130@else   // LATE_PHASE
131    indirect_parameters[indirect_parameters_index].first_instance = base_output_index;
132
133    indirect_parameters[indirect_parameters_index].base_vertex =
134        current_input[mesh_index].first_vertex_index;
135
136@if(INDEXED)
137    indirect_parameters[indirect_parameters_index].index_count =
138        current_input[mesh_index].index_count;
139@if(INDEXED)
140    indirect_parameters[indirect_parameters_index].first_index =
141        current_input[mesh_index].first_index_index;
142@if(!INDEXED)
143    indirect_parameters[indirect_parameters_index].vertex_count =
144        current_input[mesh_index].index_count;
145}