Update and extension of Adding push constants support #1369

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"

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<Arc<[u32]>>,
     },
+    SetPushConstants {
+        stages: BindingShaderStage,
+        offset: u32,
+        data: Vec<u8>,
+    },
     DrawIndexed {
         indices: Range<u32>,
         base_vertex: i32,
@@ -109,6 +115,14 @@ impl Draw {
         });
     }
 
+    pub fn set_push_constants(&mut self, stages: BindingShaderStage, offset: u32, data: Vec<u8>) {
+        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<u8>,
+    ) -> Result<(), DrawError> {
+        draw.set_push_constants(stages, offset, data);
+
+        // TODO check for issues
+        Ok(())
+    }
 }
 
 pub trait Drawable {

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<PipelineDescriptor>);
     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);

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;

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<BindGroupDescriptor>,
     pub vertex_buffer_descriptors: Vec<VertexBufferLayout>,
+    pub push_constant_ranges: Vec<PushConstantRange>,
 }
 
 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::<Vec<_>>();
+
         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<u32>,
+}

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);
+                    }
                 }
             }
         });

crates/bevy_render/src/shader/mod.rs

@@ -11,14 +11,15 @@ 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)]
 pub struct ShaderLayout {
     pub bind_groups: Vec<BindGroupDescriptor>,
     pub vertex_buffer_layout: Vec<VertexBufferLayout>,
     pub entry_point: String,
+    pub push_constant_ranges: Vec<PushConstantRange>,
 }
 
 pub const GL_VERTEX_INDEX: &str = "gl_VertexIndex";

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
+                }],
             }
         );
     }

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::<Vec<&wgpu::BindGroupLayout>>();
+        let push_constant_ranges: Vec<wgpu::PushConstantRange> = 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

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<u32>, base_vertex: i32, instances: Range<u32>) {
         self.render_pass
             .draw_indexed(indices, base_vertex, instances);

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<BindingShaderStage> 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<PushConstantRange> for wgpu::PushConstantRange {
+    fn from(val: PushConstantRange) -> Self {
+        wgpu::PushConstantRange {
+            stages: val.stages.wgpu_into(),
+            range: val.range,
+        }
+    }
+}
+
 impl WgpuFrom<TextureSampleType> for wgpu::TextureSampleType {
     fn from(texture_component_type: TextureSampleType) -> Self {
         match texture_component_type {

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)