Various SPIR-V fixes for mesh shaders

CHANGELOG.md

@@ -52,6 +52,9 @@ Bottom level categories:
 - BREAKING: Migrated from the `maxInterStageShaderComponents` limit to `maxInterStageShaderVariables`, which changes validation in a way that should not affect most programs. This follows the latest changes of the WebGPU spec. By @ErichDonGubler in [#8652](https://github.com/gfx-rs/wgpu/pull/8652), [#8792](https://github.com/gfx-rs/wgpu/pull/8792).
 - Fixed validation of the texture format in GPUDepthStencilState when neither depth nor stencil is actually enabled. By @andyleiserson in [#8766](https://github.com/gfx-rs/wgpu/pull/8766).
 
+#### 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)
+
 #### GLES
 
 - `DisplayHandle` should now be passed to `InstanceDescriptor` for correct EGL initialization on Wayland. By @MarijnS95 in [#8012](https://github.com/gfx-rs/wgpu/pull/8012)

docs/api-specs/mesh_shading.md

@@ -130,13 +130,11 @@ A task shader entry point must have a `@workgroup_size` attribute, meeting the s
 
 A task shader entry point must also have a `@payload(G)` property, where `G` is the name of a global variable in the `task_payload` address space. Each task shader workgroup has its own instance of this variable, visible to all invocations in the workgroup. Whatever value the workgroup collectively stores in that global variable becomes the **task payload**, and is provided to all invocations in the mesh shader grid dispatched for the workgroup. A task payload variable must be at least 4 bytes in size.
 
-A task shader entry point must return a `vec3<u32>` value. The return value of each workgroup's first invocation (that is, the one whose `local_invocation_index` is `0`) is taken as the size of a **mesh shader grid** to dispatch, measured in workgroups. (If the task shader entry point returns `vec3(0, 0, 0)`, then no mesh shaders are dispatched.) Mesh shader grids are described in the next section.
+A task shader entry point must return a `vec3<u32>` value decorated with `@builtin(mesh_task_size)`. The return value of each workgroup's first invocation (that is, the one whose `local_invocation_index` is `0`) is taken as the size of a **mesh shader grid** to dispatch, measured in workgroups. (If the task shader entry point returns `vec3(0, 0, 0)`, then no mesh shaders are dispatched.) Mesh shader grids are described in the next section.
 
 Each task shader workgroup dispatches an independent mesh shader grid: in mesh shader invocations, `@builtin` values like `workgroup_id` and `global_invocation_id` describe the position of the workgroup and invocation within that grid;
 and `@builtin(num_workgroups)` matches the task shader workgroup's return value. Mesh shaders dispatched for other task shader workgroups are not included in the count. If it is necessary for a mesh shader to know which task shader workgroup dispatched it, the task shader can include its own workgroup id in the task payload.
 
-Task shaders must return a value of type `vec3<u32>` decorated with `@builtin(mesh_task_size)`.
-
 Task shaders can use compute and subgroup builtin inputs, in addition to `view_index` and `draw_id`.
 
 ### Mesh shader

examples/features/src/framework.rs

@@ -260,11 +260,12 @@ 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()
             .with_display_handle(Box::new(
                 // TODO: Use event_loop.owned_display_handle() with winit 0.30
                 window.clone(),
             ));
+        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,10 @@ impl SpirvOutParameters {
             ray_query_initialization_tracking: true,
             debug_info,
             use_storage_input_output_16: self.use_storage_input_output_16,
+            task_runtime_limits: Some(spv::TaskRuntimeLimits {
+                max_mesh_workgroups_per_dim: 256,
+                max_mesh_workgroups_total: 1024,
+            }),
         }
     }
 }

naga/src/back/mod.rs

@@ -128,6 +128,7 @@ fn get_entry_points(
 /// [`EntryPoint`]: crate::EntryPoint
 /// [`Module`]: crate::Module
 /// [`Module::entry_points`]: crate::Module::entry_points
+#[derive(Clone, Copy, Debug)]
 pub enum FunctionType {
     /// A regular function.
     Function(crate::Handle<crate::Function>),

naga/src/back/spv/block.rs

@@ -222,12 +222,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,
@@ -259,13 +258,7 @@ impl Writer {
                 _ => {}
             }
         }
-        self.try_write_entry_point_task_return(
-            value_id,
-            ir_result,
-            result_members,
-            body,
-            task_payload,
-        )
+        Ok(Instruction::return_void())
     }
 }
 
@@ -3342,26 +3335,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 => {
@@ -3857,16 +3838,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(())
     }