Don't reallocate work item buffers every frame. (#17684)

We were calling `clear()` on the work item buffer table, which caused us
to deallocate all the CPU side buffers. This patch changes the logic to
instead just clear the buffers individually, but leave their backing
stores. This has two consequences:

1. To effectively retain work item buffers from frame to frame, we need
to key them off `RetainedViewEntity` values and not the render world
`Entity`, which is transient. This PR changes those buffers accordingly.

2. We need to clean up work item buffers that belong to views that went
away. Amusingly enough, we actually have a system,
`delete_old_work_item_buffers`, that tries to do this already, but it
wasn't doing anything because the `clear_batched_gpu_instance_buffers`
system already handled that. This patch actually makes the
`delete_old_work_item_buffers` system useful, by removing the clearing
behavior from `clear_batched_gpu_instance_buffers` and instead making
`delete_old_work_item_buffers` delete buffers corresponding to
nonexistent views.

On Bistro, this PR improves the performance of
`batch_and_prepare_binned_render_phase` from 61.2 us to 47.8 us, a 28%
speedup.

![Screenshot 2025-02-04
135542](https://github.com/user-attachments/assets/b0ecb551-f6c8-4677-8e4e-e39aa28115a3)

Author: pcwalton

crates/bevy_pbr/src/render/gpu_preprocess.rs

@@ -48,7 +48,7 @@ use bevy_render::{
     },
     renderer::{RenderContext, RenderDevice, RenderQueue},
     settings::WgpuFeatures,
-    view::{NoIndirectDrawing, ViewUniform, ViewUniformOffset, ViewUniforms},
+    view::{ExtractedView, NoIndirectDrawing, ViewUniform, ViewUniformOffset, ViewUniforms},
     Render, RenderApp, RenderSet,
 };
 use bevy_utils::TypeIdMap;
