Implement push constants for metal backend (#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 {