Initial naga changes for mesh shaders

CHANGELOG.md

@@ -53,6 +53,9 @@ We have merged the acceleration structure feature into the `RayQuery` feature. T
 
 By @Vecvec in [#7913](https://github.com/gfx-rs/wgpu/pull/7913).
 
+### Naga
+- Added mesh shader support to naga with `WGSL` frontend and `SPIR-V` backend. By @SupaMaggie70Incorporated in [#7930](https://github.com/gfx-rs/wgpu/pull/7930).
+
 ## v26.0.1 (2025-07-10)
 
 ### Bug Fixes

docs/api-specs/mesh_shading.md

@@ -80,19 +80,21 @@ This shader stage can be selected by marking a function with `@task`. Task shade
 
 The output of this determines how many workgroups of mesh shaders will be dispatched. Once dispatched, global id variables will be local to the task shader workgroup dispatch, and mesh shaders won't know the position of their dispatch among all mesh shader dispatches unless this is passed through the payload. The output may be zero to skip dispatching any mesh shader workgroups for the task shader workgroup.
 
-If task shaders are marked with `@payload(someVar)`, where `someVar` is global variable declared like `var<workgroup> someVar: <type>`, task shaders may write to `someVar`. This payload is passed to the mesh shader workgroup that is invoked. The mesh shader can skip declaring `@payload` to ignore this input.
+If task shaders are marked with `@payload(someVar)`, where `someVar` is global variable declared like `var<task_payload> someVar: <type>`, task shaders may use `someVar` as if it is a read-write workgroup storage variable. This payload is passed to the mesh shader workgroup that is invoked. The mesh shader can skip declaring `@payload` to ignore this input.
 
 ### Mesh shader
 This shader stage can be selected by marking a function with `@mesh`. Mesh shaders must not return anything.
 
-Mesh shaders can be marked with `@payload(someVar)` similar to task shaders. Unlike task shaders, mesh shaders cannot write to this workgroup memory. Declaring `@payload` in a pipeline with no task shader, in a pipeline with a task shader that doesn't declare `@payload`, or in a task shader with an `@payload` that is statically sized and smaller than the mesh shader payload is illegal.
+Mesh shaders can be marked with `@payload(someVar)` similar to task shaders. Unlike task shaders, mesh shaders cannot write to this memory. Declaring `@payload` in a pipeline with no task shader, in a pipeline with a task shader that doesn't declare `@payload`, or in a task shader with an `@payload` that is statically sized and smaller than the mesh shader payload is illegal.
 
-Mesh shaders must be marked with `@vertex_output(OutputType, numOutputs)`, where `numOutputs` is the maximum number of vertices to be output by a mesh shader, and `OutputType` is the data associated with vertices, similar to a standard vertex shader output.
+Mesh shaders must be marked with `@vertex_output(OutputType, numOutputs)`, where `numOutputs` is the maximum number of vertices to be output by a mesh shader, and `OutputType` is the data associated with vertices, similar to a standard vertex shader output, and must be a struct.
 
 Mesh shaders must also be marked with `@primitive_output(OutputType, numOutputs)`, which is similar to `@vertex_output` except it describes the primitive outputs.
 
 ### Mesh shader outputs
 
+Vertex outputs from mesh shaders function identically to outputs of vertex shaders, and as such must have a field with `@builtin(position)`.
+
 Primitive outputs from mesh shaders have some additional builtins they can set. These include `@builtin(cull_primitive)`, which must be a boolean value. If this is set to true, then the primitive is skipped during rendering.
 
 Mesh shader primitive outputs must also specify exactly one of `@builtin(triangle_indices)`, `@builtin(line_indices)`, or `@builtin(point_index)`. This determines the output topology of the mesh shader, and must match the output topology of the pipeline descriptor the mesh shader is used with. These must be of type `vec3<u32>`, `vec2<u32>`, and `u32` respectively. When setting this, each of the indices must be less than the number of vertices declared in `setMeshOutputs`.
@@ -105,7 +107,9 @@ The mesh shader can write to vertices using the `setVertex(idx: u32, vertex: Ver
 
 ### Fragment shader
 
-Fragment shaders may now be passed the primitive info from a mesh shader the same was as they are passed vertex inputs, for example `fn fs_main(vertex: VertexOutput, primitive: PrimitiveOutput)`. The primitive state is part of the fragment input and must match the output of the mesh shader in the pipeline.
+Fragment shaders can access vertex output data as if it is from a vertex shader. They can also access primitive output data, provided the input is decorated with `@per_primitive`. The `@per_primitive` attribute can be applied to a value directly, such as `@per_primitive @location(1) value: vec4<f32>`, to a struct such as `@per_primitive primitive_input: PrimitiveInput` where `PrimitiveInput` is a struct containing fields decorated with `@location` and `@builtin`, or to members of a struct that are themselves decorated with `@location` or `@builtin`.
+
+The primitive state is part of the fragment input and must match the output of the mesh shader in the pipeline. Using `@per_primitive` also requires enabling the mesh shader extension.
 
 ### Full example
 
@@ -128,7 +132,7 @@ struct TaskPayload {
 	colorMask: vec4<f32>,
 	visible: bool,
 }
-var<workgroup> taskPayload: TaskPayload;
+var<task_payload> taskPayload: TaskPayload;
 var<workgroup> workgroupData: f32;
 struct VertexOutput {
 	@builtin(position) position: vec4<f32>,
@@ -142,9 +146,7 @@ struct PrimitiveOutput {
 struct PrimitiveInput {
 	@location(1) colorMask: vec4<f32>,
 }
-fn test_function(input: u32) {
 
-}
 @task
 @payload(taskPayload)
 @workgroup_size(1)
@@ -163,8 +165,6 @@ fn ms_main(@builtin(local_invocation_index) index: u32, @builtin(global_invocati
 	workgroupData = 2.0;
 	var v: VertexOutput;
 
-	test_function(1);
-
 	v.position = positions[0];
 	v.color = colors[0] * taskPayload.colorMask;
 	setVertex(0, v);
@@ -184,7 +184,7 @@ fn ms_main(@builtin(local_invocation_index) index: u32, @builtin(global_invocati
 	setPrimitive(0, p);
 }
 @fragment
-fn fs_main(vertex: VertexOutput, primitive: PrimitiveInput) -> @location(0) vec4<f32> {
+fn fs_main(vertex: VertexOutput, @per_primitive primitive: PrimitiveInput) -> @location(0) vec4<f32> {
 	return vertex.color * primitive.colorMask;
 }
 ```

naga-cli/src/bin/naga.rs

@@ -64,6 +64,12 @@ struct Args {
     #[argh(option)]
     shader_model: Option<ShaderModelArg>,
 
+    /// the SPIR-V version to use if targeting SPIR-V
+    ///
+    /// For example, 1.0, 1.4, etc
+    #[argh(option)]
+    spirv_version: Option<SpirvVersionArg>,
+
     /// the shader stage, for example 'frag', 'vert', or 'compute'.
     /// if the shader stage is unspecified it will be derived from
     /// the file extension.
@@ -189,6 +195,22 @@ impl FromStr for ShaderModelArg {
     }
 }
 
+#[derive(Debug, Clone)]
+struct SpirvVersionArg(u8, u8);
+
+impl FromStr for SpirvVersionArg {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let dot = s
+            .find(".")
+            .ok_or_else(|| "Missing dot separator".to_owned())?;
+        let major = s[..dot].parse::<u8>().map_err(|e| e.to_string())?;
+        let minor = s[dot + 1..].parse::<u8>().map_err(|e| e.to_string())?;
+        Ok(Self(major, minor))
+    }
+}
+
 /// Newtype so we can implement [`FromStr`] for `ShaderSource`.
 #[derive(Debug, Clone, Copy)]
 struct ShaderStage(naga::ShaderStage);
@@ -465,6 +487,9 @@ fn run() -> anyhow::Result<()> {
     if let Some(ref version) = args.metal_version {
         params.msl.lang_version = version.0;
     }
+    if let Some(ref version) = args.spirv_version {
+        params.spv_out.lang_version = (version.0, version.1);
+    }
     params.keep_coordinate_space = args.keep_coordinate_space;
 
     params.dot.cfg_only = args.dot_cfg_only;

naga/src/back/dot/mod.rs

@@ -307,6 +307,25 @@ impl StatementGraph {
                         crate::RayQueryFunction::Terminate => "RayQueryTerminate",
                     }
                 }
+                S::MeshFunction(crate::MeshFunction::SetMeshOutputs {
+                    vertex_count,
+                    primitive_count,
+                }) => {
+                    self.dependencies.push((id, vertex_count, "vertex_count"));
+                    self.dependencies
+                        .push((id, primitive_count, "primitive_count"));
+                    "SetMeshOutputs"
+                }
+                S::MeshFunction(crate::MeshFunction::SetVertex { index, value }) => {
+                    self.dependencies.push((id, index, "index"));
+                    self.dependencies.push((id, value, "value"));
+                    "SetVertex"
+                }
+                S::MeshFunction(crate::MeshFunction::SetPrimitive { index, value }) => {
+                    self.dependencies.push((id, index, "index"));
+                    self.dependencies.push((id, value, "value"));
+                    "SetPrimitive"
+                }
                 S::SubgroupBallot { result, predicate } => {
                     if let Some(predicate) = predicate {
                         self.dependencies.push((id, predicate, "predicate"));

naga/src/back/glsl/mod.rs

@@ -139,7 +139,8 @@ impl crate::AddressSpace {
             | crate::AddressSpace::Uniform
             | crate::AddressSpace::Storage { .. }
             | crate::AddressSpace::Handle
-            | crate::AddressSpace::PushConstant => false,
+            | crate::AddressSpace::PushConstant
+            | crate::AddressSpace::TaskPayload => false,
         }
     }
 }
@@ -1299,6 +1300,9 @@ impl<'a, W: Write> Writer<'a, W> {
             crate::AddressSpace::Storage { .. } => {
                 self.write_interface_block(handle, global)?;
             }
+            crate::AddressSpace::TaskPayload => {
+                self.write_interface_block(handle, global)?;
+            }
             // A global variable in the `Function` address space is a
             // contradiction in terms.
             crate::AddressSpace::Function => unreachable!(),
@@ -2668,6 +2672,11 @@ impl<'a, W: Write> Writer<'a, W> {
                 self.write_image_atomic(ctx, image, coordinate, array_index, fun, value)?
             }
             Statement::RayQuery { .. } => unreachable!(),
+            Statement::MeshFunction(
+                crate::MeshFunction::SetMeshOutputs { .. }
+                | crate::MeshFunction::SetVertex { .. }
+                | crate::MeshFunction::SetPrimitive { .. },
+            ) => unreachable!(),
             Statement::SubgroupBallot { result, predicate } => {
                 write!(self.out, "{level}")?;
                 let res_name = Baked(result).to_string();
@@ -5242,6 +5251,15 @@ const fn glsl_built_in(built_in: crate::BuiltIn, options: VaryingOptions) -> &'s
         Bi::SubgroupId => "gl_SubgroupID",
         Bi::SubgroupSize => "gl_SubgroupSize",
         Bi::SubgroupInvocationId => "gl_SubgroupInvocationID",
+        // mesh
+        // TODO: figure out how to map these to glsl things as glsl treats them as arrays
+        Bi::CullPrimitive
+        | Bi::PointIndex
+        | Bi::LineIndices
+        | Bi::TriangleIndices
+        | Bi::MeshTaskSize => {
+            unimplemented!()
+        }
     }
 }
 
@@ -5257,6 +5275,7 @@ const fn glsl_storage_qualifier(space: crate::AddressSpace) -> Option<&'static s
         As::Handle => Some("uniform"),
         As::WorkGroup => Some("shared"),
         As::PushConstant => Some("uniform"),
+        As::TaskPayload => unreachable!(),
     }
 }
 

