diff --git a/Cargo.toml b/Cargo.toml index ee901649c3d94..a86c6511799e8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -373,6 +373,10 @@ path = "examples/shader/shader_custom_material.rs" name = "shader_defs" path = "examples/shader/shader_defs.rs" +[[example]] +name = "push_constants" +path = "examples/shader/push_constants.rs" + # Tools [[example]] name = "bevymark" diff --git a/crates/bevy_render/src/draw.rs b/crates/bevy_render/src/draw.rs index 1daf2de282df4..708871ba99d6e 100644 --- a/crates/bevy_render/src/draw.rs +++ b/crates/bevy_render/src/draw.rs @@ -1,6 +1,7 @@ use crate::{ pipeline::{ - IndexFormat, PipelineCompiler, PipelineDescriptor, PipelineLayout, PipelineSpecialization, + BindingShaderStage, IndexFormat, PipelineCompiler, PipelineDescriptor, PipelineLayout, + PipelineSpecialization, }, renderer::{ AssetRenderResourceBindings, BindGroup, BindGroupId, BufferId, RenderResource, @@ -38,6 +39,11 @@ pub enum RenderCommand { bind_group: BindGroupId, dynamic_uniform_indices: Option>, }, + SetPushConstants { + stages: BindingShaderStage, + offset: u32, + data: Vec, + }, DrawIndexed { indices: Range, base_vertex: i32, @@ -109,6 +115,14 @@ impl Draw { }); } + pub fn set_push_constants(&mut self, stages: BindingShaderStage, offset: u32, data: Vec) { + self.render_command(RenderCommand::SetPushConstants { + stages, + offset, + data, + }); + } + pub fn set_bind_group(&mut self, index: u32, bind_group: &BindGroup) { self.render_command(RenderCommand::SetBindGroup { index, @@ -345,6 +359,19 @@ impl<'a> DrawContext<'a> { } Ok(()) } + + pub fn set_push_constants( + &self, + draw: &mut Draw, + stages: BindingShaderStage, + offset: u32, + data: Vec, + ) -> Result<(), DrawError> { + draw.set_push_constants(stages, offset, data); + + // TODO check for issues + Ok(()) + } } pub trait Drawable { diff --git a/crates/bevy_render/src/pass/render_pass.rs b/crates/bevy_render/src/pass/render_pass.rs index 40f3ba36bfb8f..393f295a49da7 100644 --- a/crates/bevy_render/src/pass/render_pass.rs +++ b/crates/bevy_render/src/pass/render_pass.rs @@ -1,5 +1,5 @@ use crate::{ - pipeline::{BindGroupDescriptorId, IndexFormat, PipelineDescriptor}, + pipeline::{BindGroupDescriptorId, BindingShaderStage, IndexFormat, PipelineDescriptor}, renderer::{BindGroupId, BufferId, RenderContext}, }; use bevy_asset::Handle; @@ -9,6 +9,7 @@ pub trait RenderPass { fn get_render_context(&self) -> &dyn RenderContext; fn set_index_buffer(&mut self, buffer: BufferId, offset: u64, index_format: IndexFormat); fn set_vertex_buffer(&mut self, start_slot: u32, buffer: BufferId, offset: u64); + fn set_push_constants(&mut self, stages: BindingShaderStage, offset: u32, data: &[u8]); fn set_pipeline(&mut self, pipeline_handle: &Handle); fn set_viewport(&mut self, x: f32, y: f32, w: f32, h: f32, min_depth: f32, max_depth: f32); fn set_scissor_rect(&mut self, x: u32, y: u32, w: u32, h: u32); diff --git a/crates/bevy_render/src/pipeline/binding.rs b/crates/bevy_render/src/pipeline/binding.rs index 55fbf9bda802e..2b099fc7b45c5 100644 --- a/crates/bevy_render/src/pipeline/binding.rs +++ b/crates/bevy_render/src/pipeline/binding.rs @@ -5,6 +5,7 @@ use crate::texture::{ bitflags::bitflags! { pub struct BindingShaderStage: u32 { + const NONE = 0; const VERTEX = 1; const FRAGMENT = 2; const COMPUTE = 4; diff --git a/crates/bevy_render/src/pipeline/pipeline_layout.rs b/crates/bevy_render/src/pipeline/pipeline_layout.rs index 5a1816ca1de9a..7c6ceffd85da1 100644 --- a/crates/bevy_render/src/pipeline/pipeline_layout.rs +++ b/crates/bevy_render/src/pipeline/pipeline_layout.rs @@ -1,12 +1,13 @@ use super::{BindGroupDescriptor, VertexBufferLayout}; -use crate::shader::ShaderLayout; +use crate::{pipeline::BindingShaderStage, shader::ShaderLayout}; use bevy_utils::HashMap; -use std::hash::Hash; +use std::{hash::Hash, ops::Range}; #[derive(Clone, Debug, Default)] pub struct PipelineLayout { pub bind_groups: Vec, pub vertex_buffer_descriptors: Vec, + pub push_constant_ranges: Vec, } impl PipelineLayout { @@ -62,9 +63,15 @@ impl PipelineLayout { // with bevy and not with wgpu TODO: try removing this bind_groups_result.sort_by(|a, b| a.index.partial_cmp(&b.index).unwrap()); + let push_constant_ranges = shader_layouts + .iter() + .flat_map(|shader_layout| shader_layout.push_constant_ranges.clone()) + .collect::>(); + PipelineLayout { bind_groups: bind_groups_result, vertex_buffer_descriptors, + push_constant_ranges, } } } @@ -103,3 +110,13 @@ impl UniformProperty { } } } + +#[derive(Hash, Clone, Debug, PartialEq, Eq)] +pub struct PushConstantRange { + /// Stage push constant range is visible from. Each stage can only be served by at most one range. + /// One range can serve multiple stages however. + pub stages: BindingShaderStage, + /// Range in push constant memory to use for the stage. Must be less than [`Limits::max_push_constant_size`]. + /// Start and end must be aligned to the 4s. + pub range: Range, +} diff --git a/crates/bevy_render/src/render_graph/nodes/pass_node.rs b/crates/bevy_render/src/render_graph/nodes/pass_node.rs index 16cf5fef314dc..796b7825bc560 100644 --- a/crates/bevy_render/src/render_graph/nodes/pass_node.rs +++ b/crates/bevy_render/src/render_graph/nodes/pass_node.rs @@ -313,6 +313,13 @@ where ); draw_state.set_bind_group(index, bind_group); } + RenderCommand::SetPushConstants { + stages, + offset, + data, + } => { + render_pass.set_push_constants(stages, offset, &*data); + } } } }); diff --git a/crates/bevy_render/src/shader/mod.rs b/crates/bevy_render/src/shader/mod.rs index dee5179b79342..59c3b5ca0a215 100644 --- a/crates/bevy_render/src/shader/mod.rs +++ b/crates/bevy_render/src/shader/mod.rs @@ -11,7 +11,7 @@ pub use shader_defs::*; #[cfg(not(target_arch = "wasm32"))] pub use shader_reflect::*; -use crate::pipeline::{BindGroupDescriptor, VertexBufferLayout}; +use crate::pipeline::{BindGroupDescriptor, PushConstantRange, VertexBufferLayout}; /// Defines the memory layout of a shader #[derive(Debug, Clone, PartialEq, Eq)] @@ -19,6 +19,7 @@ pub struct ShaderLayout { pub bind_groups: Vec, pub vertex_buffer_layout: Vec, pub entry_point: String, + pub push_constant_ranges: Vec, } pub const GL_VERTEX_INDEX: &str = "gl_VertexIndex"; diff --git a/crates/bevy_render/src/shader/shader_reflect.rs b/crates/bevy_render/src/shader/shader_reflect.rs index 78a5f306f18fd..2a5c97194fa28 100644 --- a/crates/bevy_render/src/shader/shader_reflect.rs +++ b/crates/bevy_render/src/shader/shader_reflect.rs @@ -1,7 +1,7 @@ use crate::{ pipeline::{ BindGroupDescriptor, BindType, BindingDescriptor, BindingShaderStage, InputStepMode, - UniformProperty, VertexAttribute, VertexBufferLayout, VertexFormat, + PushConstantRange, UniformProperty, VertexAttribute, VertexBufferLayout, VertexFormat, }, shader::{ShaderLayout, GL_INSTANCE_INDEX, GL_VERTEX_INDEX}, texture::{TextureSampleType, TextureViewDimension}, @@ -28,6 +28,22 @@ impl ShaderLayout { bind_groups.push(bind_group); } + let mut push_constant_ranges = Vec::new(); + for push_constant_block in module.enumerate_push_constant_blocks(None).unwrap() { + let range = push_constant_block.offset..push_constant_block.size; + let mut stages = BindingShaderStage::NONE; + if shader_stage.contains(ReflectShaderStageFlags::VERTEX) { + stages.insert(BindingShaderStage::VERTEX); + } + if shader_stage.contains(ReflectShaderStageFlags::FRAGMENT) { + stages.insert(BindingShaderStage::FRAGMENT); + } + if shader_stage.contains(ReflectShaderStageFlags::COMPUTE) { + stages.insert(BindingShaderStage::COMPUTE); + } + push_constant_ranges.push(PushConstantRange { stages, range }) + } + // obtain attribute descriptors from reflection let mut vertex_attributes = Vec::new(); for input_variable in module.enumerate_input_variables(None).unwrap() { @@ -83,6 +99,7 @@ impl ShaderLayout { bind_groups, vertex_buffer_layout, entry_point: entry_point_name, + push_constant_ranges, } } Err(err) => panic!("Failed to reflect shader layout: {:?}.", err), @@ -323,6 +340,9 @@ mod tests { mat4 ViewProj; }; layout(set = 1, binding = 0) uniform texture2D Texture; + layout(push_constant) uniform SomeTestUniform { + vec3 SomeTestField; + }; void main() { v_Position = Vertex_Position; @@ -396,7 +416,11 @@ mod tests { shader_stage: BindingShaderStage::VERTEX, }] ), - ] + ], + push_constant_ranges: vec![PushConstantRange { + stages: BindingShaderStage::VERTEX, + range: 0..16 + }], } ); } diff --git a/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs b/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs index e7c5b79936279..b66b342a9f4db 100644 --- a/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs +++ b/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs @@ -444,13 +444,18 @@ impl RenderResourceContext for WgpuRenderResourceContext { .iter() .map(|bind_group| bind_group_layouts.get(&bind_group.id).unwrap()) .collect::>(); + let push_constant_ranges: Vec = layout + .push_constant_ranges + .iter() + .map(|range| range.clone().wgpu_into()) + .collect(); let pipeline_layout = self .device .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: None, bind_group_layouts: bind_group_layouts.as_slice(), - push_constant_ranges: &[], + push_constant_ranges: &push_constant_ranges, }); let owned_vertex_buffer_descriptors = layout diff --git a/crates/bevy_wgpu/src/wgpu_render_pass.rs b/crates/bevy_wgpu/src/wgpu_render_pass.rs index bfa15d2980bfd..2f6cc7cf9713f 100644 --- a/crates/bevy_wgpu/src/wgpu_render_pass.rs +++ b/crates/bevy_wgpu/src/wgpu_render_pass.rs @@ -2,7 +2,7 @@ use crate::{renderer::WgpuRenderContext, wgpu_type_converter::WgpuInto, WgpuReso use bevy_asset::Handle; use bevy_render::{ pass::RenderPass, - pipeline::{BindGroupDescriptorId, IndexFormat, PipelineDescriptor}, + pipeline::{BindGroupDescriptorId, BindingShaderStage, IndexFormat, PipelineDescriptor}, renderer::{BindGroupId, BufferId, RenderContext}, }; use bevy_utils::tracing::trace; @@ -46,6 +46,11 @@ impl<'a> RenderPass for WgpuRenderPass<'a> { .set_index_buffer(buffer.slice(offset..), index_format.wgpu_into()); } + fn set_push_constants(&mut self, stages: BindingShaderStage, offset: u32, data: &[u8]) { + self.render_pass + .set_push_constants(stages.wgpu_into(), offset, data); + } + fn draw_indexed(&mut self, indices: Range, base_vertex: i32, instances: Range) { self.render_pass .draw_indexed(indices, base_vertex, instances); diff --git a/crates/bevy_wgpu/src/wgpu_type_converter.rs b/crates/bevy_wgpu/src/wgpu_type_converter.rs index 94c1714e351bd..e2483c532c2eb 100644 --- a/crates/bevy_wgpu/src/wgpu_type_converter.rs +++ b/crates/bevy_wgpu/src/wgpu_type_converter.rs @@ -3,11 +3,11 @@ use bevy_render::{ color::Color, pass::{LoadOp, Operations}, pipeline::{ - BindType, BlendFactor, BlendOperation, BlendState, ColorTargetState, ColorWrite, - CompareFunction, CullMode, DepthBiasState, DepthStencilState, FrontFace, IndexFormat, - InputStepMode, MultisampleState, PolygonMode, PrimitiveState, PrimitiveTopology, - StencilFaceState, StencilOperation, StencilState, VertexAttribute, VertexBufferLayout, - VertexFormat, + BindType, BindingShaderStage, BlendFactor, BlendOperation, BlendState, ColorTargetState, + ColorWrite, CompareFunction, CullMode, DepthBiasState, DepthStencilState, FrontFace, + IndexFormat, InputStepMode, MultisampleState, PolygonMode, PrimitiveState, + PrimitiveTopology, PushConstantRange, StencilFaceState, StencilOperation, StencilState, + VertexAttribute, VertexBufferLayout, VertexFormat, }, renderer::BufferUsage, texture::{ @@ -231,6 +231,31 @@ impl WgpuFrom<&BindType> for wgpu::BindingType { } } +impl WgpuFrom for wgpu::ShaderStage { + fn from(val: BindingShaderStage) -> Self { + let mut wgpu_val = wgpu::ShaderStage::NONE; + if val.contains(BindingShaderStage::VERTEX) { + wgpu_val.insert(wgpu::ShaderStage::VERTEX); + } + if val.contains(BindingShaderStage::FRAGMENT) { + wgpu_val.insert(wgpu::ShaderStage::FRAGMENT); + } + if val.contains(BindingShaderStage::COMPUTE) { + wgpu_val.insert(wgpu::ShaderStage::COMPUTE); + } + wgpu_val + } +} + +impl WgpuFrom for wgpu::PushConstantRange { + fn from(val: PushConstantRange) -> Self { + wgpu::PushConstantRange { + stages: val.stages.wgpu_into(), + range: val.range, + } + } +} + impl WgpuFrom for wgpu::TextureSampleType { fn from(texture_component_type: TextureSampleType) -> Self { match texture_component_type { diff --git a/examples/README.md b/examples/README.md index f1df0b3d2f4a2..3af60610b12f1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -192,6 +192,7 @@ Example | File | Description `array_texture` | [`shader/array_texture.rs`](./shader/array_texture.rs) | Illustrates how to create a texture for use with a texture2DArray shader uniform variable `hot_shader_reloading` | [`shader/hot_shader_reloading.rs`](./shader/hot_shader_reloading.rs) | Illustrates how to load shaders such that they can be edited while the example is still running `mesh_custom_attribute` | [`shader/mesh_custom_attribute.rs`](./shader/mesh_custom_attribute.rs) | Illustrates how to add a custom attribute to a mesh and use it in a custom shader +`push_constants` | [`shader/push_constants.rs`](./shader/push_constants.rs) | Demonstrates using push constants `shader_custom_material` | [`shader/shader_custom_material.rs`](./shader/shader_custom_material.rs) | Illustrates creating a custom material and a shader that uses it `shader_defs` | [`shader/shader_defs.rs`](./shader/shader_defs.rs) | Demonstrates creating a custom material that uses "shaders defs" (a tool to selectively toggle parts of a shader) diff --git a/examples/shader/push_constants.rs b/examples/shader/push_constants.rs new file mode 100644 index 0000000000000..9ac924b970707 --- /dev/null +++ b/examples/shader/push_constants.rs @@ -0,0 +1,186 @@ +use bevy::{ + core::AsBytes, + prelude::*, + reflect::TypeUuid, + render::{ + draw::Drawable, + mesh::{shape, Indices}, + pipeline::{ + BindingShaderStage, PipelineDescriptor, PipelineSpecialization, RenderPipeline, + }, + shader::{ShaderStage, ShaderStages}, + RenderStage, + }, + utils::HashSet, + wgpu::{WgpuFeature, WgpuFeatures, WgpuLimits, WgpuOptions}, +}; + +pub const CUSTOM_DRAWABLE_PIPELINE_HANDLE: HandleUntyped = + HandleUntyped::weak_from_u64(PipelineDescriptor::TYPE_UUID, 0x137c75ab7e9ad8de); + +/// This example illustrates how to create a custom material asset that uses "shader defs" and a shader that uses that material. +/// In Bevy, "shader defs" are a way to selectively enable parts of a shader based on values set in a component or asset. +fn main() { + App::build() + .insert_resource(WgpuOptions { + features: WgpuFeatures { + features: vec![WgpuFeature::PushConstants], + }, + limits: WgpuLimits { + max_push_constant_size: 128, + ..Default::default() + }, + ..Default::default() + }) + .add_plugins(DefaultPlugins) + .register_type::() + .add_startup_system(setup.system()) + .add_system_to_stage(RenderStage::Draw, draw_custom_drawable.system()) + .run(); +} + +const VERTEX_SHADER: &str = r#" +#version 450 +layout(location = 0) in vec3 Vertex_Position; +layout(set = 0, binding = 0) uniform CameraViewProj { + mat4 ViewProj; +}; +layout(set = 1, binding = 0) uniform Transform { + mat4 Model; +}; +void main() { + gl_Position = ViewProj * Model * vec4(Vertex_Position, 1.0); +} +"#; + +const FRAGMENT_SHADER: &str = r#" +#version 450 +layout(location = 0) out vec4 o_Target; +layout(push_constant) uniform PushConstants { + vec3 Color; +}; +void main() { + o_Target = vec4(Color, 1.0); +} +"#; + +#[derive(Default, Debug, Reflect)] +#[reflect(Component)] +struct CustomDrawable; + +impl Drawable for CustomDrawable { + fn draw( + &mut self, + draw: &mut Draw, + context: &mut bevy::render::draw::DrawContext, + ) -> Result<(), bevy::render::draw::DrawError> { + context.set_push_constants( + draw, + BindingShaderStage::FRAGMENT, + 0, + Into::<[f32; 4]>::into(Color::BLUE).as_bytes().to_vec(), + ) + } +} + +fn draw_custom_drawable( + mut draw_context: bevy::render::draw::DrawContext, + msaa: Res, + meshes: Res>, + mut query: Query<( + &mut Draw, + &mut RenderPipelines, + &Handle, + &mut CustomDrawable, + &Visible, + )>, +) { + for (mut draw, mut render_pipelines, mesh_handle, mut custom_drawable, visible) in + query.iter_mut() + { + if !visible.is_visible { + continue; + } + + // don't render if the mesh isn't loaded yet + let mesh = if let Some(mesh) = meshes.get(mesh_handle) { + mesh + } else { + return; + }; + + // clear out any previous render_commands + // TODO prevent draw_render_pipelines_system from running for this + draw.clear_render_commands(); + + let mut render_pipeline = RenderPipeline::specialized( + CUSTOM_DRAWABLE_PIPELINE_HANDLE.typed(), + PipelineSpecialization { + sample_count: msaa.samples, + strip_index_format: None, + shader_specialization: Default::default(), + primitive_topology: mesh.primitive_topology(), + dynamic_bindings: render_pipelines + .bindings + .iter_dynamic_bindings() + .map(|name| name.to_string()) + .collect::>(), + vertex_buffer_layout: mesh.get_vertex_buffer_layout(), + }, + ); + render_pipeline.dynamic_bindings_generation = + render_pipelines.bindings.dynamic_bindings_generation(); + + draw_context + .set_pipeline( + &mut draw, + &render_pipeline.pipeline, + &render_pipeline.specialization, + ) + .unwrap(); + custom_drawable.draw(&mut *draw, &mut draw_context).unwrap(); + draw_context + .set_bind_groups_from_bindings(&mut draw, &mut [&mut render_pipelines.bindings]) + .unwrap(); + draw_context + .set_vertex_buffers_from_bindings(&mut draw, &[&render_pipelines.bindings]) + .unwrap(); + + match mesh.indices() { + Some(Indices::U32(indices)) => draw.draw_indexed(0..indices.len() as u32, 0, 0..1), + Some(Indices::U16(indices)) => draw.draw_indexed(0..indices.len() as u32, 0, 0..1), + None => draw.draw(0..mesh.count_vertices() as u32, 0..1), + }; + } +} + +fn setup( + mut commands: Commands, + mut pipelines: ResMut>, + mut shaders: ResMut>, + mut meshes: ResMut>, +) { + // Create a new shader pipeline + let pipeline_descriptor = PipelineDescriptor::default_config(ShaderStages { + vertex: shaders.add(Shader::from_glsl(ShaderStage::Vertex, VERTEX_SHADER)), + fragment: Some(shaders.add(Shader::from_glsl(ShaderStage::Fragment, FRAGMENT_SHADER))), + }); + let pipeline_handle = pipelines.set(CUSTOM_DRAWABLE_PIPELINE_HANDLE, pipeline_descriptor); + + commands + // cube + .spawn(MeshBundle { + mesh: meshes.add(Mesh::from(shape::Cube { size: 2.0 })), + render_pipelines: RenderPipelines::from_pipelines(vec![RenderPipeline::new( + pipeline_handle, + )]), + transform: Transform::from_xyz(-2.0, 0.0, 0.0), + ..Default::default() + }) + .with(CustomDrawable) + // camera + .spawn(PerspectiveCameraBundle { + transform: Transform::from_xyz(3.0, 5.0, -8.0).looking_at(Vec3::default(), Vec3::Y), + ..Default::default() + }); +}