Various SPIR-V fixes for mesh shaders (gfx-rs#8756)

* Initial commit adding miscellaneous changes from msl-write and hlsl-write

* Same as previous commit

* Fixed divergence issue

* Removed some unnecessary barriers

* Zero init workgroup memory

* Added limits validation

* Added changelog

* Handled overflow, removed todo

* Lets see if this fixes llvmpipe

* Also this commit fixes llvmpipe maybe

* Unfortunate but not too unexpected at this point

* Updated feature to say to use ShaderRuntimeChecks::unchecked()

* Updated snapshots and took some changes from the hlsl writer

* Snapshots

* 2 quick tweaks

* Updated framework with suggestions by Connor in gfx-rs#8752

* Moved the task runtime limits into naga::back

* Fixed soem stuff

* Fixed checks: >= into >

* Cant believe I forgot this

* Removed code using u64 to check stuff

* Removed thing with limiting it to 2<<21

* Some more work

* Fixed compiles

* Added new field to spv options

* Updated some stuff to pass around the task runtime limits in more ways

* New PR started

* Did another quick fix

* Another quick fix

* Added changelog entry

* Added some explanation docs

* Fixed a warning

* Fixed shader

* Fixed some things & added docs

* Reverted dxc thing

* Fixed thing

* Refactored to TaskDispatchLimits

* Fixed compile error

Author: inner-daemons

CHANGELOG.md

@@ -69,6 +69,7 @@ By @cwfitzgerald in [#8999](https://github.com/gfx-rs/wgpu/pull/8999).
 - Added `Dx12BackendOptions::force_shader_model` to allow using advanced features in passthrough shaders without bundling DXC. By @inner-daemons in [#8984](https://github.com/gfx-rs/wgpu/pull/8984).
 - Changed passthrough shaders to not require an entry point parameter, so that the same shader module may be used in multiple entry points. Also added support for metallib passthrough. By @inner-daemons in #8886.
 - Added `Dx12Compiler::Auto` to automatically use static or dynamic DXC if available, before falling back to FXC. By @inner-daemons in [#8882](https://github.com/gfx-rs/wgpu/pull/8882).
+- Added support for `insert_debug_marker`, `push_debug_group` and `pop_debug_group` on WebGPU. By @evilpie in [#9017](https://github.com/gfx-rs/wgpu/pull/9017).
 - Added support for `@builtin(draw_index)` to the vulkan backend. By @inner-daemons in #8883.
 - Added support for `enable primitive_index` and `@builtin(primitive_index)` with support on all platforms. By @inner-daemons in #8879.
 
@@ -127,6 +128,9 @@ By @cwfitzgerald in [#8999](https://github.com/gfx-rs/wgpu/pull/8999).
 - Fix incorrect acceptance of some swizzle selectors that are not valid for their operand, e.g. `const v = vec2<i32>(); let r = v.xyz`. By @andyleiserson in [#8949](https://github.com/gfx-rs/wgpu/pull/8949).
 - Fixed calculation of the total number of bindings in a pipeline layout when validating against device limits. By @andyleiserson in [#8997](https://github.com/gfx-rs/wgpu/pull/8997).
 
+#### 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)

benches/benches/wgpu-benchmark/shader.rs

@@ -366,14 +366,16 @@ pub fn backends(ctx: BenchmarkContext) -> anyhow::Result<Vec<SubBenchResult>> {
         let mut data = Vec::new();
         let mut writer = naga::back::spv::Writer::new(&Default::default()).unwrap();
         for input in &inputs.inner {
+            let shared_info = WriterSharedOptions {
+                mesh_output_validation: input.options.mesh_output_validation,
+                task_limits: input.options.task_limits,
+                bounds_checks_policies: input.options.bounds_check_policies,
+            };
             if input.options.targets.unwrap().contains(Targets::SPIRV) {
                 if input.filename().contains("pointer-function-arg") {
                     continue;
                 }
-                let opt = input
-                    .options
-                    .spv
-                    .to_options(input.options.bounds_check_policies, None);
+                let opt = input.options.spv.to_options(&shared_info, None);
                 if writer.set_options(&opt).is_ok() {
                     let _ = writer.write(
                         input.module.as_ref().unwrap(),

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/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,21 @@ 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[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.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);
+    }
+    if thread_id.x < 3 {
+        mesh_output.vertices[thread_id.x].position = positions[thread_id.x];
+        mesh_output.vertices[thread_id.x].color = colors[thread_id.x] * taskPayload.colorMask;
+    }
 }
 
 @fragment

naga-cli/src/bin/naga.rs

@@ -152,6 +152,34 @@ struct Args {
     /// "67108864"), the string "none", or the string "all".
     #[argh(option, default = "CapabilitiesArg(naga::valid::Capabilities::all())")]
     capabilities: CapabilitiesArg,
+
+    /// the limits on the task shader dispatch size
+    #[argh(option, default = "TaskDispatchLimitsArg(None)")]
+    task_limits: TaskDispatchLimitsArg,
+
+    /// whether or not the mesh shader output should be validated.
+    #[argh(option, default = "true")]
+    validate_mesh_output: bool,
+}
+
+/// Newtype so we can implement [`FromStr`] for `Option<TaskDispatchLimits>`.
+#[derive(Debug, Clone, Copy)]
+struct TaskDispatchLimitsArg(Option<naga::back::TaskDispatchLimits>);
+
+impl FromStr for TaskDispatchLimitsArg {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let values = s
+            .split_once(",")
+            .ok_or_else(|| format!("No comma present for --task-limits value: {s}"))?;
+        let x = values.0.parse::<u32>().map_err(|e| e.to_string())?;
+        let y = values.1.parse::<u32>().map_err(|e| e.to_string())?;
+        Ok(Self(Some(naga::back::TaskDispatchLimits {
+            max_mesh_workgroups_per_dim: x,
+            max_mesh_workgroups_total: y,
+        })))
+    }
 }
 
 /// Newtype so we can implement [`FromStr`] for `BoundsCheckPolicy`.
@@ -545,6 +573,9 @@ fn run() -> anyhow::Result<()> {
     params.compact = args.compact;
     params.capabilities = args.capabilities.0;
 
+    params.spv_out.mesh_shader_primitive_indices_clamp = args.validate_mesh_output;
+    params.spv_out.task_dispatch_limits = args.task_limits.0;
+
     if args.bulk_validate {
         return bulk_validate(&args, &params);
     }

naga-test/src/lib.rs

@@ -75,6 +75,14 @@ where
     Ok(map)
 }
 
+#[derive(Default, serde::Deserialize)]
+#[serde(default)]
+pub struct WriterSharedOptions {
+    pub mesh_output_validation: bool,
+    pub task_limits: Option<naga::back::TaskDispatchLimits>,
+    pub bounds_checks_policies: naga::proc::BoundsCheckPolicies,
+}
+
 #[derive(Default, serde::Deserialize)]
 #[serde(default)]
 pub struct WgslInParameters {
@@ -137,7 +145,7 @@ impl Default for SpirvOutParameters {
 impl SpirvOutParameters {
     pub fn to_options<'a>(
         &'a self,
-        bounds_check_policies: naga::proc::BoundsCheckPolicies,
+        shared_info: &WriterSharedOptions,
         debug_info: Option<naga::back::spv::DebugInfo<'a>>,
     ) -> naga::back::spv::Options<'a> {
         use naga::back::spv;
@@ -157,14 +165,16 @@ impl SpirvOutParameters {
             } else {
                 Some(self.capabilities.clone())
             },
-            bounds_check_policies,
+            bounds_check_policies: shared_info.bounds_checks_policies,
             fake_missing_bindings: true,
             binding_map: self.binding_map.clone(),
             zero_initialize_workgroup_memory: spv::ZeroInitializeWorkgroupMemoryMode::Polyfill,
             force_loop_bounding: true,
             ray_query_initialization_tracking: true,
             debug_info,
             use_storage_input_output_16: self.use_storage_input_output_16,
+            task_dispatch_limits: shared_info.task_limits,
+            mesh_shader_primitive_indices_clamp: shared_info.mesh_output_validation,
         }
     }
 }
@@ -235,6 +245,17 @@ pub struct Parameters {
 
     pub bounds_check_policies: naga::proc::BoundsCheckPolicies,
     pub pipeline_constants: naga::back::PipelineConstants,
+
+    pub mesh_output_validation: bool,
+    #[serde(default = "default_task_limits")]
+    pub task_limits: Option<naga::back::TaskDispatchLimits>,
+}
+
+fn default_task_limits() -> Option<naga::back::TaskDispatchLimits> {
+    Some(naga::back::TaskDispatchLimits {
+        max_mesh_workgroups_per_dim: 256,
+        max_mesh_workgroups_total: 1024,
+    })
 }
 
 /// Information about a shader input file.

naga/src/back/mod.rs

@@ -148,6 +148,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>),
@@ -391,3 +392,11 @@ pub enum RayIntersectionType {
     Triangle = 1,
     BoundingBox = 4,
 }
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
+#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
+pub struct TaskDispatchLimits {
+    pub max_mesh_workgroups_per_dim: u32,
+    pub max_mesh_workgroups_total: u32,
+}

naga/src/back/spv/block.rs

@@ -237,12 +237,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,
@@ -274,13 +273,7 @@ impl Writer {
                 _ => {}
             }
         }
-        self.try_write_entry_point_task_return(
-            value_id,
-            ir_result,
-            result_members,
-            body,
-            task_payload,
-        )
+        Ok(Instruction::return_void())
     }
 }
 
@@ -3754,26 +3747,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 => {
@@ -4242,16 +4223,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(())
     }