naga/src/back/hlsl/conv.rs

@@ -183,6 +183,9 @@ impl crate::BuiltIn {
             Self::PointSize | Self::ViewIndex | Self::PointCoord | Self::DrawID => {
                 return Err(Error::Custom(format!("Unsupported builtin {self:?}")))
             }
+            Self::CullPrimitive => "SV_CullPrimitive",
+            Self::PointIndex | Self::LineIndices | Self::TriangleIndices => unimplemented!(),
+            Self::MeshTaskSize => unreachable!(),
         })
     }
 }

naga/src/back/hlsl/mod.rs

@@ -223,7 +223,8 @@ impl crate::ShaderStage {
             Self::Vertex => "vs",
             Self::Fragment => "ps",
             Self::Compute => "cs",
-            Self::Task | Self::Mesh => unreachable!(),
+            Self::Task => "ts",
+            Self::Mesh => "ms",
         }
     }
 }

naga/src/back/hlsl/writer.rs

@@ -495,7 +495,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
 
             self.write_wrapped_functions(module, &ctx)?;
 
-            if ep.stage == ShaderStage::Compute {
+            if ep.stage.compute_like() {
                 // HLSL is calling workgroup size "num threads"
                 let num_threads = ep.workgroup_size;
                 writeln!(
@@ -941,6 +941,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 self.write_type(module, global.ty)?;
                 ""
             }
+            crate::AddressSpace::TaskPayload => unimplemented!(),
             crate::AddressSpace::Uniform => {
                 // constant buffer declarations are expected to be inlined, e.g.
                 // `cbuffer foo: register(b0) { field1: type1; }`
@@ -2459,6 +2460,19 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     writeln!(self.out, ".Abort();")?;
                 }
             },
+            Statement::MeshFunction(crate::MeshFunction::SetMeshOutputs {
+                vertex_count,
+                primitive_count,
+            }) => {
+                write!(self.out, "{level}SetMeshOutputCounts(")?;
+                self.write_expr(module, vertex_count, func_ctx)?;
+                write!(self.out, ", ")?;
+                self.write_expr(module, primitive_count, func_ctx)?;
+                write!(self.out, ");")?;
+            }
+            Statement::MeshFunction(
+                crate::MeshFunction::SetVertex { .. } | crate::MeshFunction::SetPrimitive { .. },
+            ) => unimplemented!(),
             Statement::SubgroupBallot { result, predicate } => {
                 write!(self.out, "{level}")?;
                 let name = Baked(result).to_string();
@@ -2936,7 +2950,8 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                                 crate::AddressSpace::Function
                                 | crate::AddressSpace::Private
                                 | crate::AddressSpace::WorkGroup
-                                | crate::AddressSpace::PushConstant,
+                                | crate::AddressSpace::PushConstant
+                                | crate::AddressSpace::TaskPayload,
                             )
                             | None => true,
                             Some(crate::AddressSpace::Uniform) => {

naga/src/back/msl/mod.rs

@@ -651,6 +651,10 @@ impl ResolvedBinding {
                     Bi::CullDistance | Bi::ViewIndex | Bi::DrawID => {
                         return Err(Error::UnsupportedBuiltIn(built_in))
                     }
+                    Bi::CullPrimitive => "primitive_culled",
+                    // TODO: figure out how to make this written as a function call
+                    Bi::PointIndex | Bi::LineIndices | Bi::TriangleIndices => unimplemented!(),
+                    Bi::MeshTaskSize => unreachable!(),
                 };
                 write!(out, "{name}")?;
             }

naga/src/back/msl/writer.rs

@@ -577,7 +577,8 @@ impl crate::AddressSpace {
             | Self::Private
             | Self::WorkGroup
             | Self::PushConstant
-            | Self::Handle => true,
+            | Self::Handle
+            | Self::TaskPayload => true,
             Self::Function => false,
         }
     }
