Adding push constants support

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,

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/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 {
@@ -65,6 +66,8 @@ impl PipelineLayout {
  PipelineLayout {
  bind_groups: bind_groups_result,
  vertex_buffer_descriptors,
+ // TODO: get push constant ranges from shader layout
+ push_constant_ranges: vec![],
  }
  }
 }
@@ -103,3 +106,13 @@ impl UniformProperty {
  }
  }
 }
+
+#[derive(Hash, Clone, Debug)]
+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

@@ -338,6 +338,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_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::{
@@ -230,6 +230,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 {