diff --git a/Cargo.toml b/Cargo.toml index 4b101ba572..383cb50aad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -518,6 +518,7 @@ http-body-util = "0.1" anyhow = "1" macro_rules_attribute = "0.2" accesskit = "0.17" +nonmax = "0.5" [target.'cfg(not(target_family = "wasm"))'.dev-dependencies] smol = "2" @@ -2728,6 +2729,18 @@ description = "A shader that renders a mesh multiple times in one draw call usin category = "Shaders" wasm = true +[[example]] +name = "custom_render_phase" +path = "examples/shader/custom_render_phase.rs" +doc-scrape-examples = true + +[package.metadata.example.custom_render_phase] +name = "Custom Render Phase" +description = "Shows how to make a complete render phase" +category = "Shaders" +wasm = true + + [[example]] name = "automatic_instancing" path = "examples/shader/automatic_instancing.rs" diff --git a/assets/shaders/custom_stencil.wgsl b/assets/shaders/custom_stencil.wgsl new file mode 100644 index 0000000000..6f2fa2da4f --- /dev/null +++ b/assets/shaders/custom_stencil.wgsl @@ -0,0 +1,41 @@ +//! A shader showing how to use the vertex position data to output the +//! stencil in the right position + +// First we import everything we need from bevy_pbr +// A 2d shader would be vevry similar but import from bevy_sprite instead +#import bevy_pbr::{ + mesh_functions, + view_transformations::position_world_to_clip +} + +struct Vertex { + // This is needed if you are using batching and/or gpu preprocessing + // It's a built in so you don't need to define it in the vertex layout + @builtin(instance_index) instance_index: u32, + // Like we defined for the vertex layout + // position is at location 0 + @location(0) position: vec3, +}; + +// This is the output of the vertex shader and we also use it as the input for the fragment shader +struct VertexOutput { + @builtin(position) clip_position: vec4, + @location(0) world_position: vec4, +}; + +@vertex +fn vertex(vertex: Vertex) -> VertexOutput { + var out: VertexOutput; + // This is how bevy computes the world position + // The vertex.instance_index is very important. Especially if you are using batching and gpu preprocessing + var world_from_local = mesh_functions::get_world_from_local(vertex.instance_index); + out.world_position = mesh_functions::mesh_position_local_to_world(world_from_local, vec4(vertex.position, 1.0)); + out.clip_position = position_world_to_clip(out.world_position.xyz); + return out; +} + +@fragment +fn fragment(in: VertexOutput) -> @location(0) vec4 { + // Output a red color to represent the stencil of the mesh + return vec4(1.0, 0.0, 0.0, 1.0); +} diff --git a/examples/README.md b/examples/README.md index fef76c1a19..4569d0ceac 100644 --- a/examples/README.md +++ b/examples/README.md @@ -434,6 +434,7 @@ Example | Description [Animated](../examples/shader/animate_shader.rs) | A shader that uses dynamic data like the time since startup [Array Texture](../examples/shader/array_texture.rs) | A shader that shows how to reuse the core bevy PBR shading functionality in a custom material that obtains the base color from an array texture. [Compute - Game of Life](../examples/shader/compute_shader_game_of_life.rs) | A compute shader that simulates Conway's Game of Life +[Custom Render Phase](../examples/shader/custom_render_phase.rs) | Shows how to make a complete render phase [Custom Vertex Attribute](../examples/shader/custom_vertex_attribute.rs) | A shader that reads a mesh's custom vertex attribute [Custom phase item](../examples/shader/custom_phase_item.rs) | Demonstrates how to enqueue custom draw commands in a render phase [Extended Material](../examples/shader/extended_material.rs) | A custom shader that builds on the standard material diff --git a/examples/shader/custom_render_phase.rs b/examples/shader/custom_render_phase.rs new file mode 100644 index 0000000000..4fbf105fe9 --- /dev/null +++ b/examples/shader/custom_render_phase.rs @@ -0,0 +1,629 @@ +//! This example demonstrates how to write a custom phase +//! +//! Render phases in bevy are used whenever you need to draw a group of meshes in a specific way. +//! For example, bevy's main pass has an opaque phase, a transparent phase for both 2d and 3d. +//! Sometimes, you may want to only draw a subset of meshes before or after the builtin phase. In +//! those situations you need to write your own phase. +//! +//! This example showcases how writing a custom phase to draw a stencil of a bevy mesh could look +//! like. Some shortcuts have been used for simplicity. +//! +//! This example was made for 3d, but a 2d equivalent would be almost identical. + +use std::ops::Range; + +use bevy::{ + core_pipeline::core_3d::graph::{Core3d, Node3d}, + ecs::{ + query::QueryItem, + system::{lifetimeless::SRes, SystemParamItem}, + }, + math::FloatOrd, + pbr::{ + DrawMesh, MeshInputUniform, MeshPipeline, MeshPipelineKey, MeshPipelineViewLayoutKey, + MeshUniform, RenderMeshInstances, SetMeshBindGroup, SetMeshViewBindGroup, + }, + platform_support::collections::HashSet, + prelude::*, + render::{ + batching::{ + gpu_preprocessing::{ + batch_and_prepare_sorted_render_phase, IndirectParametersMetadata, + }, + GetBatchData, GetFullBatchData, + }, + camera::ExtractedCamera, + extract_component::{ExtractComponent, ExtractComponentPlugin}, + mesh::{allocator::MeshAllocator, MeshVertexBufferLayoutRef, RenderMesh}, + render_asset::RenderAssets, + render_graph::{ + NodeRunError, RenderGraphApp, RenderGraphContext, RenderLabel, ViewNode, ViewNodeRunner, + }, + render_phase::{ + sort_phase_system, AddRenderCommand, CachedRenderPipelinePhaseItem, DrawFunctionId, + DrawFunctions, PhaseItem, PhaseItemExtraIndex, SetItemPipeline, SortedPhaseItem, + ViewSortedRenderPhases, + }, + render_resource::{ + CachedRenderPipelineId, ColorTargetState, ColorWrites, Face, FragmentState, FrontFace, + MultisampleState, PipelineCache, PolygonMode, PrimitiveState, RenderPassDescriptor, + RenderPipelineDescriptor, SpecializedMeshPipeline, SpecializedMeshPipelineError, + SpecializedMeshPipelines, TextureFormat, VertexState, + }, + renderer::RenderContext, + sync_world::MainEntity, + view::{ExtractedView, RenderVisibleEntities, RetainedViewEntity, ViewTarget}, + Extract, Render, RenderApp, RenderSet, + }, +}; +use nonmax::NonMaxU32; + +const SHADER_ASSET_PATH: &str = "shaders/custom_stencil.wgsl"; + +fn main() { + App::new() + .add_plugins((DefaultPlugins, MeshStencilPhasePlugin)) + .add_systems(Startup, setup) + .run(); +} + +fn setup( + mut commands: Commands, + mut meshes: ResMut>, + mut materials: ResMut>, +) { + // circular base + commands.spawn(( + Mesh3d(meshes.add(Circle::new(4.0))), + MeshMaterial3d(materials.add(Color::WHITE)), + Transform::from_rotation(Quat::from_rotation_x(-std::f32::consts::FRAC_PI_2)), + )); + // cube + // This cube will be rendered by the main pass, but it will also be rendered by our custom + // pass. This should result in an unlit red cube + commands.spawn(( + Mesh3d(meshes.add(Cuboid::new(1.0, 1.0, 1.0))), + MeshMaterial3d(materials.add(Color::srgb_u8(124, 144, 255))), + Transform::from_xyz(0.0, 0.5, 0.0), + // This marker component is used to identify which mesh will be used in our custom pass + // The circle doesn't have it so it won't be rendered in our pass + DrawStencil, + )); + // light + commands.spawn(( + PointLight { + shadows_enabled: true, + ..default() + }, + Transform::from_xyz(4.0, 8.0, 4.0), + )); + // camera + commands.spawn(( + Camera3d::default(), + Transform::from_xyz(-2.0, 4.5, 9.0).looking_at(Vec3::ZERO, Vec3::Y), + // disable msaa for simplicity + Msaa::Off, + )); +} + +#[derive(Component, ExtractComponent, Clone, Copy, Default)] +struct DrawStencil; + +struct MeshStencilPhasePlugin; +impl Plugin for MeshStencilPhasePlugin { + fn build(&self, app: &mut App) { + app.add_plugins((ExtractComponentPlugin::::default(),)); + // We need to get the render app from the main app + let Some(render_app) = app.get_sub_app_mut(RenderApp) else { + return; + }; + render_app + .init_resource::>() + .init_resource::>() + .add_render_command::() + .init_resource::>() + .add_systems(ExtractSchedule, extract_camera_phases) + .add_systems( + Render, + ( + queue_custom_meshes.in_set(RenderSet::QueueMeshes), + sort_phase_system::.in_set(RenderSet::PhaseSort), + batch_and_prepare_sorted_render_phase:: + .in_set(RenderSet::PrepareResources), + ), + ); + + render_app + .add_render_graph_node::>(Core3d, CustomDrawPassLabel) + // Tell the node to run after the main pass + .add_render_graph_edges(Core3d, (Node3d::MainOpaquePass, CustomDrawPassLabel)); + } + + fn finish(&self, app: &mut App) { + // We need to get the render app from the main app + let Some(render_app) = app.get_sub_app_mut(RenderApp) else { + return; + }; + // The pipeline needs the RenderDevice to be created and it's only available once plugins + // are initialized + render_app.init_resource::(); + } +} + +#[derive(Resource)] +struct StencilPipeline { + /// The base mesh pipeline defined by bevy + /// + /// Since we want to draw a stencil of an existing bevy mesh we want to reuse the default + /// pipeline as much as possible + mesh_pipeline: MeshPipeline, + /// Stores the shader used for this pipeline directly on the pipeline. + /// This isn't required, it's only done like this for simplicity. + shader_handle: Handle, +} +impl FromWorld for StencilPipeline { + fn from_world(world: &mut World) -> Self { + Self { + mesh_pipeline: MeshPipeline::from_world(world), + shader_handle: world.resource::().load(SHADER_ASSET_PATH), + } + } +} + +// For more information on how SpecializedMeshPipeline work, please look at the +// specialized_mesh_pipeline example +impl SpecializedMeshPipeline for StencilPipeline { + type Key = MeshPipelineKey; + + fn specialize( + &self, + key: Self::Key, + layout: &MeshVertexBufferLayoutRef, + ) -> Result { + // We will only use the position of the mesh in our shader so we only need to specify that + let mut vertex_attributes = Vec::new(); + if layout.0.contains(Mesh::ATTRIBUTE_POSITION) { + // Make sure this matches the shader location + vertex_attributes.push(Mesh::ATTRIBUTE_POSITION.at_shader_location(0)); + } + // This will automatically generate the correct `VertexBufferLayout` based on the vertex attributes + let vertex_buffer_layout = layout.0.get_layout(&vertex_attributes)?; + + Ok(RenderPipelineDescriptor { + label: Some("Specialized Mesh Pipeline".into()), + // We want to reuse the data from bevy so we use the same bind groups as the default + // mesh pipeline + layout: vec![ + // Bind group 0 is the view uniform + self.mesh_pipeline + .get_view_layout(MeshPipelineViewLayoutKey::from(key)) + .clone(), + // Bind group 1 is the mesh uniform + self.mesh_pipeline.mesh_layouts.model_only.clone(), + ], + push_constant_ranges: vec![], + vertex: VertexState { + shader: self.shader_handle.clone(), + shader_defs: vec![], + entry_point: "vertex".into(), + buffers: vec![vertex_buffer_layout], + }, + fragment: Some(FragmentState { + shader: self.shader_handle.clone(), + shader_defs: vec![], + entry_point: "fragment".into(), + targets: vec![Some(ColorTargetState { + format: TextureFormat::bevy_default(), + blend: None, + write_mask: ColorWrites::ALL, + })], + }), + primitive: PrimitiveState { + topology: key.primitive_topology(), + front_face: FrontFace::Ccw, + cull_mode: Some(Face::Back), + polygon_mode: PolygonMode::Fill, + ..default() + }, + depth_stencil: None, + // It's generally recommended to specialize your pipeline for MSAA, + // but it's not always possible + multisample: MultisampleState::default(), + zero_initialize_workgroup_memory: false, + }) + } +} + +// We will reuse render commands already defined by bevy to draw a 3d mesh +type DrawMesh3dStencil = ( + SetItemPipeline, + // This will set the view bindings in group 0 + SetMeshViewBindGroup<0>, + // This will set the mesh bindings in group 1 + SetMeshBindGroup<1>, + // This will draw the mesh + DrawMesh, +); + +// This is the data required per entity drawn in a custom phase in bevy. More specifically this is the +// data required when using a ViewSortedRenderPhase. This would look differently if we wanted a +// batched render phase. Sorted phases are a bit easier to implement, but a batched phase would +// look similar. +// +// If you want to see how a batched phase implementation looks, you should look at the Opaque2d +// phase. +struct Stencil3d { + pub sort_key: FloatOrd, + pub entity: (Entity, MainEntity), + pub pipeline: CachedRenderPipelineId, + pub draw_function: DrawFunctionId, + pub batch_range: Range, + pub extra_index: PhaseItemExtraIndex, + /// Whether the mesh in question is indexed (uses an index buffer in + /// addition to its vertex buffer). + pub indexed: bool, +} + +// For more information about writing a phase item, please look at the custom_phase_item example +impl PhaseItem for Stencil3d { + #[inline] + fn entity(&self) -> Entity { + self.entity.0 + } + + #[inline] + fn main_entity(&self) -> MainEntity { + self.entity.1 + } + + #[inline] + fn draw_function(&self) -> DrawFunctionId { + self.draw_function + } + + #[inline] + fn batch_range(&self) -> &Range { + &self.batch_range + } + + #[inline] + fn batch_range_mut(&mut self) -> &mut Range { + &mut self.batch_range + } + + #[inline] + fn extra_index(&self) -> PhaseItemExtraIndex { + self.extra_index.clone() + } + + #[inline] + fn batch_range_and_extra_index_mut(&mut self) -> (&mut Range, &mut PhaseItemExtraIndex) { + (&mut self.batch_range, &mut self.extra_index) + } +} + +impl SortedPhaseItem for Stencil3d { + type SortKey = FloatOrd; + + #[inline] + fn sort_key(&self) -> Self::SortKey { + self.sort_key + } + + #[inline] + fn sort(items: &mut [Self]) { + // bevy normally uses radsort instead of the std slice::sort_by_key + // radsort is a stable radix sort that performed better than `slice::sort_by_key` or `slice::sort_unstable_by_key`. + // Since it is not re-exported by bevy, we just use the std sort for the purpose of the example + items.sort_by_key(SortedPhaseItem::sort_key); + } + + #[inline] + fn indexed(&self) -> bool { + self.indexed + } +} + +impl CachedRenderPipelinePhaseItem for Stencil3d { + #[inline] + fn cached_pipeline(&self) -> CachedRenderPipelineId { + self.pipeline + } +} + +impl GetBatchData for StencilPipeline { + type Param = ( + SRes, + SRes>, + SRes, + ); + type CompareData = AssetId; + type BufferData = MeshUniform; + + fn get_batch_data( + (mesh_instances, _render_assets, mesh_allocator): &SystemParamItem, + (_entity, main_entity): (Entity, MainEntity), + ) -> Option<(Self::BufferData, Option)> { + let RenderMeshInstances::CpuBuilding(ref mesh_instances) = **mesh_instances else { + error!( + "`get_batch_data` should never be called in GPU mesh uniform \ + building mode" + ); + return None; + }; + let mesh_instance = mesh_instances.get(&main_entity)?; + let first_vertex_index = + match mesh_allocator.mesh_vertex_slice(&mesh_instance.mesh_asset_id) { + Some(mesh_vertex_slice) => mesh_vertex_slice.range.start, + None => 0, + }; + let mesh_uniform = { + let mesh_transforms = &mesh_instance.transforms; + let (local_from_world_transpose_a, local_from_world_transpose_b) = + mesh_transforms.world_from_local.inverse_transpose_3x3(); + MeshUniform { + world_from_local: mesh_transforms.world_from_local.to_transpose(), + previous_world_from_local: mesh_transforms.previous_world_from_local.to_transpose(), + lightmap_uv_rect: UVec2::ZERO, + local_from_world_transpose_a, + local_from_world_transpose_b, + flags: mesh_transforms.flags, + first_vertex_index, + current_skin_index: u32::MAX, + previous_skin_index: u32::MAX, + material_and_lightmap_bind_group_slot: 0, + } + }; + Some((mesh_uniform, None)) + } +} +impl GetFullBatchData for StencilPipeline { + type BufferInputData = MeshInputUniform; + + fn get_index_and_compare_data( + (mesh_instances, _, _): &SystemParamItem, + main_entity: MainEntity, + ) -> Option<(NonMaxU32, Option)> { + // This should only be called during GPU building. + let RenderMeshInstances::GpuBuilding(ref mesh_instances) = **mesh_instances else { + error!( + "`get_index_and_compare_data` should never be called in CPU mesh uniform building \ + mode" + ); + return None; + }; + let mesh_instance = mesh_instances.get(&main_entity)?; + Some(( + mesh_instance.current_uniform_index, + mesh_instance + .should_batch() + .then_some(mesh_instance.mesh_asset_id), + )) + } + + fn get_binned_batch_data( + (mesh_instances, _render_assets, mesh_allocator): &SystemParamItem, + main_entity: MainEntity, + ) -> Option { + let RenderMeshInstances::CpuBuilding(ref mesh_instances) = **mesh_instances else { + error!( + "`get_binned_batch_data` should never be called in GPU mesh uniform building mode" + ); + return None; + }; + let mesh_instance = mesh_instances.get(&main_entity)?; + let first_vertex_index = + match mesh_allocator.mesh_vertex_slice(&mesh_instance.mesh_asset_id) { + Some(mesh_vertex_slice) => mesh_vertex_slice.range.start, + None => 0, + }; + + Some(MeshUniform::new( + &mesh_instance.transforms, + first_vertex_index, + mesh_instance.material_bindings_index.slot, + None, + None, + None, + )) + } + + fn write_batch_indirect_parameters_metadata( + mesh_index: u32, + indexed: bool, + base_output_index: u32, + batch_set_index: Option, + indirect_parameters_buffers: &mut bevy_render::batching::gpu_preprocessing::IndirectParametersBuffers, + indirect_parameters_offset: u32, + ) { + // Note that `IndirectParameters` covers both of these structures, even + // though they actually have distinct layouts. See the comment above that + // type for more information. + let indirect_parameters = IndirectParametersMetadata { + mesh_index, + base_output_index, + batch_set_index: match batch_set_index { + None => !0, + Some(batch_set_index) => u32::from(batch_set_index), + }, + early_instance_count: 0, + late_instance_count: 0, + }; + + if indexed { + indirect_parameters_buffers + .set_indexed(indirect_parameters_offset, indirect_parameters); + } else { + indirect_parameters_buffers + .set_non_indexed(indirect_parameters_offset, indirect_parameters); + } + } + + fn get_binned_index( + _param: &SystemParamItem, + _query_item: MainEntity, + ) -> Option { + None + } +} + +// When defining a phase, we need to extract it from the main world and add it to a resource +// that will be used by the render world. We need to give that resource all views that will use +// that phase +fn extract_camera_phases( + mut stencil_phases: ResMut>, + cameras: Extract>>, + mut live_entities: Local>, +) { + live_entities.clear(); + for (main_entity, camera) in &cameras { + if !camera.is_active { + continue; + } + // This is the main camera, so we use the first subview index (0) + let retained_view_entity = RetainedViewEntity::new(main_entity.into(), None, 0); + + stencil_phases.insert_or_clear(retained_view_entity); + live_entities.insert(retained_view_entity); + } + + // Clear out all dead views. + stencil_phases.retain(|camera_entity, _| live_entities.contains(camera_entity)); +} + +// This is a very important step when writing a custom phase. +// +// This system determines which meshes will be added to the phase. +fn queue_custom_meshes( + custom_draw_functions: Res>, + mut pipelines: ResMut>, + pipeline_cache: Res, + custom_draw_pipeline: Res, + render_meshes: Res>, + render_mesh_instances: Res, + mut custom_render_phases: ResMut>, + mut views: Query<(&ExtractedView, &RenderVisibleEntities, &Msaa)>, + has_marker: Query<(), With>, +) { + for (view, visible_entities, msaa) in &mut views { + let Some(custom_phase) = custom_render_phases.get_mut(&view.retained_view_entity) else { + continue; + }; + let draw_custom = custom_draw_functions.read().id::(); + + // Create the key based on the view. + // In this case we only care about MSAA and HDR + let view_key = MeshPipelineKey::from_msaa_samples(msaa.samples()) + | MeshPipelineKey::from_hdr(view.hdr); + + let rangefinder = view.rangefinder3d(); + // Since our phase can work on any 3d mesh we can reuse the default mesh 3d filter + for (render_entity, visible_entity) in visible_entities.iter::() { + // We only want meshes with the marker component to be queued to our phase. + if has_marker.get(*render_entity).is_err() { + continue; + } + let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*visible_entity) + else { + continue; + }; + let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id) else { + continue; + }; + + // Specialize the key for the current mesh entity + // For this example we only specialize based on the mesh topology + // but you could have more complex keys and that's where you'd need to create those keys + let mut mesh_key = view_key; + mesh_key |= MeshPipelineKey::from_primitive_topology(mesh.primitive_topology()); + + let pipeline_id = pipelines.specialize( + &pipeline_cache, + &custom_draw_pipeline, + mesh_key, + &mesh.layout, + ); + let pipeline_id = match pipeline_id { + Ok(id) => id, + Err(err) => { + error!("{}", err); + continue; + } + }; + let distance = rangefinder.distance_translation(&mesh_instance.translation); + // At this point we have all the data we need to create a phase item and add it to our + // phase + custom_phase.add(Stencil3d { + // Sort the data based on the distance to the view + sort_key: FloatOrd(distance), + entity: (*render_entity, *visible_entity), + pipeline: pipeline_id, + draw_function: draw_custom, + // Sorted phase items aren't batched + batch_range: 0..1, + extra_index: PhaseItemExtraIndex::None, + indexed: mesh.indexed(), + }); + } + } +} + +// Render label used to order our render graph node that will render our phase +#[derive(RenderLabel, Debug, Clone, Hash, PartialEq, Eq)] +struct CustomDrawPassLabel; + +#[derive(Default)] +struct CustomDrawNode; +impl ViewNode for CustomDrawNode { + type ViewQuery = ( + &'static ExtractedCamera, + &'static ExtractedView, + &'static ViewTarget, + ); + + fn run<'w>( + &self, + graph: &mut RenderGraphContext, + render_context: &mut RenderContext<'w>, + (camera, view, target): QueryItem<'w, Self::ViewQuery>, + world: &'w World, + ) -> Result<(), NodeRunError> { + // First, we need to get our phases resource + let Some(stencil_phases) = world.get_resource::>() else { + return Ok(()); + }; + + // Get the view entity from the graph + let view_entity = graph.view_entity(); + + // Get the phase for the current view running our node + let Some(stencil_phase) = stencil_phases.get(&view.retained_view_entity) else { + return Ok(()); + }; + + // Render pass setup + let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor { + label: Some("stencil pass"), + // For the purpose of the example, we will write directly to the view target. A real + // stencil pass would write to a custom texture and that texture would be used in later + // passes to render custom effects using it. + color_attachments: &[Some(target.get_color_attachment())], + // We don't bind any depth buffer for this pass + depth_stencil_attachment: None, + timestamp_writes: None, + occlusion_query_set: None, + }); + + if let Some(viewport) = camera.viewport.as_ref() { + render_pass.set_camera_viewport(viewport); + } + + // Render the phase + // This will execute each draw functions of each phase items queued in this phase + if let Err(err) = stencil_phase.render(&mut render_pass, world, view_entity) { + error!("Error encountered while rendering the stencil phase {err:?}"); + } + + Ok(()) + } +}