bevy_render/render_resource/
pipeline_specializer.rs

1use crate::{
2    mesh::{MeshVertexBufferLayoutRef, MissingVertexAttributeError, VertexBufferLayout},
3    render_resource::{
4        CachedComputePipelineId, CachedRenderPipelineId, ComputePipelineDescriptor, PipelineCache,
5        RenderPipelineDescriptor,
6    },
7};
8use bevy_ecs::system::Resource;
9use bevy_utils::{
10    default,
11    hashbrown::hash_map::{RawEntryMut, VacantEntry},
12    tracing::error,
13    Entry, HashMap,
14};
15use core::{fmt::Debug, hash::Hash};
16use derive_more::derive::{Display, Error, From};
17
18pub trait SpecializedRenderPipeline {
19    type Key: Clone + Hash + PartialEq + Eq;
20    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor;
21}
22
23#[derive(Resource)]
24pub struct SpecializedRenderPipelines<S: SpecializedRenderPipeline> {
25    cache: HashMap<S::Key, CachedRenderPipelineId>,
26}
27
28impl<S: SpecializedRenderPipeline> Default for SpecializedRenderPipelines<S> {
29    fn default() -> Self {
30        Self { cache: default() }
31    }
32}
33
34impl<S: SpecializedRenderPipeline> SpecializedRenderPipelines<S> {
35    pub fn specialize(
36        &mut self,
37        cache: &PipelineCache,
38        specialize_pipeline: &S,
39        key: S::Key,
40    ) -> CachedRenderPipelineId {
41        *self.cache.entry(key.clone()).or_insert_with(|| {
42            let descriptor = specialize_pipeline.specialize(key);
43            cache.queue_render_pipeline(descriptor)
44        })
45    }
46}
47
48pub trait SpecializedComputePipeline {
49    type Key: Clone + Hash + PartialEq + Eq;
50    fn specialize(&self, key: Self::Key) -> ComputePipelineDescriptor;
51}
52
53#[derive(Resource)]
54pub struct SpecializedComputePipelines<S: SpecializedComputePipeline> {
55    cache: HashMap<S::Key, CachedComputePipelineId>,
56}
57
58impl<S: SpecializedComputePipeline> Default for SpecializedComputePipelines<S> {
59    fn default() -> Self {
60        Self { cache: default() }
61    }
62}
63
64impl<S: SpecializedComputePipeline> SpecializedComputePipelines<S> {
65    pub fn specialize(
66        &mut self,
67        cache: &PipelineCache,
68        specialize_pipeline: &S,
69        key: S::Key,
70    ) -> CachedComputePipelineId {
71        *self.cache.entry(key.clone()).or_insert_with(|| {
72            let descriptor = specialize_pipeline.specialize(key);
73            cache.queue_compute_pipeline(descriptor)
74        })
75    }
76}
77
78pub trait SpecializedMeshPipeline {
79    type Key: Clone + Hash + PartialEq + Eq;
80    fn specialize(
81        &self,
82        key: Self::Key,
83        layout: &MeshVertexBufferLayoutRef,
84    ) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError>;
85}
86
87#[derive(Resource)]
88pub struct SpecializedMeshPipelines<S: SpecializedMeshPipeline> {
89    mesh_layout_cache: HashMap<(MeshVertexBufferLayoutRef, S::Key), CachedRenderPipelineId>,
90    vertex_layout_cache: VertexLayoutCache<S>,
91}
92
93pub type VertexLayoutCache<S> = HashMap<
94    VertexBufferLayout,
95    HashMap<<S as SpecializedMeshPipeline>::Key, CachedRenderPipelineId>,
96>;
97
98impl<S: SpecializedMeshPipeline> Default for SpecializedMeshPipelines<S> {
99    fn default() -> Self {
100        Self {
101            mesh_layout_cache: Default::default(),
102            vertex_layout_cache: Default::default(),
103        }
104    }
105}
106
107impl<S: SpecializedMeshPipeline> SpecializedMeshPipelines<S> {
108    #[inline]
109    pub fn specialize(
110        &mut self,
111        cache: &PipelineCache,
112        specialize_pipeline: &S,
113        key: S::Key,
114        layout: &MeshVertexBufferLayoutRef,
115    ) -> Result<CachedRenderPipelineId, SpecializedMeshPipelineError> {
116        return match self.mesh_layout_cache.entry((layout.clone(), key.clone())) {
117            Entry::Occupied(entry) => Ok(*entry.into_mut()),
118            Entry::Vacant(entry) => specialize_slow(
119                &mut self.vertex_layout_cache,
120                cache,
121                specialize_pipeline,
122                key,
123                layout,
124                entry,
125            ),
126        };
127
128        #[cold]
129        fn specialize_slow<S>(
130            vertex_layout_cache: &mut VertexLayoutCache<S>,
131            cache: &PipelineCache,
132            specialize_pipeline: &S,
133            key: S::Key,
134            layout: &MeshVertexBufferLayoutRef,
135            entry: VacantEntry<(MeshVertexBufferLayoutRef, S::Key), CachedRenderPipelineId>,
136        ) -> Result<CachedRenderPipelineId, SpecializedMeshPipelineError>
137        where
138            S: SpecializedMeshPipeline,
139        {
140            let descriptor = specialize_pipeline
141                .specialize(key.clone(), layout)
142                .map_err(|mut err| {
143                    {
144                        let SpecializedMeshPipelineError::MissingVertexAttribute(err) = &mut err;
145                        err.pipeline_type = Some(core::any::type_name::<S>());
146                    }
147                    err
148                })?;
149            // Different MeshVertexBufferLayouts can produce the same final VertexBufferLayout
150            // We want compatible vertex buffer layouts to use the same pipelines, so we must "deduplicate" them
151            let layout_map = match vertex_layout_cache
152                .raw_entry_mut()
153                .from_key(&descriptor.vertex.buffers[0])
154            {
155                RawEntryMut::Occupied(entry) => entry.into_mut(),
156                RawEntryMut::Vacant(entry) => {
157                    entry
158                        .insert(descriptor.vertex.buffers[0].clone(), Default::default())
159                        .1
160                }
161            };
162            Ok(*entry.insert(match layout_map.entry(key) {
163                Entry::Occupied(entry) => {
164                    if cfg!(debug_assertions) {
165                        let stored_descriptor = cache.get_render_pipeline_descriptor(*entry.get());
166                        if stored_descriptor != &descriptor {
167                            error!(
168                                "The cached pipeline descriptor for {} is not \
169                                    equal to the generated descriptor for the given key. \
170                                    This means the SpecializePipeline implementation uses \
171                                    unused' MeshVertexBufferLayout information to specialize \
172                                    the pipeline. This is not allowed because it would invalidate \
173                                    the pipeline cache.",
174                                core::any::type_name::<S>()
175                            );
176                        }
177                    }
178                    *entry.into_mut()
179                }
180                Entry::Vacant(entry) => *entry.insert(cache.queue_render_pipeline(descriptor)),
181            }))
182        }
183    }
184}
185
186#[derive(Error, Display, Debug, From)]
187pub enum SpecializedMeshPipelineError {
188    MissingVertexAttribute(MissingVertexAttributeError),
189}
bevy_render/render_resource/pipeline_specializer.rs

bevy_render/render_resource/
pipeline_specializer.rs
