[msl] map push constants to buffers

src/back/msl/mod.rs

@@ -36,10 +36,12 @@ mod writer;
 
 pub use writer::Writer;
 
+pub type Slot = u8;
+
 #[derive(Clone, Debug, PartialEq)]
 #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
 pub enum BindSamplerTarget {
-    Resource(u8),
+    Resource(Slot),
     Inline(Handle<sampler::InlineSampler>),
 }
 
@@ -47,9 +49,9 @@ pub enum BindSamplerTarget {
 #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
 pub struct BindTarget {
     #[cfg_attr(feature = "deserialize", serde(default))]
-    pub buffer: Option<u8>,
+    pub buffer: Option<Slot>,
     #[cfg_attr(feature = "deserialize", serde(default))]
-    pub texture: Option<u8>,
+    pub texture: Option<Slot>,
     #[cfg_attr(feature = "deserialize", serde(default))]
     pub sampler: Option<BindSamplerTarget>,
     #[cfg_attr(feature = "deserialize", serde(default))]
@@ -67,6 +69,18 @@ pub struct BindSource {
 
 pub type BindingMap = FastHashMap<BindSource, BindTarget>;
 
+#[derive(Clone, Debug, Default, Hash, Eq, Ord, PartialEq, PartialOrd)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
+#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
+pub struct PushConstantsMap {
+    #[cfg_attr(feature = "deserialize", serde(default))]
+    pub vs_buffer: Option<Slot>,
+    #[cfg_attr(feature = "deserialize", serde(default))]
+    pub fs_buffer: Option<Slot>,
+    #[cfg_attr(feature = "deserialize", serde(default))]
+    pub cs_buffer: Option<Slot>,
+}
+
 enum ResolvedBinding {
     BuiltIn(crate::BuiltIn),
     Attribute(u32),
@@ -101,6 +115,8 @@ pub enum Error {
 pub enum EntryPointError {
     #[error("mapping of {0:?} is missing")]
     MissingBinding(BindSource),
+    #[error("mapping for push constants at stage {0:?} is missing")]
+    MissingPushConstants(crate::ShaderStage),
 }
 
 #[derive(Clone, Copy, Debug)]
@@ -118,6 +134,8 @@ pub struct Options {
     pub lang_version: (u8, u8),
     /// Binding model mapping to Metal.
     pub binding_map: BindingMap,
+    /// Push constants mapping to Metal.
+    pub push_constants_map: PushConstantsMap,
     /// Samplers to be inlined into the code.
     pub inline_samplers: Arena<sampler::InlineSampler>,
     /// Make it possible to link different stages via SPIRV-Cross.
@@ -131,6 +149,7 @@ impl Default for Options {
         Options {
             lang_version: (1, 0),
             binding_map: BindingMap::default(),
+            push_constants_map: PushConstantsMap::default(),
             inline_samplers: Arena::new(),
             spirv_cross_compatibility: false,
             fake_missing_bindings: true,
@@ -185,7 +204,7 @@ impl Options {
         }
     }
 
-    fn resolve_global_binding(
+    fn resolve_resource_binding(
         &self,
         stage: crate::ShaderStage,
         res_binding: &crate::ResourceBinding,
@@ -204,6 +223,30 @@ impl Options {
             None => Err(EntryPointError::MissingBinding(source)),
         }
     }
+
+    fn resolve_push_constants(
+        &self,
+        stage: crate::ShaderStage,
+    ) -> Result<ResolvedBinding, EntryPointError> {
+        let slot = match stage {
+            crate::ShaderStage::Vertex => self.push_constants_map.vs_buffer,
+            crate::ShaderStage::Fragment => self.push_constants_map.fs_buffer,
+            crate::ShaderStage::Compute => self.push_constants_map.cs_buffer,
+        };
+        match slot {
+            Some(slot) => Ok(ResolvedBinding::Resource(BindTarget {
+                buffer: Some(slot),
+                texture: None,
+                sampler: None,
+                mutable: false,
+            })),
+            None if self.fake_missing_bindings => Ok(ResolvedBinding::User {
+                prefix: "fake",
+                index: 0,
+            }),
+            None => Err(EntryPointError::MissingPushConstants(stage)),
+        }
+    }
 }
 
 impl ResolvedBinding {

src/back/msl/writer.rs

@@ -321,6 +321,7 @@ impl crate::StorageClass {
             crate::StorageClass::Uniform
             | crate::StorageClass::Storage
             | crate::StorageClass::Private
+            | crate::StorageClass::PushConstant
             | crate::StorageClass::Handle => true,
             _ => false,
         }
@@ -1751,7 +1752,12 @@ impl<W: Write> Writer<W> {
                     .find_map(|(var_handle, var)| {
                         if !fun_info[var_handle].is_empty() {
                             if let Some(ref br) = var.binding {
-                                if let Err(e) = options.resolve_global_binding(ep.stage, br) {
+                                if let Err(e) = options.resolve_resource_binding(ep.stage, br) {
+                                    return Some(e);
+                                }
+                            }
+                            if var.class == crate::StorageClass::PushConstant {
+                                if let Err(e) = options.resolve_push_constants(ep.stage) {
                                     return Some(e);
                                 }
                             }
@@ -1920,10 +1926,16 @@ impl<W: Write> Writer<W> {
                 if usage.is_empty() || var.class == crate::StorageClass::Private {
                     continue;
                 }
-                let resolved = var
-                    .binding
-                    .as_ref()
-                    .map(|binding| options.resolve_global_binding(ep.stage, binding).unwrap());
+                // the resolves have already been checked for `!fake_missing_bindings` case
+                let resolved = match var.class {
+                    crate::StorageClass::PushConstant => {
+                        options.resolve_push_constants(ep.stage).ok()
+                    }
+                    crate::StorageClass::WorkGroup => None,
+                    _ => options
+                        .resolve_resource_binding(ep.stage, var.binding.as_ref().unwrap())
+                        .ok(),
+                };
                 if let Some(ref resolved) = resolved {
                     // Inline samplers are be defined in the EP body
                     if resolved.as_inline_sampler(options).is_some() {
@@ -1997,7 +2009,7 @@ impl<W: Write> Writer<W> {
                     };
                 } else if let Some(ref binding) = var.binding {
                     // write an inline sampler
-                    let resolved = options.resolve_global_binding(ep.stage, binding).unwrap();
+                    let resolved = options.resolve_resource_binding(ep.stage, binding).unwrap();
                     if let Some(sampler) = resolved.as_inline_sampler(options) {
                         let name = &self.names[&NameKey::GlobalVariable(handle)];
                         writeln!(

tests/in/boids.param.ron

@@ -12,6 +12,8 @@
 			(stage: Compute, group: 0, binding: 1): (buffer: Some(1), mutable: true),
 			(stage: Compute, group: 0, binding: 2): (buffer: Some(2), mutable: true),
 		},
+		push_constants_map: (
+		),
 		inline_samplers: [],
 		spirv_cross_compatibility: false,
 		fake_missing_bindings: false,

tests/in/skybox.param.ron

@@ -12,6 +12,8 @@
 			(stage: Fragment, group: 0, binding: 1): (texture: Some(0)),
 			(stage: Fragment, group: 0, binding: 2): (sampler: Some(Inline(1))),
 		},
+		push_constants_map: (
+		),
 		inline_samplers: [
 			(
 				coord: Normalized,