@@ -100,7 +100,7 @@ pub struct GpuMeshPreprocessPlugin {
 pub struct EarlyGpuPreprocessNode {
     view_query: QueryState<
         (
-            Entity,
+            Read<ExtractedView>,
             Option<Read<PreprocessBindGroups>>,
             Option<Read<ViewUniformOffset>>,
             Has<NoIndirectDrawing>,
@@ -120,7 +120,11 @@ pub struct EarlyGpuPreprocessNode {
 /// metadata for the subsequent [`LatePrepassBuildIndirectParametersNode`].
 pub struct LateGpuPreprocessNode {
     view_query: QueryState<
-        (Entity, Read<PreprocessBindGroups>, Read<ViewUniformOffset>),
+        (
+            Read<ExtractedView>,
+            Read<PreprocessBindGroups>,
+            Read<ViewUniformOffset>,
+        ),
         (
             Without<SkipGpuPreprocess>,
             Without<NoIndirectDrawing>,
@@ -580,7 +584,8 @@ impl Node for EarlyGpuPreprocessNode {
             };
 
             // Grab the work item buffers for this view.
-            let Some(phase_work_item_buffers) = index_buffers.get(&view) else {
+            let Some(phase_work_item_buffers) = index_buffers.get(&view.retained_view_entity)
+            else {
                 warn!("The preprocessing index buffer wasn't present");
                 continue;
             };
@@ -791,7 +796,8 @@ impl Node for LateGpuPreprocessNode {
         // Run the compute passes.
         for (view, bind_groups, view_uniform_offset) in self.view_query.iter_manual(world) {
             // Grab the work item buffers for this view.
-            let Some(phase_work_item_buffers) = work_item_buffers.get(&view) else {
+            let Some(phase_work_item_buffers) = work_item_buffers.get(&view.retained_view_entity)
+            else {
                 warn!("The preprocessing index buffer wasn't present");
                 continue;
             };
@@ -1619,6 +1625,7 @@ impl BuildIndirectParametersPipeline {
 )]
 pub fn prepare_preprocess_bind_groups(
     mut commands: Commands,
+    views: Query<(Entity, &ExtractedView)>,
     view_depth_pyramids: Query<(&ViewDepthPyramid, &PreviousViewUniformOffset)>,
     render_device: Res<RenderDevice>,
     batched_instance_buffers: Res<BatchedInstanceBuffers<MeshUniform, MeshInputUniform>>,
@@ -1651,14 +1658,19 @@ pub fn prepare_preprocess_bind_groups(
     let mut any_indirect = false;
 
     // Loop over each view.
-    for (view, phase_work_item_buffers) in work_item_buffers {
+    for (view_entity, view) in &views {
+        let Some(phase_work_item_buffers) = work_item_buffers.get(&view.retained_view_entity)
+        else {
+            continue;
+        };
+
         let mut bind_groups = TypeIdMap::default();
 
         // Loop over each phase.
         for (&phase_id, work_item_buffers) in phase_work_item_buffers {
             // Create the `PreprocessBindGroupBuilder`.
             let preprocess_bind_group_builder = PreprocessBindGroupBuilder {
-                view: *view,
+                view: view_entity,
                 late_indexed_indirect_parameters_buffer,
                 late_non_indexed_indirect_parameters_buffer,
                 render_device: &render_device,
@@ -1719,7 +1731,7 @@ pub fn prepare_preprocess_bind_groups(
 
         // Save the bind groups.
         commands
-            .entity(*view)
+            .entity(view_entity)
             .insert(PreprocessBindGroups(bind_groups));
     }
 

crates/bevy_render/src/batching/gpu_preprocessing.rs

@@ -4,7 +4,6 @@ use core::any::TypeId;
 
 use bevy_app::{App, Plugin};
 use bevy_ecs::{
-    entity::{hash_map::EntityHashMap, Entity},
     query::{Has, With},
     resource::Resource,
     schedule::IntoSystemConfigs as _,
@@ -13,7 +12,7 @@ use bevy_ecs::{
 };
 use bevy_encase_derive::ShaderType;
 use bevy_math::UVec4;
-use bevy_platform_support::collections::hash_map::Entry;
+use bevy_platform_support::collections::{hash_map::Entry, HashMap, HashSet};
 use bevy_utils::{default, TypeIdMap};
 use bytemuck::{Pod, Zeroable};
 use nonmax::NonMaxU32;
@@ -30,7 +29,7 @@ use crate::{
     },
     render_resource::{Buffer, BufferVec, GpuArrayBufferable, RawBufferVec, UninitBufferVec},
     renderer::{RenderAdapter, RenderDevice, RenderQueue},
-    view::{ExtractedView, NoIndirectDrawing},
+    view::{ExtractedView, NoIndirectDrawing, RetainedViewEntity},
     Render, RenderApp, RenderSet,
 };
 
@@ -157,7 +156,7 @@ where
     /// corresponds to each instance.
     ///
     /// This is keyed off each view. Each view has a separate buffer.
-    pub work_item_buffers: EntityHashMap<TypeIdMap<PreprocessWorkItemBuffers>>,
+    pub work_item_buffers: HashMap<RetainedViewEntity, TypeIdMap<PreprocessWorkItemBuffers>>,
 
     /// The uniform data inputs for the current frame.
     ///
@@ -409,8 +408,8 @@ impl Default for LatePreprocessWorkItemIndirectParameters {
 /// You may need to call this function if you're implementing your own custom
 /// render phases. See the `specialized_mesh_pipeline` example.
 pub fn get_or_create_work_item_buffer<'a, I>(
-    work_item_buffers: &'a mut EntityHashMap<TypeIdMap<PreprocessWorkItemBuffers>>,
-    view: Entity,
+    work_item_buffers: &'a mut HashMap<RetainedViewEntity, TypeIdMap<PreprocessWorkItemBuffers>>,
+    view: RetainedViewEntity,
     no_indirect_drawing: bool,
     gpu_occlusion_culling: bool,
     late_indexed_indirect_parameters_buffer: &'_ mut RawBufferVec<
@@ -493,6 +492,28 @@ impl PreprocessWorkItemBuffers {
             }
         }
     }
+
+    /// Clears out the GPU work item buffers in preparation for a new frame.
+    pub fn clear(&mut self) {
+        match *self {
+            PreprocessWorkItemBuffers::Direct(ref mut buffer) => {
+                buffer.clear();
+            }
+            PreprocessWorkItemBuffers::Indirect {
+                indexed: ref mut indexed_buffer,
+                non_indexed: ref mut non_indexed_buffer,
+                ref mut gpu_occlusion_culling,
+            } => {
+                indexed_buffer.clear();
+                non_indexed_buffer.clear();
+
+                if let Some(ref mut gpu_occlusion_culling) = *gpu_occlusion_culling {
+                    gpu_occlusion_culling.late_indexed.clear();
+                    gpu_occlusion_culling.late_non_indexed.clear();
+                }
+            }
+        }
+    }
 }
 
 /// One invocation of the preprocessing shader: i.e. one mesh instance in a
@@ -942,7 +963,7 @@ where
     pub fn new() -> Self {
         BatchedInstanceBuffers {
             data_buffer: UninitBufferVec::new(BufferUsages::STORAGE),
-            work_item_buffers: EntityHashMap::default(),
+            work_item_buffers: HashMap::default(),
             current_input_buffer: InstanceInputUniformBuffer::new(),
             previous_input_buffer: InstanceInputUniformBuffer::new(),
             late_indexed_indirect_parameters_buffer: RawBufferVec::new(
@@ -968,7 +989,14 @@ where
         self.data_buffer.clear();
         self.late_indexed_indirect_parameters_buffer.clear();
         self.late_non_indexed_indirect_parameters_buffer.clear();
-        self.work_item_buffers.clear();
+
+        // Clear each individual set of buffers, but don't depopulate the hash
+        // table. We want to avoid reallocating these vectors every frame.
+        for view_work_item_buffers in self.work_item_buffers.values_mut() {
+            for phase_work_item_buffers in view_work_item_buffers.values_mut() {
+                phase_work_item_buffers.clear();
+            }
+        }
     }
 }
 
@@ -1074,13 +1102,17 @@ pub fn delete_old_work_item_buffers<GFBD>(
     mut gpu_batched_instance_buffers: ResMut<
         BatchedInstanceBuffers<GFBD::BufferData, GFBD::BufferInputData>,
     >,
-    extracted_views: Query<Entity, With<ExtractedView>>,
+    extracted_views: Query<&ExtractedView>,
 ) where
     GFBD: GetFullBatchData,
 {
+    let retained_view_entities: HashSet<_> = extracted_views
+        .iter()
+        .map(|extracted_view| extracted_view.retained_view_entity)
+        .collect();
     gpu_batched_instance_buffers
         .work_item_buffers
-        .retain(|entity, _| extracted_views.contains(*entity));
+        .retain(|retained_view_entity, _| retained_view_entities.contains(retained_view_entity));
 }
 
 /// Batch the items in a sorted render phase, when GPU instance buffer building
@@ -1091,7 +1123,6 @@ pub fn batch_and_prepare_sorted_render_phase<I, GFBD>(
     mut indirect_parameters_buffers: ResMut<IndirectParametersBuffers>,
     mut sorted_render_phases: ResMut<ViewSortedRenderPhases<I>>,
     mut views: Query<(
-        Entity,
         &ExtractedView,
         Has<NoIndirectDrawing>,
         Has<OcclusionCulling>,
@@ -1110,15 +1141,15 @@ pub fn batch_and_prepare_sorted_render_phase<I, GFBD>(
         ..
     } = gpu_array_buffer.into_inner();
 
-    for (view, extracted_view, no_indirect_drawing, gpu_occlusion_culling) in &mut views {
+    for (extracted_view, no_indirect_drawing, gpu_occlusion_culling) in &mut views {
         let Some(phase) = sorted_render_phases.get_mut(&extracted_view.retained_view_entity) else {
             continue;
         };
 
         // Create the work item buffer if necessary.
         let work_item_buffer = get_or_create_work_item_buffer::<I>(
             work_item_buffers,
-            view,
+            extracted_view.retained_view_entity,
             no_indirect_drawing,
             gpu_occlusion_culling,
             late_indexed_indirect_parameters_buffer,
@@ -1248,7 +1279,6 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
     mut binned_render_phases: ResMut<ViewBinnedRenderPhases<BPI>>,
     mut views: Query<
         (
-            Entity,
             &ExtractedView,
             Has<NoIndirectDrawing>,
             Has<OcclusionCulling>,
@@ -1270,7 +1300,7 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
         ..
     } = gpu_array_buffer.into_inner();
 
-    for (view, extracted_view, no_indirect_drawing, gpu_occlusion_culling) in &mut views {
+    for (extracted_view, no_indirect_drawing, gpu_occlusion_culling) in &mut views {
         let Some(phase) = binned_render_phases.get_mut(&extracted_view.retained_view_entity) else {
             continue;
         };
@@ -1279,7 +1309,7 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
         // used this frame.
         let work_item_buffer = get_or_create_work_item_buffer::<BPI>(
             work_item_buffers,
-            view,
+            extracted_view.retained_view_entity,
             no_indirect_drawing,
             gpu_occlusion_culling,
             late_indexed_indirect_parameters_buffer,

examples/shader/specialized_mesh_pipeline.rs

@@ -280,7 +280,6 @@ fn queue_custom_mesh_pipeline(
     ),
     mut specialized_mesh_pipelines: ResMut<SpecializedMeshPipelines<CustomMeshPipeline>>,
     views: Query<(
-        Entity,
         &RenderVisibleEntities,
         &ExtractedView,
         &Msaa,
@@ -318,14 +317,8 @@ fn queue_custom_mesh_pipeline(
     // Render phases are per-view, so we need to iterate over all views so that
     // the entity appears in them. (In this example, we have only one view, but
     // it's good practice to loop over all views anyway.)
-    for (
-        view_entity,
-        view_visible_entities,
-        view,
-        msaa,
-        no_indirect_drawing,
-        gpu_occlusion_culling,
-    ) in views.iter()
+    for (view_visible_entities, view, msaa, no_indirect_drawing, gpu_occlusion_culling) in
+        views.iter()
     {
         let Some(opaque_phase) = opaque_render_phases.get_mut(&view.retained_view_entity) else {
             continue;
@@ -336,7 +329,7 @@ fn queue_custom_mesh_pipeline(
         // enabled.
         let work_item_buffer = gpu_preprocessing::get_or_create_work_item_buffer::<Opaque3d>(
             work_item_buffers,
-            view_entity,
+            view.retained_view_entity,
             no_indirect_drawing,
             gpu_occlusion_culling,
             late_indexed_indirect_parameters_buffer,