Implement push constants for metal backend (gfx-rs#2314)

* Implement push constants for metal backend

* Clear push constants on reset; get upload size from layout

wgpu-hal/src/metal/adapter.rs

@@ -918,6 +918,7 @@ impl super::PrivateCapabilities {
  | F::MAPPABLE_PRIMARY_BUFFERS
  | F::VERTEX_WRITABLE_STORAGE
  | F::TEXTURE_ADAPTER_SPECIFIC_FORMAT_FEATURES
+ | F::PUSH_CONSTANTS
  | F::POLYGON_MODE_LINE
  | F::CLEAR_COMMANDS
  | F::TEXTURE_FORMAT_16BIT_NORM;

wgpu-hal/src/metal/command.rs

@@ -16,6 +16,7 @@ impl Default for super::CommandState {
  stage_infos: Default::default(),
  storage_buffer_length_map: Default::default(),
  work_group_memory_sizes: Vec::new(),
+ push_constants: Vec::new(),
  }
  }
 }
@@ -61,6 +62,7 @@ impl super::CommandState {
  self.stage_infos.fs.clear();
  self.stage_infos.cs.clear();
  self.work_group_memory_sizes.clear();
+ self.push_constants.clear();
  }
 
  fn make_sizes_buffer_update<'a>(
@@ -587,12 +589,41 @@ impl crate::CommandEncoder<super::Api> for super::CommandEncoder {
 
  unsafe fn set_push_constants(
  &mut self,
- _layout: &super::PipelineLayout,
- _stages: wgt::ShaderStages,
- _offset: u32,
- _data: &[u32],
+ layout: &super::PipelineLayout,
+ stages: wgt::ShaderStages,
+ offset: u32,
+ data: &[u32],
  ) {
- //TODO
+ let state_pc = &mut self.state.push_constants;
+ if state_pc.len() < layout.total_push_constants as usize {
+ state_pc.resize(layout.total_push_constants as usize, 0);
+ }
+ assert_eq!(offset as usize % WORD_SIZE, 0);
+
+ let offset = offset as usize / WORD_SIZE;
+ state_pc[offset..offset + data.len()].copy_from_slice(data);
+
+ if stages.contains(wgt::ShaderStages::COMPUTE) {
+ self.state.compute.as_ref().unwrap().set_bytes(
+ layout.push_constants_infos.cs.unwrap().buffer_index as _,
+ (layout.total_push_constants as usize * WORD_SIZE) as _,
+ state_pc.as_ptr() as _,
+ )
+ }
+ if stages.contains(wgt::ShaderStages::VERTEX) {
+ self.state.render.as_ref().unwrap().set_vertex_bytes(
+ layout.push_constants_infos.vs.unwrap().buffer_index as _,
+ (layout.total_push_constants as usize * WORD_SIZE) as _,
+ state_pc.as_ptr() as _,
+ )
+ }
+ if stages.contains(wgt::ShaderStages::FRAGMENT) {
+ self.state.render.as_ref().unwrap().set_fragment_bytes(
+ layout.push_constants_infos.fs.unwrap().buffer_index as _,
+ (layout.total_push_constants as usize * WORD_SIZE) as _,
+ state_pc.as_ptr() as _,
+ )
+ }
  }
 
  unsafe fn insert_debug_marker(&mut self, label: &str) {

wgpu-hal/src/metal/device.rs

@@ -471,6 +471,7 @@ impl crate::Device<super::Api> for super::Device {
  let mut bind_group_infos = arrayvec::ArrayVec::new();
 
  // First, place the push constants
+ let mut total_push_constants = 0;
  for info in stage_data.iter_mut() {
  for pcr in desc.push_constant_ranges {
  if pcr.stages.contains(map_naga_stage(info.stage)) {
@@ -492,6 +493,8 @@ impl crate::Device<super::Api> for super::Device {
  info.pc_buffer = Some(info.counters.buffers);
  info.counters.buffers += 1;
  }
+
+ total_push_constants = total_push_constants.max(info.pc_limit);
  }
 
  // Second, place the described resources
@@ -641,6 +644,7 @@ impl crate::Device<super::Api> for super::Device {
  image: naga::proc::BoundsCheckPolicy::ReadZeroSkipWrite,
  },
  },
+ total_push_constants,
  })
  }
  unsafe fn destroy_pipeline_layout(&self, _pipeline_layout: super::PipelineLayout) {}

wgpu-hal/src/metal/mod.rs

@@ -524,6 +524,7 @@ pub struct PipelineLayout {
  bind_group_infos: ArrayVec<BindGroupLayoutInfo, { crate::MAX_BIND_GROUPS }>,
  push_constants_infos: MultiStageData<Option<PushConstantsInfo>>,
  total_counters: MultiStageResourceCounters,
+ total_push_constants: u32,
 }
 
 trait AsNative {
@@ -709,6 +710,7 @@ struct CommandState {
  stage_infos: MultiStageData<PipelineStageInfo>,
  storage_buffer_length_map: fxhash::FxHashMap<naga::ResourceBinding, wgt::BufferSize>,
  work_group_memory_sizes: Vec<u32>,
+ push_constants: Vec<u32>,
 }
 
 pub struct CommandEncoder {