Unpack bins belonging to multidrawable batch sets on the GPU instead of on the CPU.

crates/bevy_pbr/src/render/mesh_preprocess.wgsl

@@ -15,7 +15,7 @@
 // respectively.
 
 #import bevy_pbr::mesh_preprocess_types::{
-    IndirectParametersCpuMetadata, IndirectParametersGpuMetadata, MeshInput
+    IndirectParametersCpuMetadata, IndirectParametersGpuMetadata, MeshInput, PreprocessWorkItem
 }
 #import bevy_pbr::mesh_types::{
     Mesh, MESH_FLAGS_AABB_BASED_VISIBILITY_RANGE_BIT, MESH_FLAGS_NO_FRUSTUM_CULLING_BIT,
@@ -43,17 +43,6 @@ struct MeshCullingData {
     aabb_half_extents: vec4<f32>,
 }
 
-// One invocation of this compute shader: i.e. one mesh instance in a view.
-struct PreprocessWorkItem {
-    // The index of the `MeshInput` in the `current_input` buffer that we read
-    // from.
-    input_index: u32,
-    // In direct mode, the index of the `Mesh` in `output` that we write to. In
-    // indirect mode, the index of the `IndirectParameters` in
-    // `indirect_parameters` that we write to.
-    output_or_indirect_parameters_index: u32,
-}
-
 // The parameters for the indirect compute dispatch for the late mesh
 // preprocessing phase.
 struct LatePreprocessWorkItemIndirectParameters {

crates/bevy_pbr/src/render/unpack_bins.wgsl

@@ -0,0 +1,85 @@
+// A compute shader that unpacks bins.
+//
+// This shader runs before mesh preprocessing in order to generate
+// `PreprocessWorkItem`s from cached bins. A single dispatch of this shader
+// corresponds to a single batch set, and one invocation of this shader
+// corresponds to one binned entity. Each shader invocation builds the work item
+// for its entity by copying over the mesh input uniform index and calculating
+// the position of the command in the indirect parameters buffer that will draw
+// that entity.
+
+#import bevy_pbr::mesh_preprocess_types::PreprocessWorkItem
+
+// Information needed to unpack bins belonging to a single batch set.
+struct BinUnpackingMetadata {
+    // The index of the first `PreprocessWorkItem` that this compute shader
+    // dispatch is to write to.
+    base_output_work_item_index: u32,
+    // The index of the first GPU indirect parameters command for this batch
+    // set.
+    base_indirect_parameters_index: u32,
+    // The number of binned mesh instances in the `binned_mesh_instances`
+    // array.
+    binned_mesh_instance_count: u32,
+    // Padding.
+    pad_a: u32,
+    // Padding.
+    pad_b: array<vec4<u32>, 15>,
+};
+
+// One mesh instance in a bin.
+//
+// This corresponds to the CPU-side `GpuRenderBinnedMeshInstance` structure.
+//
+// Note that this structure isn't sorted within the
+// `binned_mesh_instances` buffer. Instances from the same bin aren't
+// guaranteed to be adjacent to one another.
+struct BinnedMeshInstance {
+    // The index of the `MeshInputUniform` corresponding to this mesh instance
+    // in the `MeshInputUniform` buffer.
+    input_uniform_index: u32,
+    // The index of the bin that this mesh instance belongs to.
+    bin_index: u32,
+};
+
+// Metadata for the entire batch set.
+@group(0) @binding(0) var<uniform> bin_unpacking_metadata: BinUnpackingMetadata;
+
+// The input array of `BinnedMeshInstance`s.
+//
+// Note that this array isn't sorted.
+@group(0) @binding(1) var<storage> binned_mesh_instances: array<BinnedMeshInstance>;
+
+// The output list of `PreprocessWorkItem`s.
+@group(0) @binding(2) var<storage, read_write> preprocess_work_items: array<PreprocessWorkItem>;
+
+// A mapping from each `bin_index` to the index of the GPU indirect parameters
+// for this bin, relative to the start of the indirect parameters for this batch
+// set.
+@group(0) @binding(3) var<storage> bin_index_to_indirect_parameters_offset: array<u32>;
+
+@compute
+@workgroup_size(64)
+fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
+    // Figure out which instance we're looking at.
+    let global_id = global_invocation_id.x;
+    if (global_id >= bin_unpacking_metadata.binned_mesh_instance_count) {
+        return;
+    }
+
+    // Unpack the `BinnedMeshInstance`.
+    let input_uniform_index = binned_mesh_instances[global_id].input_uniform_index;
+    let bin_index = binned_mesh_instances[global_id].bin_index;
+
+    // Look up the indirect parameters index for this bin, relative to the first
+    // indirect parameters offset for this batch set.
+    let indirect_parameters_offset = bin_index_to_indirect_parameters_offset[bin_index];
+
+    // Determine the location we should write the work item to.
+    let output_index = bin_unpacking_metadata.base_output_work_item_index + global_id;
+
+    // Write out the resulting work item.
+    preprocess_work_items[output_index].input_index = input_uniform_index;
+    preprocess_work_items[output_index].output_or_indirect_parameters_index =
+        bin_unpacking_metadata.base_indirect_parameters_index + indirect_parameters_offset;
+}

crates/bevy_render/Cargo.toml

@@ -126,6 +126,10 @@ bitflags = "2"
 itertools = "0.14"
 weak-table = "0.3"
 
+[dev-dependencies]
+proptest = "1"
+proptest-derive = "0.2"
+
 [target.'cfg(all(target_arch = "wasm32", target_feature = "atomics"))'.dependencies]
 send_wrapper = { version = "0.6.0" }