Reduce LayoutCache::get and BindGroupLayoutCache::get heap allocation

crates/bevy_render/src/render_resource/pipeline_cache.rs

@@ -2,6 +2,7 @@ use bevy_material::descriptor::{
     BindGroupLayoutDescriptor, CachedComputePipelineId, CachedRenderPipelineId,
     ComputePipelineDescriptor, PipelineDescriptor, RenderPipelineDescriptor,
 };
+use smallvec::SmallVec;
 
 use crate::{
     render_resource::*,
@@ -16,7 +17,7 @@ use bevy_ecs::{
     system::{Res, ResMut},
 };
 use bevy_log::error;
-use bevy_platform::collections::{HashMap, HashSet};
+use bevy_platform::collections::{hash_map::RawEntryMut, HashMap, HashSet};
 use bevy_shader::{
     CachedPipelineId, Shader, ShaderCache, ShaderCacheError, ShaderCacheSource, ShaderDefVal,
     ValidateShader,
@@ -79,8 +80,15 @@ impl CachedPipelineState {
     }
 }
 
+// The default webgpu `max_bind_groups` is 4. Most desktop GPUs support 8.
+// Since the element size we used below is small, we inline 8 on the stack.
+const BIND_GROUP_LAYOUTS_INLINE_CAPACITY: usize = 8;
+
 type ImmediateSize = u32;
-type LayoutCacheKey = (Vec<BindGroupLayoutId>, ImmediateSize);
+type LayoutCacheKey = (
+    SmallVec<[BindGroupLayoutId; BIND_GROUP_LAYOUTS_INLINE_CAPACITY]>,
+    ImmediateSize,
+);
 #[derive(Default)]
 struct LayoutCache {
     layouts: HashMap<LayoutCacheKey, Arc<WgpuWrapper<PipelineLayout>>>,
@@ -101,7 +109,7 @@ impl LayoutCache {
                     .iter()
                     .map(BindGroupLayout::value)
                     .map(Some)
-                    .collect::<Vec<_>>();
+                    .collect::<SmallVec<[_; BIND_GROUP_LAYOUTS_INLINE_CAPACITY]>>();
                 Arc::new(WgpuWrapper::new(render_device.create_pipeline_layout(
                     &PipelineLayoutDescriptor {
                         bind_group_layouts: &bind_group_layouts,
@@ -175,15 +183,17 @@ impl BindGroupLayoutCache {
     fn get(
         &mut self,
         render_device: &RenderDevice,
-        descriptor: BindGroupLayoutDescriptor,
+        descriptor: &BindGroupLayoutDescriptor,
     ) -> BindGroupLayout {
-        self.bgls
-            .entry(descriptor)
-            .or_insert_with_key(|descriptor| {
-                render_device
-                    .create_bind_group_layout(descriptor.label.as_ref(), &descriptor.entries)
-            })
-            .clone()
+        match self.bgls.raw_entry_mut().from_key(descriptor) {
+            RawEntryMut::Occupied(entry) => entry.into_mut(),
+            RawEntryMut::Vacant(slot) => {
+                let created = render_device
+                    .create_bind_group_layout(descriptor.label.as_ref(), &descriptor.entries);
+                slot.insert(descriptor.clone(), created).1
+            }
+        }
+        .clone()
     }
 }
 
@@ -442,7 +452,7 @@ impl PipelineCache {
         self.bindgroup_layout_cache
             .lock()
             .unwrap()
-            .get(&self.device, bind_group_layout_descriptor.clone())
+            .get(&self.device, bind_group_layout_descriptor)
     }
 
     /// Inserts a [`Shader`] into this cache with the provided [`AssetId`].
@@ -478,9 +488,9 @@ impl PipelineCache {
             .layout
             .iter()
             .map(|bind_group_layout_descriptor| {
-                bindgroup_layout_cache.get(&self.device, bind_group_layout_descriptor.clone())
+                bindgroup_layout_cache.get(&self.device, bind_group_layout_descriptor)
             })
-            .collect::<Vec<_>>();
+            .collect::<SmallVec<[_; BIND_GROUP_LAYOUTS_INLINE_CAPACITY]>>();
 
         create_pipeline_task(
             async move {
@@ -586,9 +596,9 @@ impl PipelineCache {
             .layout
             .iter()
             .map(|bind_group_layout_descriptor| {
-                bindgroup_layout_cache.get(&self.device, bind_group_layout_descriptor.clone())
+                bindgroup_layout_cache.get(&self.device, bind_group_layout_descriptor)
             })
-            .collect::<Vec<_>>();
+            .collect::<SmallVec<[_; BIND_GROUP_LAYOUTS_INLINE_CAPACITY]>>();
 
         create_pipeline_task(
             async move {