docs/bevy/examples/shader_advanced/specialized_mesh_pipeline.rs 18 K · 447 lines · raw · history

1 //! Demonstrates how to define and use specialized mesh pipeline
2 //!
3 //! This example shows how to use the built-in [`SpecializedMeshPipeline`]
4 //! functionality with a custom [`RenderCommand`](bevy::render::render_phase::RenderCommand)
5 //! to allow custom mesh rendering with more flexibility than the material api.
6 //!
7 //! [`SpecializedMeshPipeline`] let's you customize the entire pipeline used when rendering a mesh.
8
9 use bevy::{
10 asset::RenderAssetUsages,
11 camera::visibility::{self, VisibilityClass},
12 core_pipeline::core_3d::{Opaque3d, Opaque3dBatchSetKey, Opaque3dBinKey, CORE_3D_DEPTH_FORMAT},
13 ecs::change_detection::Tick,
14 math::{vec3, vec4},
15 mesh::{
16 Indices, MeshAttributeCompressionFlags, MeshCompressionArgs, MeshVertexBufferLayoutRef,
17 PrimitiveTopology,
18 },
19 pbr::{
20 DrawMesh, MeshPipeline, MeshPipelineKey, MeshPipelineSystems, MeshPipelineViewLayoutKey,
21 RenderMeshInstances, SetMeshBindGroup, SetMeshViewBindGroup, SetMeshViewEmptyBindGroup,
22 ViewKeyCache,
23 },
24 prelude::*,
25 render::{
26 batching::gpu_preprocessing::GpuPreprocessingSupport,
27 camera::{DirtySpecializations, PendingQueues},
28 extract_component::{ExtractComponent, ExtractComponentPlugin},
29 mesh::{allocator::MeshAllocator, RenderMesh},
30 render_asset::RenderAssets,
31 render_phase::{
32 AddRenderCommand, BinnedRenderPhaseType, DrawFunctions, SetItemPipeline,
33 ViewBinnedRenderPhases,
34 },
35 render_resource::{
36 ColorTargetState, ColorWrites, CompareFunction, DepthStencilState, Face, FragmentState,
37 FrontFace, MultisampleState, PipelineCache, PolygonMode, PrimitiveState,
38 RenderPipelineDescriptor, SpecializedMeshPipeline, SpecializedMeshPipelineError,
39 SpecializedMeshPipelines, VertexState,
40 },
41 sync_world::MainEntityHashSet,
42 view::{ExtractedView, RenderVisibleEntities},
43 Render, RenderApp, RenderStartup, RenderSystems,
44 },
45 };
46
47 const SHADER_ASSET_PATH: &str = "shaders/specialized_mesh_pipeline.wesl";
48
49 fn main() {
50 App::new()
51 .add_plugins(DefaultPlugins)
52 .add_plugins(CustomRenderedMeshPipelinePlugin)
53 .add_systems(Startup, setup)
54 .run();
55 }
56
57 /// Spawns the objects in the scene.
58 fn setup(mut commands: Commands, mut meshes: ResMut<Assets<Mesh>>) {
59 // Build a custom triangle mesh with colors
60 // We define a custom mesh because the examples only uses a limited
61 // set of vertex attributes for simplicity
62 let mesh = Mesh::new(
63 PrimitiveTopology::TriangleList,
64 RenderAssetUsages::default(),
65 )
66 .with_inserted_indices(Indices::U32(vec![0, 1, 2]))
67 .with_inserted_attribute(
68 Mesh::ATTRIBUTE_POSITION,
69 vec![
70 vec3(-0.5, -0.5, 0.0),
71 vec3(0.5, -0.5, 0.0),
72 vec3(0.0, 0.25, 0.0),
73 ],
74 )
75 .with_inserted_attribute(
76 Mesh::ATTRIBUTE_COLOR,
77 vec![
78 vec4(1.0, 0.0, 0.0, 1.0),
79 vec4(0.0, 1.0, 0.0, 1.0),
80 vec4(0.0, 0.0, 1.0, 1.0),
81 ],
82 );
83
84 // spawn 3 triangles to show that batching works
85 for (x, y) in [-0.5, 0.0, 0.5].into_iter().zip([-0.25, 0.5, -0.25]) {
86 // Spawn an entity with all the required components for it to be rendered with our custom pipeline
87 commands.spawn((
88 // We use a marker component to identify the mesh that will be rendered
89 // with our specialized pipeline
90 CustomRenderedEntity,
91 // We need to add the mesh handle to the entity
92 Mesh3d(meshes.add(mesh.clone())),
93 Transform::from_xyz(x, y, 0.0),
94 ));
95 }
96 commands.spawn((
97 // We use a marker component to identify the mesh that will be rendered
98 // with our specialized pipeline
99 CustomRenderedEntity,
100 // We need to add the mesh handle to the entity
101 Mesh3d(
102 meshes.add(
103 Rectangle::new(0.5, 0.5)
104 .mesh()
105 .build()
106 .with_removed_attribute(Mesh::ATTRIBUTE_NORMAL)
107 .with_removed_attribute(Mesh::ATTRIBUTE_UV_0)
108 .compressed_mesh(&MeshCompressionArgs::regular())
109 .unwrap(),
110 ),
111 ),
112 Transform::from_xyz(0.0, -1.0 / 3.0, -1.0),
113 ));
114
115 // Spawn the camera.
116 commands.spawn((
117 Camera3d::default(),
118 // Move the camera back a bit to see all the triangles
119 Transform::from_xyz(0.0, 0.0, 3.0).looking_at(Vec3::ZERO, Vec3::Y),
120 ));
121 }
122
123 // When writing custom rendering code it's generally recommended to use a plugin.
124 // The main reason for this is that it gives you access to the finish() hook
125 // which is called after rendering resources are initialized.
126 struct CustomRenderedMeshPipelinePlugin;
127 impl Plugin for CustomRenderedMeshPipelinePlugin {
128 fn build(&self, app: &mut App) {
129 app.add_plugins(ExtractComponentPlugin::<CustomRenderedEntity>::default());
130
131 // We make sure to add these to the render app, not the main app.
132 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
133 return;
134 };
135 render_app
136 // This is needed to tell bevy about your custom pipeline
137 .init_resource::<SpecializedMeshPipelines<CustomMeshPipeline>>()
138 .init_resource::<PendingCustomMeshQueues>()
139 // We need to use a custom draw command so we need to register it
140 .add_render_command::<Opaque3d, DrawSpecializedPipelineCommands>()
141 .add_systems(
142 RenderStartup,
143 init_custom_mesh_pipeline.after(MeshPipelineSystems),
144 )
145 .add_systems(
146 Render,
147 queue_custom_mesh_pipeline.in_set(RenderSystems::Queue),
148 );
149 }
150 }
151
152 /// A marker component that represents an entity that is to be rendered using
153 /// our specialized pipeline.
154 ///
155 /// Note the [`ExtractComponent`] trait implementation: this is necessary to
156 /// tell Bevy that this object should be pulled into the render world. Also note
157 /// the `on_add` hook, which is needed to tell Bevy's `check_visibility` system
158 /// that entities with this component need to be examined for visibility.
159 #[derive(Clone, Component, ExtractComponent)]
160 #[require(VisibilityClass)]
161 #[component(on_add = visibility::add_visibility_class::<CustomRenderedEntity>)]
162 #[extract_app(RenderApp)]
163 struct CustomRenderedEntity;
164
165 /// The custom draw commands that Bevy executes for each entity we enqueue into
166 /// the render phase.
167 type DrawSpecializedPipelineCommands = (
168 // Set the pipeline
169 SetItemPipeline,
170 // Set the view uniform at bind group 0
171 SetMeshViewBindGroup<0>,
172 // Set an empty material bind group at bind group 1
173 SetMeshViewEmptyBindGroup<1>,
174 // Set the mesh uniform at bind group 2
175 SetMeshBindGroup<2>,
176 // Draw the mesh
177 DrawMesh,
178 );
179
180 // This contains the state needed to specialize a mesh pipeline
181 #[derive(Resource)]
182 struct CustomMeshPipeline {
183 /// The base mesh pipeline defined by bevy
184 ///
185 /// This isn't required, but if you want to use a bevy `Mesh` it's easier when you
186 /// have access to the base `MeshPipeline` that bevy already defines
187 mesh_pipeline: MeshPipeline,
188 /// Stores the shader used for this pipeline directly on the pipeline.
189 /// This isn't required, it's only done like this for simplicity.
190 shader_handle: Handle<Shader>,
191 }
192
193 fn init_custom_mesh_pipeline(
194 mut commands: Commands,
195 asset_server: Res<AssetServer>,
196 mesh_pipeline: Res<MeshPipeline>,
197 ) {
198 // Load the shader
199 let shader_handle: Handle<Shader> = asset_server.load(SHADER_ASSET_PATH);
200 commands.insert_resource(CustomMeshPipeline {
201 mesh_pipeline: mesh_pipeline.clone(),
202 shader_handle,
203 });
204 }
205
206 impl SpecializedMeshPipeline for CustomMeshPipeline {
207 /// Pipeline use keys to determine how to specialize it.
208 /// The key is also used by the pipeline cache to determine if
209 /// it needs to create a new pipeline or not
210 ///
211 /// In this example we just use the base `MeshPipelineKey` defined by bevy, but this could be anything.
212 /// For example, if you want to make a pipeline with a procedural shader you could add the Handle<Shader> to the key.
213 type Key = MeshPipelineKey;
214
215 fn specialize(
216 &self,
217 mesh_key: Self::Key,
218 layout: &MeshVertexBufferLayoutRef,
219 ) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError> {
220 let mut shader_defs = Vec::new();
221 // Define the vertex attributes based on a standard bevy [`Mesh`]
222 let mut vertex_attributes = Vec::new();
223 if layout.0.contains(Mesh::ATTRIBUTE_POSITION) {
224 // Handle compressed vertex positions.
225 if layout
226 .0
227 .get_attribute_compression()
228 .contains(MeshAttributeCompressionFlags::COMPRESS_POSITION)
229 {
230 shader_defs.push("VERTEX_POSITIONS_COMPRESSED".into());
231 }
232 // Make sure this matches the shader location
233 vertex_attributes.push(Mesh::ATTRIBUTE_POSITION.at_shader_location(0));
234 }
235 if layout.0.contains(Mesh::ATTRIBUTE_COLOR) {
236 shader_defs.push("VERTEX_COLORS".into());
237 // Make sure this matches the shader location
238 vertex_attributes.push(Mesh::ATTRIBUTE_COLOR.at_shader_location(1));
239 }
240 // This will automatically generate the correct `VertexBufferLayout` based on the vertex attributes
241 let vertex_buffer_layout = layout.0.get_layout(&vertex_attributes)?;
242
243 let view_layout = self
244 .mesh_pipeline
245 .get_view_layout(MeshPipelineViewLayoutKey::from(mesh_key));
246
247 Ok(RenderPipelineDescriptor {
248 label: Some("Specialized Mesh Pipeline".into()),
249 layout: vec![
250 view_layout.main_layout,
251 view_layout.empty_layout,
252 self.mesh_pipeline.mesh_layouts.model_only.clone(),
253 ],
254 vertex: VertexState {
255 shader: self.shader_handle.clone(),
256 shader_defs: shader_defs.clone(),
257 // Customize how to store the meshes' vertex attributes in the vertex buffer
258 buffers: vec![vertex_buffer_layout],
259 ..default()
260 },
261 fragment: Some(FragmentState {
262 shader: self.shader_handle.clone(),
263 shader_defs,
264 targets: vec![Some(ColorTargetState {
265 // This isn't required, but bevy supports rendering different formats
266 // so it's generally recommended to specialize the pipeline for that
267 format: mesh_key.target_format(),
268 // For this example we only use opaque meshes,
269 // but if you wanted to use alpha blending you would need to set it here
270 blend: None,
271 write_mask: ColorWrites::ALL,
272 })],
273 ..default()
274 }),
275 primitive: PrimitiveState {
276 topology: mesh_key.primitive_topology(),
277 strip_index_format: mesh_key.strip_index_format(),
278 front_face: FrontFace::Ccw,
279 cull_mode: Some(Face::Back),
280 polygon_mode: PolygonMode::Fill,
281 ..default()
282 },
283 // Note that if your view has no depth buffer this will need to be
284 // changed.
285 depth_stencil: Some(DepthStencilState {
286 format: CORE_3D_DEPTH_FORMAT,
287 depth_write_enabled: Some(true),
288 depth_compare: Some(CompareFunction::GreaterEqual),
289 stencil: default(),
290 bias: default(),
291 }),
292 // It's generally recommended to specialize your pipeline for MSAA,
293 // but it's not always possible
294 multisample: MultisampleState {
295 count: mesh_key.msaa_samples(),
296 ..default()
297 },
298 ..default()
299 })
300 }
301 }
302
303 /// A resource that stores meshes that couldn't be specialized yet because their
304 /// materials hadn't loaded.
305 ///
306 /// See the documentation for [`PendingQueues`] for more information.
307 #[derive(Default, Deref, DerefMut, Resource)]
308 struct PendingCustomMeshQueues(pub PendingQueues);
309
310 /// A render-world system that enqueues the entity with custom rendering into
311 /// the opaque render phases of each view.
312 fn queue_custom_mesh_pipeline(
313 pipeline_cache: Res<PipelineCache>,
314 custom_mesh_pipeline: Res<CustomMeshPipeline>,
315 (mut opaque_render_phases, opaque_draw_functions): (
316 ResMut<ViewBinnedRenderPhases<Opaque3d>>,
317 Res<DrawFunctions<Opaque3d>>,
318 ),
319 mut specialized_mesh_pipelines: ResMut<SpecializedMeshPipelines<CustomMeshPipeline>>,
320 views: Query<(&RenderVisibleEntities, &ExtractedView)>,
321 view_key_cache: Res<ViewKeyCache>,
322 (render_meshes, render_mesh_instances): (
323 Res<RenderAssets<RenderMesh>>,
324 Res<RenderMeshInstances>,
325 ),
326 mut change_tick: Local<Tick>,
327 mesh_allocator: Res<MeshAllocator>,
328 gpu_preprocessing_support: Res<GpuPreprocessingSupport>,
329 dirty_specializations: Res<DirtySpecializations>,
330 mut pending_custom_mesh_queues: ResMut<PendingCustomMeshQueues>,
331 mut mesh_instances_queued_this_iteration_scratch_space: Local<MainEntityHashSet>,
332 ) {
333 // Get the id for our custom draw function
334 let draw_function = opaque_draw_functions
335 .read()
336 .id::<DrawSpecializedPipelineCommands>();
337
338 // Render phases are per-view, so we need to iterate over all views so that
339 // the entity appears in them. (In this example, we have only one view, but
340 // it's good practice to loop over all views anyway.)
341 for (view_visible_entities, view) in views.iter() {
342 let Some(opaque_phase) = opaque_render_phases.get_mut(&view.retained_view_entity) else {
343 continue;
344 };
345
346 let Some(&view_key) = view_key_cache.get(&view.retained_view_entity) else {
347 continue;
348 };
349
350 let Some(render_visible_mesh_entities) =
351 view_visible_entities.get::<CustomRenderedEntity>()
352 else {
353 continue;
354 };
355
356 // Initialize the pending queues.
357 let view_pending_custom_mesh_queues =
358 pending_custom_mesh_queues.prepare_for_new_frame(view.retained_view_entity);
359
360 // First, remove meshes that need to be respecialized, and those that were removed, from the bins.
361 for &main_entity in dirty_specializations
362 .iter_to_dequeue(view.retained_view_entity, render_visible_mesh_entities)
363 {
364 opaque_phase.remove(main_entity);
365 }
366
367 // Find all the custom rendered entities that are visible from this
368 // view.
369 for (render_entity, visible_entity) in dirty_specializations.iter_to_queue(
370 view.retained_view_entity,
371 render_visible_mesh_entities,
372 &view_pending_custom_mesh_queues.prev_frame,
373 &mut mesh_instances_queued_this_iteration_scratch_space,
374 ) {
375 // Get the mesh instance
376 let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*visible_entity)
377 else {
378 // We couldn't fetch the mesh, probably because it hasn't been
379 // loaded yet. Add the entity to the list of pending custom
380 // meshes and bail.
381 view_pending_custom_mesh_queues
382 .current_frame
383 .insert((*render_entity, *visible_entity));
384 continue;
385 };
386
387 // Get the mesh data
388 let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id()) else {
389 continue;
390 };
391
392 let Some(mesh_slabs) = mesh_allocator.mesh_slabs(&mesh_instance.mesh_asset_id()) else {
393 continue;
394 };
395
396 // Specialize the key for the current mesh entity
397 // For this example we only specialize based on the mesh topology
398 // but you could have more complex keys and that's where you'd need to create those keys
399 let mut mesh_key = view_key;
400 mesh_key |= MeshPipelineKey::from_primitive_topology_and_strip_index(
401 mesh.primitive_topology(),
402 mesh.index_format(),
403 );
404
405 // Finally, we can specialize the pipeline based on the key
406 let pipeline_id = specialized_mesh_pipelines
407 .specialize(
408 &pipeline_cache,
409 &custom_mesh_pipeline,
410 mesh_key,
411 &mesh.layout,
412 )
413 // This should never happen with this example, but if your pipeline
414 // specialization can fail you need to handle the error here
415 .expect("Failed to specialize mesh pipeline");
416
417 // Bump the change tick so that Bevy is forced to rebuild the bin.
418 let next_change_tick = change_tick.get() + 1;
419 change_tick.set(next_change_tick);
420
421 // Add the mesh with our specialized pipeline
422 opaque_phase.add(
423 Opaque3dBatchSetKey {
424 draw_function,
425 pipeline: pipeline_id,
426 material_bind_group_index: None,
427 slabs: mesh_slabs,
428 lightmap_slab: None,
429 },
430 // For this example we can use the mesh asset id as the bin key,
431 // but you can use any asset_id as a key
432 Opaque3dBinKey {
433 asset_id: mesh_instance.mesh_asset_id().into(),
434 },
435 (*render_entity, *visible_entity),
436 mesh_instance.current_uniform_index,
437 // This example supports batching and multi draw indirect,
438 // but if your pipeline doesn't support it you can use
439 // `BinnedRenderPhaseType::UnbatchableMesh`
440 BinnedRenderPhaseType::mesh(
441 mesh_instance.should_batch(),
442 &gpu_preprocessing_support,
443 ),
444 );
445 }
446 }
447 }