Add mesh shading info to naga IR

docs/api-specs/mesh_shading.md

@@ -2,13 +2,37 @@
 
 🧪Experimental🧪
 
-`wgpu` supports an experimental version of mesh shading. Currently `naga` has no support for mesh shaders beyond recognizing the additional shader stages.
+`wgpu` supports an experimental version of mesh shading when `Features::EXPERIMENTAL_MESH_SHADER` is enabled.
+Currently `naga` has no support for parsing or writing mesh shaders.
 For this reason, all shaders must be created with `Device::create_shader_module_passthrough`.
 
 **Note**: The features documented here may have major bugs in them and are expected to be subject
 to breaking changes, suggestions for the API exposed by this should be posted on [the mesh-shading issue](https://github.com/gfx-rs/wgpu/issues/7197).
 
-***This is not*** a thorough explanation of mesh shading and how it works. Those wishing to understand mesh shading more broadly should look elsewhere first.
+## Mesh shaders overview
+
+### What are mesh shaders
+Mesh shaders are a new kind of rasterization pipeline intended to address some of the shortfalls with the vertex shader pipeline. The core idea of mesh shaders is that the GPU decides how to render the many small parts of a scene instead of the CPU issuing a draw call for every small part or issuing an inefficient monolithic draw call for a large part of the scene.
+
+Mesh shaders are specifically designed to be used with **meshlet rendering**, a technique where every object is split into many subobjects called meshlets that are each rendered with their own parameters. With the standard vertex pipeline, each draw call specifies an exact number of primitives to render and the same parameters for all vertex shaders on an entire object (or even multiple objects). This doesn't leave room for different LODs for different parts of an object, for example a closer part having more detail, nor does it allow culling smaller sections (or primitives) of objects. With mesh shaders, each task workgroup might get assigned to a single object. It can then analyze the different meshlets(sections) of that object, determine which are visible and should actually be rendered, and for those meshlets determine what LOD to use based on the distance from the camera. It can then dispatch a mesh workgroup for each meshlet, with each mesh workgroup then reading the data for that specific LOD of its meshlet, determining which and how many vertices and primitives to output, determining which remaining primitives need to be culled, and passing the resulting primitives to the rasterizer.
+
+Mesh shaders are most effective in scenes with many polygons. They can allow skipping processing of entire groups of primitives that are facing away from the camera or otherwise occluded, which reduces the number of primitives that need to be processed by more than half in most cases, and they can reduce the number of primitives that need to be processed for more distant objects. Scenes that are not bottlenecked by geometry (perhaps instead by fragment processing or post processing) will not see much benefit from using them.
+
+Mesh shaders were first shown off in [NVIDIA's asteroids demo](https://www.youtube.com/watch?v=CRfZYJ_sk5E). Now, they form the basis for [Unreal Engine's Nanite](https://www.unrealengine.com/en-US/blog/unreal-engine-5-is-now-available-in-preview#Nanite).
+
+### Mesh shader pipeline
+A mesh shader pipeline is just like a standard render pipeline, except that the vertex shader stage is replaced by a mesh shader stage (and optionally a task shader stage). This functions as follows:
+
+* If there is a task shader stage, task shader workgroups are invoked first, with the number of workgroups determined by the draw call. Each task shader workgroup outputs a workgroup size and a task payload. A dispatch group of mesh shaders with the given workgroup size is then invoked with the task payload as a parameter.
+* Otherwise, a single dispatch group of mesh shaders with workgroup size given by the draw call is invoked.
+* Each mesh shader dispatch group functions exactly as a compute dispatch group, except that it has special outputs and may take a task payload as input. Mesh dispatch groups invoked by different task shader workgroups cannot interact.
+* Each workgroup within the mesh shader dispatch group can output vertices and primitives
+  * It determines how many vertices and primitives to write and then sets those vertices and primitives.
+  * Primitives have an indices field which determines the indices of the vertices of that primitive. The indices are based on the output of that mesh shader workgroup only; there is no sharing of vertices across workgroups (no vertex or index buffer equivalents).
+  * Primitives can then be culled by setting the appropriate builtin
+  * Each vertex output functions exactly as the output from a vertex shader would
+  * There can also be per-primitive outputs passed to fragment shaders; these are not interpolated or based on the vertices of the primitive in any way.
+* Once all of the primitives are written, those that weren't culled are are rasterized. From this point forward, the only difference from a standard render pipeline is that there may be some per-primitive inputs passed to fragment shaders.
 
 ## `wgpu` API
 
@@ -79,32 +103,36 @@ 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.
+Task shaders must be 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.
 
 ### 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, this is optional, and mesh shaders cannot write to this memory. Declaring `@payload` in a pipeline with no task shader or in a task shader with an `@payload` that is statically sized and differently than the mesh shader payload is illegal. The `@payload` attribute can only be ignored in pipelines that don't have a task shader.
 
-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
 
-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.
+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. All non-builtin primitive outputs must be decorated with `@per_primitive`.
 
 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`.
 
 Additionally, the `@location` attributes from the vertex and primitive outputs can't overlap.
 
-Before setting any vertices or indices, or exiting, the mesh shader must call `setMeshOutputs(numVertices: u32, numIndices: u32)`, which declares the number of vertices and indices that will be written to. These must be less than the corresponding maximums set in `@vertex_output` and `@primitive_output`. The mesh shader must then write to exactly these numbers of vertices and primitives.
+Before exiting, the mesh shader must call `setMeshOutputs(numVertices: u32, numIndices: u32)`, which declares the number of vertices and indices that will be written to. These must be less than the corresponding maximums set in `@vertex_output` and `@primitive_output`. The mesh shader must then write to exactly this range of vertices and primitives. A varying member with `@per_primitive` cannot be used in function interfaces except as a primitive output for mesh shaders or as input for fragment shaders.
 
 The mesh shader can write to vertices using the `setVertex(idx: u32, vertex: VertexOutput)` where `VertexOutput` is replaced with the vertex type declared in `@vertex_output`, and `idx` is the index of the vertex to write. Similarly, the mesh shader can write to vertices using `setPrimitive(idx: u32, primitive: PrimitiveOutput)`. These can be written to multiple times, however unsynchronized writes are undefined behavior. The primitives and indices are shared across the entire mesh shader workgroup.
 
 ### 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. Additionally, the locations of vertex and primitive input cannot overlap.
 
 ### Full example
 
@@ -114,9 +142,9 @@ The following is a full example of WGSL shaders that could be used to create a m
 enable mesh_shading;
 
 const positions = array(
-	vec4(0.,-1.,0.,1.),
-	vec4(-1.,1.,0.,1.),
-	vec4(1.,1.,0.,1.)
+	vec4(0.,1.,0.,1.),
+	vec4(-1.,-1.,0.,1.),
+	vec4(1.,-1.,0.,1.)
 );
 const colors = array(
 	vec4(0.,1.,0.,1.),
@@ -127,7 +155,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>,
@@ -136,14 +164,12 @@ struct VertexOutput {
 struct PrimitiveOutput {
 	@builtin(triangle_indices) index: vec3<u32>,
 	@builtin(cull_primitive) cull: bool,
-	@location(1) colorMask: vec4<f32>,
+	@per_primitive @location(1) colorMask: vec4<f32>,
 }
 struct PrimitiveInput {
-	@location(1) colorMask: vec4<f32>,
+	@per_primitive @location(1) colorMask: vec4<f32>,
 }
-fn test_function(input: u32) {
 
-}
 @task
 @payload(taskPayload)
 @workgroup_size(1)
@@ -162,8 +188,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);

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/features.rs

@@ -610,6 +610,7 @@ impl<W> Writer<'_, W> {
                     interpolation,
                     sampling,
                     blend_src,
+                    per_primitive: _,
                 } => {
                     if interpolation == Some(Interpolation::Linear) {
                         self.features.request(Features::NOPERSPECTIVE_QUALIFIER);

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,
         }
     }
 }
@@ -1300,6 +1301,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!(),
@@ -1614,6 +1618,7 @@ impl<'a, W: Write> Writer<'a, W> {
                 interpolation,
                 sampling,
                 blend_src,
+                per_primitive: _,
             } => (location, interpolation, sampling, blend_src),
             crate::Binding::BuiltIn(built_in) => {
                 match built_in {
@@ -1732,6 +1737,7 @@ impl<'a, W: Write> Writer<'a, W> {
                 interpolation: None,
                 sampling: None,
                 blend_src,
+                per_primitive: false,
             },
             stage: self.entry_point.stage,
             options: VaryingOptions::from_writer_options(self.options, output),
@@ -2669,6 +2675,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();
@@ -5247,6 +5258,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!()
+        }
     }
 }
 
@@ -5262,6 +5282,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

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