[hal/metal] Mesh Shaders

CHANGELOG.md

@@ -201,6 +201,9 @@ By @cwfitzgerald in [#8162](https://github.com/gfx-rs/wgpu/pull/8162).
 
 - Allow disabling waiting for latency waitable object. By @marcpabst in [#7400](https://github.com/gfx-rs/wgpu/pull/7400)
 
+#### Metal
+- Add support for mesh shaders. By @SupaMaggie70Incorporated in [#8139](https://github.com/gfx-rs/wgpu/pull/8139)
+
 ### Bug Fixes
 
 #### General

examples/features/src/mesh_shader/mod.rs

@@ -33,13 +33,25 @@ fn compile_glsl(
     }
 }
 
+fn compile_msl(device: &wgpu::Device, entry: &str) -> wgpu::ShaderModule {
+    unsafe {
+        device.create_shader_module_passthrough(wgpu::ShaderModuleDescriptorPassthrough {
+            entry_point: entry.to_owned(),
+            label: None,
+            msl: Some(std::borrow::Cow::Borrowed(include_str!("shader.metal"))),
+            num_workgroups: (1, 1, 1),
+            ..Default::default()
+        })
+    }
+}
+
 pub struct Example {
     pipeline: wgpu::RenderPipeline,
 }
 impl crate::framework::Example for Example {
     fn init(
         config: &wgpu::SurfaceConfiguration,
-        _adapter: &wgpu::Adapter,
+        adapter: &wgpu::Adapter,
         device: &wgpu::Device,
         _queue: &wgpu::Queue,
     ) -> Self {
@@ -48,11 +60,19 @@ impl crate::framework::Example for Example {
             bind_group_layouts: &[],
             push_constant_ranges: &[],
         });
-        let (ts, ms, fs) = (
-            compile_glsl(device, include_bytes!("shader.task"), "task"),
-            compile_glsl(device, include_bytes!("shader.mesh"), "mesh"),
-            compile_glsl(device, include_bytes!("shader.frag"), "frag"),
-        );
+        let (ts, ms, fs) = if adapter.get_info().backend == wgpu::Backend::Metal {
+            (
+                compile_msl(device, "taskShader"),
+                compile_msl(device, "meshShader"),
+                compile_msl(device, "fragShader"),
+            )
+        } else {
+            (
+                compile_glsl(device, include_bytes!("shader.task"), "task"),
+                compile_glsl(device, include_bytes!("shader.mesh"), "mesh"),
+                compile_glsl(device, include_bytes!("shader.frag"), "frag"),
+            )
+        };
         let pipeline = device.create_mesh_pipeline(&wgpu::MeshPipelineDescriptor {
             label: None,
             layout: Some(&pipeline_layout),

examples/features/src/mesh_shader/shader.metal

@@ -0,0 +1,77 @@
+using namespace metal;
+
+struct OutVertex {
+    float4 Position [[position]];
+    float4 Color [[user(locn0)]];
+};
+
+struct OutPrimitive {
+    float4 ColorMask [[flat]] [[user(locn1)]];
+    bool CullPrimitive [[primitive_culled]];
+};
+
+struct InVertex {
+};
+
+struct InPrimitive {
+    float4 ColorMask [[flat]] [[user(locn1)]];
+};
+
+struct FragmentIn {
+    float4 Color [[user(locn0)]];
+    float4 ColorMask [[flat]] [[user(locn1)]];
+};
+
+struct PayloadData {
+    float4 ColorMask;
+    bool Visible;
+};
+
+using Meshlet = metal::mesh<OutVertex, OutPrimitive, 3, 1, topology::triangle>;
+
+
+constant float4 positions[3] = {
+    float4(0.0, 1.0, 0.0, 1.0),
+    float4(-1.0, -1.0, 0.0, 1.0),
+    float4(1.0, -1.0, 0.0, 1.0)
+};
+
+constant float4 colors[3] = {
+    float4(0.0, 1.0, 0.0, 1.0),
+    float4(0.0, 0.0, 1.0, 1.0),
+    float4(1.0, 0.0, 0.0, 1.0)
+};
+
+
+[[object]]
+void taskShader(uint3 tid [[thread_position_in_grid]], object_data PayloadData &outPayload [[payload]], mesh_grid_properties grid) {
+    outPayload.ColorMask = float4(1.0, 1.0, 0.0, 1.0);
+    outPayload.Visible = true;
+    grid.set_threadgroups_per_grid(uint3(3, 1, 1));
+}
+
+[[mesh]]
+void meshShader(
+    object_data PayloadData const& payload [[payload]],
+    Meshlet out
+)
+{
+    out.set_primitive_count(1);
+
+    for(int i = 0;i < 3;i++) {
+        OutVertex vert;
+        vert.Position = positions[i];
+        vert.Color = colors[i] * payload.ColorMask;
+        out.set_vertex(i, vert);
+        out.set_index(i, i);
+    }
+
+    OutPrimitive prim;
+    prim.ColorMask = float4(1.0, 0.0, 0.0, 1.0);
+    prim.CullPrimitive = !payload.Visible;
+    out.set_primitive(0, prim);
+}
+
+fragment float4 fragShader(FragmentIn data [[stage_in]]) {
+    return data.Color * data.ColorMask;
+}

tests/tests/wgpu-gpu/mesh_shader/mod.rs

@@ -86,6 +86,9 @@ fn mesh_pipeline_build(
     frag: Option<&[u8]>,
     draw: bool,
 ) {
+    if ctx.adapter.get_info().backend != wgpu::Backend::Vulkan {
+        return;
+    }
     let device = &ctx.device;
     let (_depth_image, depth_view, depth_state) = create_depth(device);
     let task = task.map(|t| compile_glsl(device, t, "task"));
@@ -160,6 +163,9 @@ pub enum DrawType {
 }
 
 fn mesh_draw(ctx: &TestingContext, draw_type: DrawType) {
+    if ctx.adapter.get_info().backend != wgpu::Backend::Vulkan {
+        return;
+    }
     let device = &ctx.device;
     let (_depth_image, depth_view, depth_state) = create_depth(device);
     let task = compile_glsl(device, BASIC_TASK, "task");

wgpu-hal/src/metal/adapter.rs

@@ -902,6 +902,8 @@ impl super::PrivateCapabilities {
                 && (device.supports_family(MTLGPUFamily::Apple7)
                     || device.supports_family(MTLGPUFamily::Mac2)),
             supports_shared_event: version.at_least((10, 14), (12, 0), os_is_mac),
+            mesh_shaders: device.supports_family(MTLGPUFamily::Apple7)
+                || device.supports_family(MTLGPUFamily::Mac2),
         }
     }
 
@@ -1001,6 +1003,8 @@ impl super::PrivateCapabilities {
             features.insert(F::SUBGROUP | F::SUBGROUP_BARRIER);
         }
 
+        features.set(F::EXPERIMENTAL_MESH_SHADER, self.mesh_shaders);
+
         features
     }
 
@@ -1077,10 +1081,11 @@ impl super::PrivateCapabilities {
                 max_buffer_size: self.max_buffer_size,
                 max_non_sampler_bindings: u32::MAX,
 
-                max_task_workgroup_total_count: 0,
-                max_task_workgroups_per_dimension: 0,
+                // See https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf, Maximum threadgroups per mesh shader grid
+                max_task_workgroup_total_count: 1024,
+                max_task_workgroups_per_dimension: 1024,
                 max_mesh_multiview_count: 0,
-                max_mesh_output_layers: 0,
+                max_mesh_output_layers: self.max_texture_layers as u32,
 
                 max_blas_primitive_count: 0, // When added: 2^28 from https://developer.apple.com/documentation/metal/mtlaccelerationstructureusage/extendedlimits
                 max_blas_geometry_count: 0,  // When added: 2^24