docs/bevy/examples/shader_advanced/custom_phase_item.rs 15.3 K · 424 lines · raw · history

1 //! Demonstrates how to enqueue custom draw commands in a render phase.
2 //!
3 //! This example shows how to use the built-in
4 //! [`bevy_render::render_phase::BinnedRenderPhase`] functionality with a
5 //! custom [`RenderCommand`] to allow inserting arbitrary GPU drawing logic
6 //! into Bevy's pipeline. This is not the only way to add custom rendering code
7 //! into Bevy — render nodes are another, lower-level method — but it does allow
8 //! for better reuse of parts of Bevy's built-in mesh rendering logic.
9
10 use bevy::{
11 camera::{
12 primitives::Aabb,
13 visibility::{self, VisibilityClass},
14 },
15 core_pipeline::core_3d::{Opaque3d, Opaque3dBatchSetKey, Opaque3dBinKey, CORE_3D_DEPTH_FORMAT},
16 ecs::{
17 query::ROQueryItem,
18 system::{lifetimeless::SRes, SystemParamItem},
19 },
20 mesh::VertexBufferLayout,
21 prelude::*,
22 render::{
23 camera::{DirtySpecializations, PendingQueues},
24 extract_component::{ExtractComponent, ExtractComponentPlugin},
25 mesh::allocator::MeshSlabs,
26 render_phase::{
27 AddRenderCommand, BinnedRenderPhaseType, DrawFunctions, InputUniformIndex, PhaseItem,
28 RenderCommand, RenderCommandResult, SetItemPipeline, TrackedRenderPass,
29 ViewBinnedRenderPhases,
30 },
31 render_resource::{
32 BufferUsages, Canonical, ColorTargetState, ColorWrites, CompareFunction,
33 DepthStencilState, FragmentState, IndexFormat, PipelineCache, RawBufferVec,
34 RenderPipeline, RenderPipelineDescriptor, Specializer, SpecializerKey, TextureFormat,
35 Variants, VertexAttribute, VertexFormat, VertexState, VertexStepMode,
36 },
37 renderer::{RenderDevice, RenderQueue},
38 sync_world::MainEntityHashSet,
39 view::{ExtractedView, RenderVisibleEntities},
40 Render, RenderApp, RenderSystems,
41 },
42 };
43 use bytemuck::{Pod, Zeroable};
44
45 /// A marker component that represents an entity that is to be rendered using
46 /// our custom phase item.
47 ///
48 /// Note the [`ExtractComponent`] trait implementation: this is necessary to
49 /// tell Bevy that this object should be pulled into the render world. Also note
50 /// the `on_add` hook, which is needed to tell Bevy's `check_visibility` system
51 /// that entities with this component need to be examined for visibility.
52 #[derive(Clone, Component, ExtractComponent)]
53 #[require(VisibilityClass)]
54 #[component(on_add = visibility::add_visibility_class::<CustomRenderedEntity>)]
55 #[extract_app(RenderApp)]
56 struct CustomRenderedEntity;
57
58 /// A [`RenderCommand`] that binds the vertex and index buffers and issues the
59 /// draw command for our custom phase item.
60 struct DrawCustomPhaseItem;
61
62 impl<P> RenderCommand<P> for DrawCustomPhaseItem
63 where
64 P: PhaseItem,
65 {
66 type Param = SRes<CustomPhaseItemBuffers>;
67
68 type ViewQuery = ();
69
70 type ItemQuery = ();
71
72 fn render<'w>(
73 _: &P,
74 _: ROQueryItem<'w, '_, Self::ViewQuery>,
75 _: Option<ROQueryItem<'w, '_, Self::ItemQuery>>,
76 custom_phase_item_buffers: SystemParamItem<'w, '_, Self::Param>,
77 pass: &mut TrackedRenderPass<'w>,
78 ) -> RenderCommandResult {
79 // Borrow check workaround.
80 let custom_phase_item_buffers = custom_phase_item_buffers.into_inner();
81
82 // Tell the GPU where the vertices are.
83 pass.set_vertex_buffer(
84 0,
85 custom_phase_item_buffers
86 .vertices
87 .buffer()
88 .unwrap()
89 .slice(..),
90 );
91
92 // Tell the GPU where the indices are.
93 pass.set_index_buffer(
94 custom_phase_item_buffers
95 .indices
96 .buffer()
97 .unwrap()
98 .slice(..),
99 IndexFormat::Uint32,
100 );
101
102 // Draw one triangle (3 vertices).
103 pass.draw_indexed(0..3, 0, 0..1);
104
105 RenderCommandResult::Success
106 }
107 }
108
109 /// The GPU vertex and index buffers for our custom phase item.
110 ///
111 /// As the custom phase item is a single triangle, these are uploaded once and
112 /// then left alone.
113 #[derive(Resource)]
114 struct CustomPhaseItemBuffers {
115 /// The vertices for the single triangle.
116 ///
117 /// This is a [`RawBufferVec`] because that's the simplest and fastest type
118 /// of GPU buffer, and [`Vertex`] objects are simple.
119 vertices: RawBufferVec<Vertex>,
120
121 /// The indices of the single triangle.
122 ///
123 /// As above, this is a [`RawBufferVec`] because `u32` values have trivial
124 /// size and alignment.
125 indices: RawBufferVec<u32>,
126 }
127
128 /// The CPU-side structure that describes a single vertex of the triangle.
129 #[derive(Clone, Copy, Pod, Zeroable)]
130 #[repr(C)]
131 struct Vertex {
132 /// The 3D position of the triangle vertex.
133 position: Vec3,
134 /// Padding.
135 pad0: u32,
136 /// The color of the triangle vertex.
137 color: Vec3,
138 /// Padding.
139 pad1: u32,
140 }
141
142 impl Vertex {
143 /// Creates a new vertex structure.
144 const fn new(position: Vec3, color: Vec3) -> Vertex {
145 Vertex {
146 position,
147 color,
148 pad0: 0,
149 pad1: 0,
150 }
151 }
152 }
153
154 /// The custom draw commands that Bevy executes for each entity we enqueue into
155 /// the render phase.
156 type DrawCustomPhaseItemCommands = (SetItemPipeline, DrawCustomPhaseItem);
157
158 /// A single triangle's worth of vertices, for demonstration purposes.
159 static VERTICES: [Vertex; 3] = [
160 Vertex::new(vec3(-0.866, -0.5, 0.5), vec3(1.0, 0.0, 0.0)),
161 Vertex::new(vec3(0.866, -0.5, 0.5), vec3(0.0, 1.0, 0.0)),
162 Vertex::new(vec3(0.0, 1.0, 0.5), vec3(0.0, 0.0, 1.0)),
163 ];
164
165 /// The entry point.
166 fn main() {
167 let mut app = App::new();
168 app.add_plugins(DefaultPlugins)
169 .add_plugins(ExtractComponentPlugin::<CustomRenderedEntity>::default())
170 .add_systems(Startup, setup);
171
172 // We make sure to add these to the render app, not the main app.
173 app.sub_app_mut(RenderApp)
174 .init_resource::<CustomPhasePipeline>()
175 .init_resource::<PendingCustomPhaseItemQueues>()
176 .add_render_command::<Opaque3d, DrawCustomPhaseItemCommands>()
177 .add_systems(
178 Render,
179 prepare_custom_phase_item_buffers.in_set(RenderSystems::Prepare),
180 )
181 .add_systems(Render, queue_custom_phase_item.in_set(RenderSystems::Queue));
182
183 app.run();
184 }
185
186 /// Spawns the objects in the scene.
187 fn setup(mut commands: Commands) {
188 // Spawn a single entity that has custom rendering. It'll be extracted into
189 // the render world via [`ExtractComponent`].
190 commands.spawn((
191 Visibility::default(),
192 Transform::default(),
193 // This `Aabb` is necessary for the visibility checks to work.
194 Aabb {
195 center: Vec3A::ZERO,
196 half_extents: Vec3A::splat(0.5),
197 },
198 CustomRenderedEntity,
199 ));
200
201 // Spawn the camera.
202 commands.spawn((
203 Camera3d::default(),
204 Transform::from_xyz(0.0, 0.0, 1.0).looking_at(Vec3::ZERO, Vec3::Y),
205 ));
206 }
207
208 /// Creates the [`CustomPhaseItemBuffers`] resource.
209 ///
210 /// This must be done in a startup system because it needs the [`RenderDevice`]
211 /// and [`RenderQueue`] to exist, and they don't until [`App::run`] is called.
212 fn prepare_custom_phase_item_buffers(mut commands: Commands) {
213 commands.init_resource::<CustomPhaseItemBuffers>();
214 }
215
216 /// A resource that holds entities that couldn't be specialized and/or queued
217 /// yet because their dependent assets haven't loaded yet.
218 ///
219 /// In this particular example, entities with custom rendering can always be
220 /// specialized, so this resource goes unused in practice. However, we still
221 /// need it, because [`DirtySpecializations`] requires such a resource.
222 ///
223 /// See the documentation of [`PendingQueues`] for more information.
224 #[derive(Default, Deref, DerefMut, Resource)]
225 pub struct PendingCustomPhaseItemQueues(pub PendingQueues);
226
227 /// A render-world system that enqueues the entity with custom rendering into
228 /// the opaque render phases of each view.
229 fn queue_custom_phase_item(
230 pipeline_cache: Res<PipelineCache>,
231 mut pipeline: ResMut<CustomPhasePipeline>,
232 mut opaque_render_phases: ResMut<ViewBinnedRenderPhases<Opaque3d>>,
233 opaque_draw_functions: Res<DrawFunctions<Opaque3d>>,
234 views: Query<(&ExtractedView, &RenderVisibleEntities, &Msaa)>,
235 dirty_specializations: Res<DirtySpecializations>,
236 mut pending_custom_phase_item_queues: ResMut<PendingCustomPhaseItemQueues>,
237 mut mesh_instances_queued_this_iteration_scratch_space: Local<MainEntityHashSet>,
238 ) {
239 let draw_custom_phase_item = opaque_draw_functions
240 .read()
241 .id::<DrawCustomPhaseItemCommands>();
242
243 // Render phases are per-view, so we need to iterate over all views so that
244 // the entity appears in them. (In this example, we have only one view, but
245 // it's good practice to loop over all views anyway.)
246 for (view, view_visible_entities, msaa) in views.iter() {
247 let Some(opaque_phase) = opaque_render_phases.get_mut(&view.retained_view_entity) else {
248 continue;
249 };
250
251 // Fetch the list of visible entities in the `CustomRenderedEntity`
252 // class. If there are no such entities, then we have no entities to
253 // render, and we're done.
254 let Some(render_visible_mesh_entities) =
255 view_visible_entities.get::<CustomRenderedEntity>()
256 else {
257 continue;
258 };
259
260 let view_pending_custom_phase_item_queues =
261 pending_custom_phase_item_queues.prepare_for_new_frame(view.retained_view_entity);
262
263 // First, remove meshes that need to be respecialized, and those that
264 // were removed, from the bins.
265 for &main_entity in dirty_specializations
266 .iter_to_dequeue(view.retained_view_entity, render_visible_mesh_entities)
267 {
268 opaque_phase.remove(main_entity);
269 }
270
271 // Find all the custom rendered entities that are visible from this
272 // view.
273 for (render_entity, main_entity) in dirty_specializations.iter_to_queue(
274 view.retained_view_entity,
275 render_visible_mesh_entities,
276 &view_pending_custom_phase_item_queues.prev_frame,
277 &mut mesh_instances_queued_this_iteration_scratch_space,
278 ) {
279 // Ordinarily, the [`SpecializedRenderPipeline::Key`] would contain
280 // some per-view settings, such as whether the view is HDR, but for
281 // simplicity's sake we simply hard-code the view's characteristics,
282 // with the exception of number of MSAA samples.
283 let Ok(pipeline_id) = pipeline
284 .variants
285 .specialize(&pipeline_cache, CustomPhaseKey(*msaa))
286 else {
287 continue;
288 };
289
290 // Add the custom render item. We use the
291 // [`BinnedRenderPhaseType::NonMesh`] type to skip the special
292 // handling that Bevy has for meshes (preprocessing, indirect
293 // draws, etc.)
294 //
295 // The asset ID is arbitrary; we simply use [`AssetId::default`],
296 // but you can use anything you like. Note that the asset ID need
297 // not be the ID of a [`Mesh`].
298 opaque_phase.add(
299 Opaque3dBatchSetKey {
300 draw_function: draw_custom_phase_item,
301 pipeline: pipeline_id,
302 material_bind_group_index: None,
303 lightmap_slab: None,
304 slabs: MeshSlabs::default(),
305 },
306 Opaque3dBinKey {
307 asset_id: AssetId::<Mesh>::default().untyped(),
308 },
309 (*render_entity, *main_entity),
310 InputUniformIndex::default(),
311 BinnedRenderPhaseType::NonMesh,
312 );
313 }
314 }
315 }
316
317 struct CustomPhaseSpecializer;
318
319 #[derive(Resource)]
320 struct CustomPhasePipeline {
321 /// the `variants` collection holds onto the shader handle through the base descriptor
322 variants: Variants<RenderPipeline, CustomPhaseSpecializer>,
323 }
324
325 impl FromWorld for CustomPhasePipeline {
326 fn from_world(world: &mut World) -> Self {
327 let asset_server = world.resource::<AssetServer>();
328 let shader = asset_server.load("shaders/custom_phase_item.wesl");
329
330 let base_descriptor = RenderPipelineDescriptor {
331 label: Some("custom render pipeline".into()),
332 vertex: VertexState {
333 shader: shader.clone(),
334 buffers: vec![VertexBufferLayout {
335 array_stride: size_of::<Vertex>() as u64,
336 step_mode: VertexStepMode::Vertex,
337 // This needs to match the layout of [`Vertex`].
338 attributes: vec![
339 VertexAttribute {
340 format: VertexFormat::Float32x3,
341 offset: 0,
342 shader_location: 0,
343 },
344 VertexAttribute {
345 format: VertexFormat::Float32x3,
346 offset: 16,
347 shader_location: 1,
348 },
349 ],
350 }],
351 ..default()
352 },
353 fragment: Some(FragmentState {
354 shader: shader.clone(),
355 targets: vec![Some(ColorTargetState {
356 // Ordinarily, you'd want to check whether the view has the
357 // HDR format and substitute the appropriate texture format
358 // here, but we omit that for simplicity.
359 format: TextureFormat::Rgba8UnormSrgb,
360 blend: None,
361 write_mask: ColorWrites::ALL,
362 })],
363 ..default()
364 }),
365 // Note that if your view has no depth buffer this will need to be
366 // changed.
367 depth_stencil: Some(DepthStencilState {
368 format: CORE_3D_DEPTH_FORMAT,
369 depth_write_enabled: Some(false),
370 depth_compare: Some(CompareFunction::Always),
371 stencil: default(),
372 bias: default(),
373 }),
374 ..default()
375 };
376
377 let variants = Variants::new(CustomPhaseSpecializer, base_descriptor);
378
379 Self { variants }
380 }
381 }
382
383 #[derive(Copy, Clone, PartialEq, Eq, Hash, SpecializerKey)]
384 struct CustomPhaseKey(Msaa);
385
386 impl Specializer<RenderPipeline> for CustomPhaseSpecializer {
387 type Key = CustomPhaseKey;
388
389 fn specialize(
390 &self,
391 key: Self::Key,
392 descriptor: &mut RenderPipelineDescriptor,
393 ) -> Result<Canonical<Self::Key>, BevyError> {
394 descriptor.multisample.count = key.0.samples();
395 Ok(key)
396 }
397 }
398
399 impl FromWorld for CustomPhaseItemBuffers {
400 fn from_world(world: &mut World) -> Self {
401 let render_device = world.resource::<RenderDevice>();
402 let render_queue = world.resource::<RenderQueue>();
403
404 // Create the vertex and index buffers.
405 let mut vbo = RawBufferVec::new(BufferUsages::VERTEX);
406 let mut ibo = RawBufferVec::new(BufferUsages::INDEX);
407
408 for vertex in &VERTICES {
409 vbo.push(*vertex);
410 }
411 for index in 0..3 {
412 ibo.push(index);
413 }
414
415 // These two lines are required in order to trigger the upload to GPU.
416 vbo.write_buffer(render_device, render_queue);
417 ibo.write_buffer(render_device, render_queue);
418
419 CustomPhaseItemBuffers {
420 vertices: vbo,
421 indices: ibo,
422 }
423 }
424 }