bevy_render/render_resource/
pipeline_specializer.rs

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

bevy_render/render_resource/
pipeline_specializer.rs
