Skip to main content

bevy_render/render_resource/
pipeline_specializer.rs

1use bevy_material::descriptor::{
2    CachedComputePipelineId, CachedRenderPipelineId, ComputePipelineDescriptor,
3    RenderPipelineDescriptor,
4};
5
6use crate::render_resource::PipelineCache;
7use bevy_ecs::resource::Resource;
8use bevy_log::error;
9use bevy_material::specialize::SpecializedMeshPipelineError;
10use bevy_mesh::{MeshVertexBufferLayoutRef, VertexBufferLayout};
11use bevy_platform::{
12    collections::{
13        hash_map::{Entry, RawEntryMut, VacantEntry},
14        HashMap,
15    },
16    hash::FixedHasher,
17};
18use bevy_utils::default;
19use core::hash::{BuildHasher, Hash, Hasher};
20
21/// A trait that allows constructing different variants of a render pipeline from a key.
22///
23/// Note: This is intended for modifying your pipeline descriptor on the basis of a key. If your key
24/// contains no data then you don't need to specialize. For example, if you are using the
25/// [`AsBindGroup`](crate::render_resource::AsBindGroup) without the `#[bind_group_data]` attribute,
26/// you don't need to specialize. Instead, create the pipeline directly from [`PipelineCache`] and
27/// store its ID.
28///
29/// See [`SpecializedRenderPipelines`] for more info.
30pub trait SpecializedRenderPipeline {
31    /// The key that defines each "variant" of the render pipeline.
32    type Key: Clone + Hash + PartialEq + Eq;
33
34    /// Construct a new render pipeline based on the provided key.
35    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor;
36}
37
38/// A convenience cache for creating different variants of a render pipeline based on some key.
39///
40/// Some render pipelines may need to be configured differently depending on the exact situation.
41/// This cache allows constructing different render pipelines for each situation based on a key,
42/// making it easy to A) construct the necessary pipelines, and B) reuse already constructed
43/// pipelines.
44///
45/// Note: This is intended for modifying your pipeline descriptor on the basis of a key. If your key
46/// contains no data then you don't need to specialize. For example, if you are using the
47/// [`AsBindGroup`](crate::render_resource::AsBindGroup) without the `#[bind_group_data]` attribute,
48/// you don't need to specialize. Instead, create the pipeline directly from [`PipelineCache`] and
49/// store its ID.
50#[derive(Resource)]
51pub struct SpecializedRenderPipelines<S: SpecializedRenderPipeline> {
52    cache: HashMap<S::Key, CachedRenderPipelineId>,
53}
54
55impl<S: SpecializedRenderPipeline> Default for SpecializedRenderPipelines<S> {
56    fn default() -> Self {
57        Self { cache: default() }
58    }
59}
60
61impl<S: SpecializedRenderPipeline> SpecializedRenderPipelines<S> {
62    /// Get or create a specialized instance of the pipeline corresponding to `key`.
63    pub fn specialize(
64        &mut self,
65        cache: &PipelineCache,
66        pipeline_specializer: &S,
67        key: S::Key,
68    ) -> CachedRenderPipelineId {
69        *self.cache.entry(key.clone()).or_insert_with(|| {
70            let descriptor = pipeline_specializer.specialize(key);
71            cache.queue_render_pipeline(descriptor)
72        })
73    }
74}
75
76/// A trait that allows constructing different variants of a compute pipeline from a key.
77///
78/// Note: This is intended for modifying your pipeline descriptor on the basis of a key. If your key
79/// contains no data then you don't need to specialize. For example, if you are using the
80/// [`AsBindGroup`](crate::render_resource::AsBindGroup) without the `#[bind_group_data]` attribute,
81/// you don't need to specialize. Instead, create the pipeline directly from [`PipelineCache`] and
82/// store its ID.
83///
84/// See [`SpecializedComputePipelines`] for more info.
85pub trait SpecializedComputePipeline {
86    /// The key that defines each "variant" of the compute pipeline.
87    type Key: Clone + Hash + PartialEq + Eq;
88
89    /// Construct a new compute pipeline based on the provided key.
90    fn specialize(&self, key: Self::Key) -> ComputePipelineDescriptor;
91}
92
93/// A convenience cache for creating different variants of a compute pipeline based on some key.
94///
95/// Some compute pipelines may need to be configured differently depending on the exact situation.
96/// This cache allows constructing different compute pipelines for each situation based on a key,
97/// making it easy to A) construct the necessary pipelines, and B) reuse already constructed
98/// pipelines.
99///
100/// Note: This is intended for modifying your pipeline descriptor on the basis of a key. If your key
101/// contains no data then you don't need to specialize. For example, if you are using the
102/// [`AsBindGroup`](crate::render_resource::AsBindGroup) without the `#[bind_group_data]` attribute,
103/// you don't need to specialize. Instead, create the pipeline directly from [`PipelineCache`] and
104/// store its ID.
105#[derive(Resource)]
106pub struct SpecializedComputePipelines<S: SpecializedComputePipeline> {
107    cache: HashMap<S::Key, CachedComputePipelineId>,
108}
109
110impl<S: SpecializedComputePipeline> Default for SpecializedComputePipelines<S> {
111    fn default() -> Self {
112        Self { cache: default() }
113    }
114}
115
116impl<S: SpecializedComputePipeline> SpecializedComputePipelines<S> {
117    /// Get or create a specialized instance of the pipeline corresponding to `key`.
118    pub fn specialize(
119        &mut self,
120        cache: &PipelineCache,
121        specialize_pipeline: &S,
122        key: S::Key,
123    ) -> CachedComputePipelineId {
124        *self.cache.entry(key.clone()).or_insert_with(|| {
125            let descriptor = specialize_pipeline.specialize(key);
126            cache.queue_compute_pipeline(descriptor)
127        })
128    }
129}
130
131/// A trait that allows constructing different variants of a render pipeline from a key and the
132/// particular mesh's vertex buffer layout.
133///
134/// See [`SpecializedMeshPipelines`] for more info.
135pub trait SpecializedMeshPipeline {
136    /// The key that defines each "variant" of the render pipeline.
137    type Key: Clone + Hash + PartialEq + Eq;
138
139    /// Construct a new render pipeline based on the provided key and vertex layout.
140    ///
141    /// The returned pipeline descriptor should have a single vertex buffer, which is derived from
142    /// `layout`.
143    fn specialize(
144        &self,
145        key: Self::Key,
146        layout: &MeshVertexBufferLayoutRef,
147    ) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError>;
148}
149
150/// A cache of different variants of a render pipeline based on a key and the particular mesh's
151/// vertex buffer layout.
152#[derive(Resource)]
153pub struct SpecializedMeshPipelines<S: SpecializedMeshPipeline> {
154    mesh_layout_cache: HashMap<(MeshVertexBufferLayoutRef, S::Key), CachedRenderPipelineId>,
155    vertex_layout_cache: VertexLayoutCache<S>,
156}
157
158type VertexLayoutCache<S> = HashMap<
159    VertexBufferLayout,
160    HashMap<<S as SpecializedMeshPipeline>::Key, CachedRenderPipelineId>,
161>;
162
163impl<S: SpecializedMeshPipeline> Default for SpecializedMeshPipelines<S> {
164    fn default() -> Self {
165        Self {
166            mesh_layout_cache: Default::default(),
167            vertex_layout_cache: Default::default(),
168        }
169    }
170}
171
172impl<S: SpecializedMeshPipeline> SpecializedMeshPipelines<S> {
173    #[inline]
174    pub fn get_pipeline(
175        &self,
176        key: &S::Key,
177        layout: &MeshVertexBufferLayoutRef,
178    ) -> Option<CachedRenderPipelineId> {
179        let hash = {
180            let mut hasher = self.mesh_layout_cache.hasher().build_hasher();
181            layout.hash(&mut hasher);
182            key.hash(&mut hasher);
183            hasher.finish()
184        };
185
186        self.mesh_layout_cache
187            .raw_entry()
188            .from_hash(hash, |(l, k)| l == layout && k == key)
189            .map(|(_, id)| *id)
190    }
191
192    /// Construct a new render pipeline based on the provided key and the mesh's vertex buffer
193    /// layout.
194    #[inline]
195    pub fn specialize(
196        &mut self,
197        cache: &PipelineCache,
198        pipeline_specializer: &S,
199        key: S::Key,
200        layout: &MeshVertexBufferLayoutRef,
201    ) -> Result<CachedRenderPipelineId, SpecializedMeshPipelineError> {
202        return match self.mesh_layout_cache.entry((layout.clone(), key.clone())) {
203            Entry::Occupied(entry) => Ok(*entry.into_mut()),
204            Entry::Vacant(entry) => specialize_slow(
205                &mut self.vertex_layout_cache,
206                cache,
207                pipeline_specializer,
208                key,
209                layout,
210                entry,
211            ),
212        };
213
214        #[cold]
215        fn specialize_slow<S>(
216            vertex_layout_cache: &mut VertexLayoutCache<S>,
217            cache: &PipelineCache,
218            specialize_pipeline: &S,
219            key: S::Key,
220            layout: &MeshVertexBufferLayoutRef,
221            entry: VacantEntry<
222                (MeshVertexBufferLayoutRef, S::Key),
223                CachedRenderPipelineId,
224                FixedHasher,
225            >,
226        ) -> Result<CachedRenderPipelineId, SpecializedMeshPipelineError>
227        where
228            S: SpecializedMeshPipeline,
229        {
230            let descriptor = specialize_pipeline
231                .specialize(key.clone(), layout)
232                .map_err(|mut err| {
233                    {
234                        let SpecializedMeshPipelineError::MissingVertexAttribute(err) = &mut err;
235                        err.pipeline_type = Some(core::any::type_name::<S>());
236                    }
237                    err
238                })?;
239            // Different MeshVertexBufferLayouts can produce the same final VertexBufferLayout
240            // We want compatible vertex buffer layouts to use the same pipelines, so we must "deduplicate" them
241            let layout_map = match vertex_layout_cache
242                .raw_entry_mut()
243                .from_key(&descriptor.vertex.buffers[0])
244            {
245                RawEntryMut::Occupied(entry) => entry.into_mut(),
246                RawEntryMut::Vacant(entry) => {
247                    entry
248                        .insert(descriptor.vertex.buffers[0].clone(), Default::default())
249                        .1
250                }
251            };
252            Ok(*entry.insert(match layout_map.entry(key) {
253                Entry::Occupied(entry) => {
254                    if cfg!(debug_assertions) {
255                        let stored_descriptor = cache.get_render_pipeline_descriptor(*entry.get());
256                        if stored_descriptor != &descriptor {
257                            error!(
258                                "The cached pipeline descriptor for {} is not \
259                                    equal to the generated descriptor for the given key. \
260                                    This means the SpecializePipeline implementation uses \
261                                    unused' MeshVertexBufferLayout information to specialize \
262                                    the pipeline. This is not allowed because it would invalidate \
263                                    the pipeline cache.",
264                                core::any::type_name::<S>()
265                            );
266                        }
267                    }
268                    *entry.into_mut()
269                }
270                Entry::Vacant(entry) => *entry.insert(cache.queue_render_pipeline(descriptor)),
271            }))
272        }
273    }
274}