Various SPIR-V fixes for mesh shaders

CHANGELOG.md

@@ -41,6 +41,11 @@ Bottom level categories:
 
 ## Unreleased
 
+### Bug Fixes
+
+#### Vulkan
+- Fixed a variety of mesh shader SPIR-V writer issues from the original implementation. By @inner-daemons in [#8756](https://github.com/gfx-rs/wgpu/pull/8756)
+
 ## v28.0.0 (2025-12-17)
 
 ### Major Changes

examples/features/src/framework.rs

@@ -260,7 +260,8 @@ impl ExampleContext {
     async fn init_async<E: Example>(surface: &mut SurfaceWrapper, window: Arc<Window>) -> Self {
         log::info!("Initializing wgpu...");
 
-        let instance_descriptor = wgpu::InstanceDescriptor::from_env_or_default();
+        let mut instance_descriptor = wgpu::InstanceDescriptor::from_env_or_default();
+        instance_descriptor.backend_options.dx12.shader_compiler = wgpu::Dx12Compiler::StaticDxc;
         let instance = wgpu::Instance::new(&instance_descriptor);
         surface.pre_adapter(&instance, window);
         let adapter = get_adapter_with_capabilities_or_from_env(

examples/features/src/mesh_shader/mod.rs

@@ -1,9 +1,15 @@
 // Same as in mesh shader tests
 fn compile_wgsl(device: &wgpu::Device) -> wgpu::ShaderModule {
-    device.create_shader_module(wgpu::ShaderModuleDescriptor {
-        label: None,
-        source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()),
-    })
+    // Workgroup memory zero initialization can be expensive for mesh shaders
+    unsafe {
+        device.create_shader_module_trusted(
+            wgpu::ShaderModuleDescriptor {
+                label: None,
+                source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()),
+            },
+            wgpu::ShaderRuntimeChecks::unchecked(),
+        )
+    }
 }
 fn compile_hlsl(device: &wgpu::Device, entry: &str, stage_str: &str) -> wgpu::ShaderModule {
     let out_path = format!(

examples/features/src/mesh_shader/shader.wgsl

@@ -33,12 +33,15 @@ var<workgroup> workgroupData: f32;
 
 @task
 @payload(taskPayload)
-@workgroup_size(1)
-fn ts_main() -> @builtin(mesh_task_size) vec3<u32> {
-    workgroupData = 1.0;
-    taskPayload.colorMask = vec4(1.0, 1.0, 0.0, 1.0);
-    taskPayload.visible = true;
-    return vec3(1, 1, 1);
+@workgroup_size(64)
+fn ts_main(@builtin(local_invocation_id) thread_id: vec3<u32>) -> @builtin(mesh_task_size) vec3<u32> {
+    if thread_id.x == 0 {
+        workgroupData = 1.0;
+        taskPayload.colorMask = vec4(1.0, 1.0, 0.0, 1.0);
+        taskPayload.visible = true;
+        return vec3(1, 1, 1);
+    }
+    return vec3(0, 0, 0);
 }
 
 struct MeshOutput {
@@ -52,24 +55,27 @@ var<workgroup> mesh_output: MeshOutput;
 
 @mesh(mesh_output)
 @payload(taskPayload)
-@workgroup_size(1)
-fn ms_main() {
-    mesh_output.vertex_count = 3;
-    mesh_output.primitive_count = 1;
-    workgroupData = 2.0;
+@workgroup_size(64)
+fn ms_main(@builtin(local_invocation_id) thread_id: vec3<u32>) {
+    if thread_id.x == 0 {
+        mesh_output.vertex_count = 3;
+        mesh_output.primitive_count = 1;
+        workgroupData = 2.0;
 
-    mesh_output.vertices[0].position = positions[0];
-    mesh_output.vertices[0].color = colors[0] * taskPayload.colorMask;
+        mesh_output.vertices[0].position = positions[0];
+        mesh_output.vertices[0].color = colors[0] * taskPayload.colorMask;
 
-    mesh_output.vertices[1].position = positions[1];
-    mesh_output.vertices[1].color = colors[1] * taskPayload.colorMask;
+        mesh_output.vertices[1].position = positions[1];
+        mesh_output.vertices[1].color = colors[1] * taskPayload.colorMask;
 
-    mesh_output.vertices[2].position = positions[2];
-    mesh_output.vertices[2].color = colors[2] * taskPayload.colorMask;
+        mesh_output.vertices[2].position = positions[2];
+        mesh_output.vertices[2].color = colors[2] * taskPayload.colorMask;
 
-    mesh_output.primitives[0].indices = vec3<u32>(0, 1, 2);
-    mesh_output.primitives[0].cull = !taskPayload.visible;
-    mesh_output.primitives[0].colorMask = vec4<f32>(1.0, 0.0, 1.0, 1.0);
+        mesh_output.primitives[0].indices = vec3<u32>(0, 1, 2);
+        mesh_output.primitives[0].cull = !taskPayload.visible;
+        mesh_output.primitives[0].colorMask = vec4<f32>(1.0, 0.0, 1.0, 1.0);
+        return;
+    }
 }
 
 @fragment

naga-test/src/lib.rs

@@ -164,6 +164,11 @@ impl SpirvOutParameters {
             ray_query_initialization_tracking: true,
             debug_info,
             use_storage_input_output_16: self.use_storage_input_output_16,
+            // Choose
+            task_runtime_limits: Some(spv::TaskRuntimeLimits {
+                max_mesh_workgroups_per_dim: 256,
+                max_mesh_workgroups_total: 1024,
+            }),
         }
     }
 }

naga/src/back/spv/block.rs

@@ -221,12 +221,11 @@ impl Writer {
         ir_result: &crate::FunctionResult,
         result_members: &[ResultMember],
         body: &mut Vec<Instruction>,
-        task_payload: Option<Word>,
     ) -> Result<Instruction, Error> {
         for (index, res_member) in result_members.iter().enumerate() {
             // This isn't a real builtin, and is handled elsewhere
             if res_member.built_in == Some(crate::BuiltIn::MeshTaskSize) {
-                continue;
+                return Ok(Instruction::return_value(value_id));
             }
             let member_value_id = match ir_result.binding {
                 Some(_) => value_id,
@@ -258,13 +257,7 @@ impl Writer {
                 _ => {}
             }
         }
-        self.try_write_entry_point_task_return(
-            value_id,
-            ir_result,
-            result_members,
-            body,
-            task_payload,
-        )
+        Ok(Instruction::return_void())
     }
 }
 
@@ -3267,26 +3260,14 @@ impl BlockContext<'_> {
                             self.ir_function.result.as_ref().unwrap(),
                             &context.results,
                             &mut block.body,
-                            context.task_payload_variable_id,
                         )?,
                         None => Instruction::return_value(value_id),
                     };
                     self.function.consume(block, instruction);
                     return Ok(BlockExitDisposition::Discarded);
                 }
                 Statement::Return { value: None } => {
-                    if let Some(super::EntryPointContext {
-                        mesh_state: Some(ref mesh_state),
-                        ..
-                    }) = self.function.entry_point_context
-                    {
-                        self.function.consume(
-                            block,
-                            Instruction::branch(mesh_state.entry_point_epilogue_id),
-                        );
-                    } else {
-                        self.function.consume(block, Instruction::return_void());
-                    }
+                    self.function.consume(block, Instruction::return_void());
                     return Ok(BlockExitDisposition::Discarded);
                 }
                 Statement::Kill => {
@@ -3746,16 +3727,6 @@ impl BlockContext<'_> {
             LoopContext::default(),
             debug_info,
         )?;
-        if let Some(super::EntryPointContext {
-            mesh_state: Some(ref mesh_state),
-            ..
-        }) = self.function.entry_point_context
-        {
-            let mut block = Block::new(mesh_state.entry_point_epilogue_id);
-            self.writer
-                .write_mesh_shader_return(mesh_state, &mut block)?;
-            self.function.consume(block, Instruction::return_void());
-        }
 
         Ok(())
     }