Switch Binding Arrays on Metal to Argument Buffers (#6751)

Author: cwfitzgerald

.github/workflows/ci.yml

@@ -48,7 +48,7 @@ env:
   CARGO_INCREMENTAL: false
   CARGO_TERM_COLOR: always
   WGPU_DX12_COMPILER: dxc
-  RUST_LOG: info
+  RUST_LOG: debug
   RUST_BACKTRACE: full
   PKG_CONFIG_ALLOW_CROSS: 1 # allow android to work
   RUSTFLAGS: -D warnings

Cargo.lock

@@ -142,9 +142,9 @@ wgpu-types = { version = "23.0.0", path = "./wgpu-types" }
 winit = { version = "0.29", features = ["android-native-activity"] }
 
 # Metal dependencies
+metal = { version = "0.30.0", git = "https://github.com/gfx-rs/metal-rs.git", rev = "ef768ff9d7" }
 block = "0.1"
 core-graphics-types = "0.1"
-metal = { version = "0.30.0" }
 objc = "0.2.5"
 
 # Vulkan dependencies

Cargo.toml

@@ -62,6 +62,15 @@ impl BindGroupState {
 fn run_bench(ctx: &mut Criterion) {
     let state = Lazy::new(BindGroupState::new);
 
+    if !state
+        .device_state
+        .device
+        .features()
+        .contains(wgpu::Features::TEXTURE_BINDING_ARRAY)
+    {
+        return;
+    }
+
     let mut group = ctx.benchmark_group("Bind Group Creation");
 
     for count in [5, 50, 500, 5_000, 50_000] {

benches/benches/bind_groups.rs

@@ -341,4 +341,5 @@ pub const RESERVED: &[&str] = &[
     "DefaultConstructible",
     super::writer::FREXP_FUNCTION,
     super::writer::MODF_FUNCTION,
+    super::writer::ARGUMENT_BUFFER_WRAPPER_STRUCT,
 ];

naga/src/back/msl/keywords.rs

@@ -59,8 +59,6 @@ pub struct BindTarget {
     pub buffer: Option<Slot>,
     pub texture: Option<Slot>,
     pub sampler: Option<BindSamplerTarget>,
-    /// If the binding is an unsized binding array, this overrides the size.
-    pub binding_array_size: Option<u32>,
     pub mutable: bool,
 }
 

naga/src/back/msl/mod.rs

@@ -36,6 +36,14 @@ const RAY_QUERY_FUN_MAP_INTERSECTION: &str = "_map_intersection_type";
 pub(crate) const ATOMIC_COMP_EXCH_FUNCTION: &str = "naga_atomic_compare_exchange_weak_explicit";
 pub(crate) const MODF_FUNCTION: &str = "naga_modf";
 pub(crate) const FREXP_FUNCTION: &str = "naga_frexp";
+/// For some reason, Metal does not let you have `metal::texture<..>*` as a buffer argument.
+/// However, if you put that texture inside a struct, everything is totally fine. This
+/// baffles me to no end.
+///
+/// As such, we wrap all argument buffers in a struct that has a single generic `<T>` field.
+/// This allows `NagaArgumentBufferWrapper<metal::texture<..>>*` to work. The astute among
+/// you have noticed that this should be exactly the same to the compiler, and you're correct.
+pub(crate) const ARGUMENT_BUFFER_WRAPPER_STRUCT: &str = "NagaArgumentBufferWrapper";
 
 /// Write the Metal name for a Naga numeric type: scalar, vector, or matrix.
 ///
@@ -275,24 +283,17 @@ impl Display for TypeContext<'_> {
             crate::TypeInner::RayQuery => {
                 write!(out, "{RAY_QUERY_TYPE}")
             }
-            crate::TypeInner::BindingArray { base, size } => {
+            crate::TypeInner::BindingArray { base, .. } => {
                 let base_tyname = Self {
                     handle: base,
                     first_time: false,
                     ..*self
                 };
 
-                if let Some(&super::ResolvedBinding::Resource(super::BindTarget {
-                    binding_array_size: Some(override_size),
-                    ..
-                })) = self.binding
-                {
-                    write!(out, "{NAMESPACE}::array<{base_tyname}, {override_size}>")
-                } else if let crate::ArraySize::Constant(size) = size {
-                    write!(out, "{NAMESPACE}::array<{base_tyname}, {size}>")
-                } else {
-                    unreachable!("metal requires all arrays be constant sized");
-                }
+                write!(
+                    out,
+                    "constant {ARGUMENT_BUFFER_WRAPPER_STRUCT}<{base_tyname}>*"
+                )
             }
         }
     }
@@ -2552,6 +2553,8 @@ impl<W: Write> Writer<W> {
             } => true,
             _ => false,
         };
+        let accessing_wrapped_binding_array =
+            matches!(*base_ty, crate::TypeInner::BindingArray { .. });
 
         self.put_access_chain(base, policy, context)?;
         if accessing_wrapped_array {
@@ -2588,6 +2591,10 @@ impl<W: Write> Writer<W> {
 
         write!(self.out, "]")?;
 
+        if accessing_wrapped_binding_array {
+            write!(self.out, ".{WRAPPED_ARRAY_FIELD}")?;
+        }
+
         Ok(())
     }
 
@@ -3701,7 +3708,18 @@ impl<W: Write> Writer<W> {
     }
 
     fn write_type_defs(&mut self, module: &crate::Module) -> BackendResult {
+        let mut generated_argument_buffer_wrapper = false;
         for (handle, ty) in module.types.iter() {
+            if let crate::TypeInner::BindingArray { .. } = ty.inner {
+                if !generated_argument_buffer_wrapper {
+                    writeln!(self.out, "template <typename T>")?;
+                    writeln!(self.out, "struct {ARGUMENT_BUFFER_WRAPPER_STRUCT} {{")?;
+                    writeln!(self.out, "{}T {WRAPPED_ARRAY_FIELD};", back::INDENT)?;
+                    writeln!(self.out, "}};")?;
+                    generated_argument_buffer_wrapper = true;
+                }
+            }
+
             if !ty.needs_alias() {
                 continue;
             }
@@ -5132,13 +5150,10 @@ template <typename A>
                             let target = options.get_resource_binding_target(ep, br);
                             let good = match target {
                                 Some(target) => {
-                                    let binding_ty = match module.types[var.ty].inner {
-                                        crate::TypeInner::BindingArray { base, .. } => {
-                                            &module.types[base].inner
-                                        }
-                                        ref ty => ty,
-                                    };
-                                    match *binding_ty {
+                                    // We intentionally don't dereference binding_arrays here,
+                                    // so that binding arrays fall to the buffer location.
+
+                                    match module.types[var.ty].inner {
                                         crate::TypeInner::Image { .. } => target.texture.is_some(),
                                         crate::TypeInner::Sampler { .. } => {
                                             target.sampler.is_some()

naga/src/back/msl/writer.rs

@@ -19,11 +19,11 @@
 		restrict_indexing: true
 	),
 	msl: (
-		lang_version: (2, 0),
+		lang_version: (3, 0),
 		per_entry_point_map: {
 			"main": (
 				resources: {
-					(group: 0, binding: 0): (texture: Some(0), binding_array_size: Some(10), mutable: false),
+					(group: 0, binding: 0): (buffer: Some(0), binding_array_size: Some(10), mutable: false),
 				},
 				sizes_buffer: None,
 			)