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

bevy_render/render_resource/
pipeline_specializer.rs
