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