@@ -590,6 +591,7 @@ impl crate::AddressSpace {
             // may end up with "const" even if the binding is read-write,
             // and that should be OK.
             Self::Storage { .. } => true,
+            Self::TaskPayload => unimplemented!(),
             // These should always be read-write.
             Self::Private | Self::WorkGroup => false,
             // These translate to `constant` address space, no need for qualifiers.
@@ -606,6 +608,7 @@ impl crate::AddressSpace {
             Self::Storage { .. } => Some("device"),
             Self::Private | Self::Function => Some("thread"),
             Self::WorkGroup => Some("threadgroup"),
+            Self::TaskPayload => Some("object_data"),
         }
     }
 }
@@ -4019,6 +4022,14 @@ impl<W: Write> Writer<W> {
                         }
                     }
                 }
+                // TODO: write emitters for these
+                crate::Statement::MeshFunction(crate::MeshFunction::SetMeshOutputs { .. }) => {
+                    unimplemented!()
+                }
+                crate::Statement::MeshFunction(
+                    crate::MeshFunction::SetVertex { .. }
+                    | crate::MeshFunction::SetPrimitive { .. },
+                ) => unimplemented!(),
                 crate::Statement::SubgroupBallot { result, predicate } => {
                     write!(self.out, "{level}")?;
                     let name = self.namer.call("");
@@ -6168,7 +6179,7 @@ template <typename A>
                     LocationMode::Uniform,
                     false,
                 ),
-                crate::ShaderStage::Task | crate::ShaderStage::Mesh => unreachable!(),
+                crate::ShaderStage::Task | crate::ShaderStage::Mesh => unimplemented!(),
             };
 
             // Should this entry point be modified to do vertex pulling?
@@ -6231,6 +6242,9 @@ template <typename A>
                                 break;
                             }
                         }
+                        crate::AddressSpace::TaskPayload => {
+                            unimplemented!()
+                        }
                         crate::AddressSpace::Function
                         | crate::AddressSpace::Private
                         | crate::AddressSpace::WorkGroup => {}
@@ -7157,7 +7171,7 @@ mod workgroup_mem_init {
             fun_info: &valid::FunctionInfo,
         ) -> bool {
             options.zero_initialize_workgroup_memory
-                && ep.stage == crate::ShaderStage::Compute
+                && ep.stage.compute_like()
                 && module.global_variables.iter().any(|(handle, var)| {
                     !fun_info[handle].is_empty() && var.space == crate::AddressSpace::WorkGroup
                 })