Add a not-yet-working vertex pulling flag to Metal pipelines.

This is an early effort to add infrastructure to support vertex pulling
transformation of Metal shaders. It is *not* a working transformation
that generates valid, useful shaders. It includes:

1) It adds a experimental_vertex_pulling_transform flag to
msl::PipelineOptions. This flag defaults to false but can be forcibly
set to true by naga tests.
2) When the flag is set, generated msl vertex shaders are passed an
additional vertex id parameter, plus an additional parameter for each
bound vertex buffer.
3) A new naga test is added which exercises this flag and demonstrates
the effect of the transform.

Future work will make the transformed shaders valid, and add tests that
transformed shaders produce correct results.

naga/src/back/msl/mod.rs

@@ -232,6 +232,19 @@ pub struct PipelineOptions {
     ///
     /// Enable this for vertex shaders with point primitive topologies.
     pub allow_and_force_point_size: bool,
+
+    /// Experimental
+    /// If set, when generating the Metal vertex shader, transform it
+    /// to receive the vertex buffers, lengths, and vertex id as args,
+    /// and bounds-check the vertex id and use the index into the
+    /// vertex buffers to access attributes, rather than using Metal's
+    /// [[stage-in]] assembled attribute data.
+    pub experimental_vertex_pulling_transform: bool,
+
+    /// Only used if experimental_vertex_pulling_transform is set.
+    /// This is a Vec of the vertex buffer bound ids, as referenced
+    /// by the vertex shader.
+    pub vertex_buffer_indexes: Vec<u32>,
 }
 
 impl Options {

naga/src/back/msl/writer.rs

@@ -3865,6 +3865,7 @@ impl<W: Write> Writer<W> {
             // `<fun>Output`.
             let stage_out_name = format!("{fun_name}Output");
             let result_member_name = self.namer.call("member");
+            let result_return_statement: &str;
             let result_type_name = match fun.result {
                 Some(ref result) => {
                     let mut result_members = Vec::new();
@@ -3934,9 +3935,13 @@ impl<W: Write> Writer<W> {
                         )?;
                     }
                     writeln!(self.out, "}};")?;
+                    result_return_statement = "return {}";
                     &stage_out_name
                 }
-                None => "void",
+                None => {
+                    result_return_statement = "return";
+                    "void"
+                }
             };
 
             // Write the entry point function's name, and begin its argument list.
@@ -4171,16 +4176,34 @@ impl<W: Write> Writer<W> {
                 writeln!(self.out)?;
             }
 
+            if pipeline_options.experimental_vertex_pulling_transform {
+                let separator = if is_first_argument {
+                    is_first_argument = false;
+                    ' '
+                } else {
+                    ','
+                };
+
+                // Write the [[vertex_id]] argument.
+                let v_id = self.namer.call("v_id");
+                writeln!(self.out, "{separator} uint {v_id} [[vertex_id]]")?;
+
+                // Read the pipeline options we specified earlier, output one
+                // argument for every vertex buffer.
+                for index in pipeline_options.vertex_buffer_indexes.iter() {
+                    writeln!(
+                        self.out,
+                        "{separator} constant some_type *some_name [[buffer({index})]]"
+                    )?;
+                }
+            }
+
             // If this entry uses any variable-length arrays, their sizes are
             // passed as a final struct-typed argument.
             if supports_array_length {
                 // this is checked earlier
                 let resolved = options.resolve_sizes_buffer(ep).unwrap();
-                let separator = if module.global_variables.is_empty() {
-                    ' '
-                } else {
-                    ','
-                };
+                let separator = if is_first_argument { ' ' } else { ',' };
                 write!(
                     self.out,
                     "{separator} constant _mslBufferSizes& _buffer_sizes",
@@ -4192,6 +4215,16 @@ impl<W: Write> Writer<W> {
             // end of the entry point argument list
             writeln!(self.out, ") {{")?;
 
+            if pipeline_options.experimental_vertex_pulling_transform {
+                // Output the bounds check.
+                writeln!(
+                    self.out,
+                    "{}if (false) {{{};}}",
+                    back::Level(1),
+                    result_return_statement
+                )?;
+            }
+
             if need_workgroup_variables_initialization {
                 self.write_workgroup_variables_initialization(
                     module,

naga/tests/in/interface.param.ron

@@ -27,5 +27,7 @@
 	),
 	msl_pipeline: (
 		allow_and_force_point_size: true,
+    experimental_vertex_pulling_transform: false,
+    vertex_buffer_indexes: [],
 	),
 )

naga/tests/in/vertex-pulling-transform.param.ron

@@ -0,0 +1,7 @@
+(
+	msl_pipeline: (
+		allow_and_force_point_size: false,
+    experimental_vertex_pulling_transform: true,
+    vertex_buffer_indexes: [],
+	),
+)

naga/tests/in/vertex-pulling-transform.wgsl

@@ -0,0 +1,29 @@
+struct VertexOutput {
+  @builtin(position) position: vec4<f32>,
+  @location(0) color: vec4<f32>,
+  @location(1) texcoord: vec2<f32>,
+}
+
+struct VertexInput {
+  @location(0) position: vec4<f32>,
+  @location(1) normal: vec3<f32>,
+  @location(2) texcoord: vec2<f32>,
+}
+
+@group(0) @binding(0) var<uniform> mvp_matrix: mat4x4<f32>;
+
+@vertex
+fn render_vertex(v_in: VertexInput) -> VertexOutput
+{
+  var v_out: VertexOutput;
+  v_out.position = v_in.position * mvp_matrix;
+  v_out.color = do_lighting(v_in.position,
+                            v_in.normal);
+  v_out.texcoord = v_in.texcoord;
+  return v_out;
+}
+
+fn do_lighting(position: vec4<f32>, normal: vec3<f32>) -> vec4<f32> {
+  // blah blah blah
+  return vec4<f32>(0);
+}

naga/tests/out/msl/vertex-pulling-transform.msl

@@ -0,0 +1,51 @@
+// language: metal1.0
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using metal::uint;
+
+struct VertexOutput {
+    metal::float4 position;
+    metal::float4 color;
+    metal::float2 texcoord;
+};
+struct VertexInput {
+    metal::float4 position;
+    metal::float3 normal;
+    metal::float2 texcoord;
+};
+
+metal::float4 do_lighting(
+    metal::float4 position,
+    metal::float3 normal
+) {
+    return metal::float4(0.0);
+}
+
+struct render_vertexInput {
+    metal::float4 position [[attribute(0)]];
+    metal::float3 normal [[attribute(1)]];
+    metal::float2 texcoord [[attribute(2)]];
+};
+struct render_vertexOutput {
+    metal::float4 position [[position]];
+    metal::float4 color [[user(loc0), center_perspective]];
+    metal::float2 texcoord [[user(loc1), center_perspective]];
+};
+vertex render_vertexOutput render_vertex(
+  render_vertexInput varyings [[stage_in]]
+, constant metal::float4x4& mvp_matrix [[user(fake0)]]
+, uint v_id [[vertex_id]]
+) {
+    if (false) {return {};}
+    const VertexInput v_in = { varyings.position, varyings.normal, varyings.texcoord };
+    VertexOutput v_out = {};
+    metal::float4x4 _e5 = mvp_matrix;
+    v_out.position = v_in.position * _e5;
+    metal::float4 _e10 = do_lighting(v_in.position, v_in.normal);
+    v_out.color = _e10;
+    v_out.texcoord = v_in.texcoord;
+    VertexOutput _e13 = v_out;
+    const auto _tmp = _e13;
+    return render_vertexOutput { _tmp.position, _tmp.color, _tmp.texcoord };
+}

naga/tests/snapshots.rs

@@ -802,6 +802,7 @@ fn convert_wgsl() {
             "abstract-types-operators",
             Targets::SPIRV | Targets::METAL | Targets::GLSL | Targets::WGSL,
         ),
+        ("vertex-pulling-transform", Targets::METAL),
     ];
 
     for &(name, targets) in inputs.iter() {

wgpu-hal/src/metal/device.rs

@@ -63,6 +63,7 @@ impl super::Device {
     fn load_shader(
         &self,
         stage: &crate::ProgrammableStage<super::Api>,
+        vertex_buffer_count: usize,
         layout: &super::PipelineLayout,
         primitive_class: metal::MTLPrimitiveTopologyClass,
         naga_stage: naga::ShaderStage,
@@ -107,11 +108,16 @@ impl super::Device {
             zero_initialize_workgroup_memory: true,
         };
 
+        let last_vertex_buffer_index = self.shared.private_caps.max_vertex_buffers - 1;
         let pipeline_options = naga::back::msl::PipelineOptions {
             allow_and_force_point_size: match primitive_class {
                 metal::MTLPrimitiveTopologyClass::Point => true,
                 _ => false,
             },
+            experimental_vertex_pulling_transform: false,
+            vertex_buffer_indexes: ((last_vertex_buffer_index - vertex_buffer_count as u32)
+                ..last_vertex_buffer_index)
+                .collect(),
         };
 
         let (source, info) = naga::back::msl::write_string(
@@ -819,6 +825,7 @@ impl crate::Device<super::Api> for super::Device {
             let (vs_lib, vs_info) = {
                 let vs = self.load_shader(
                     &desc.vertex_stage,
+                    desc.vertex_buffers.len(),
                     desc.layout,
                     primitive_class,
                     naga::ShaderStage::Vertex,
@@ -846,6 +853,7 @@ impl crate::Device<super::Api> for super::Device {
                 Some(ref stage) => {
                     let fs = self.load_shader(
                         stage,
+                        0,
                         desc.layout,
                         primitive_class,
                         naga::ShaderStage::Fragment,
@@ -1038,6 +1046,7 @@ impl crate::Device<super::Api> for super::Device {
 
             let cs = self.load_shader(
                 &desc.stage,
+                0,
                 desc.layout,
                 metal::MTLPrimitiveTopologyClass::Unspecified,
                 naga::ShaderStage::Compute,