review wip

Author: jimblandy

docs/api-specs/mesh_shading.md

@@ -12,6 +12,7 @@ to breaking changes, suggestions for the API exposed by this should be posted on
 ## 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.
@@ -21,7 +22,14 @@ Mesh shaders are most effective in scenes with many polygons. They can allow ski
 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:
+
+A mesh draw command like `RenderPass::draw_mesh_tasks` uses a `RenderPipeline` created using
+`Device::create_mesh_pipeline`, which runs the following stages:
+
+- First, an optional **task shader stage** decides how many mesh shader grids to dispatch
+
+
+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.
@@ -99,9 +107,17 @@ Using any of these features in a `wgsl` program will require adding the `enable
 Two new shader stages will be added to `WGSL`. Fragment shaders are also modified slightly. Both task shaders and mesh shaders are allowed to use any compute-specific functionality, such as subgroup operations.
 
 ### Task shader
-This shader stage can be selected by marking a function with `@task`. Task shaders must return a `vec3<u32>` as their output type. Similar to compute shaders, task shaders run in a workgroup. The output must be uniform across all threads in a workgroup.
 
-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.
+A function with the `@task` attribute is a **task shader stage entry point**. A mesh shader pipeline may optionally specify a task shader entry point; if it does, then mesh draw commands using that pipeline dispatch a **mesh shader grid** of invocations running the task shader entry point to compute the sizes of mesh shader grids to dispatch, and to provide a payload value for each dispatch.
+
+Dispatching a task shader runs a grid of workgroups, like a compute shader dispatch. The mesh draw command determines the number of workgroups along each axis of the grid.
+
+A task shader entry point must return a `vec3<u32>`. All invocations within a task shader workgroup must return the same value, but different workgroups may return different values. After each task shader workgroup finishes, the implementation dispatches a grid of mesh shaders, whose size in workgroups is determined by the task shader workgroup's return value. If the task shader returns `vec3(0, 0, 0)`, then no mesh shaders are dispatched.
+
+Each task shader workgroup produces an independent dispatch grid of mesh shaders: `@builtin` values like `workgroup_id`, `global_invocation_id` describe the position of the workgroup and invocation within that grid. Mesh shaders dispatched for other task shader workgroups are not included in the count.
+Similarly, `@builtin(num_workgroups)` matches the task shader workgroup's return value.
+
+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.
 
 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.
 
@@ -210,4 +226,4 @@ fn ms_main(@builtin(local_invocation_index) index: u32, @builtin(global_invocati
 fn fs_main(vertex: VertexOutput, primitive: PrimitiveInput) -> @location(0) vec4<f32> {
 	return vertex.color * primitive.colorMask;
 }
-```
+```

naga/src/ir/mod.rs

@@ -320,14 +320,21 @@ pub enum ConservativeDepth {
 #[cfg_attr(feature = "serialize", derive(Serialize))]
 #[cfg_attr(feature = "deserialize", derive(Deserialize))]
 #[cfg_attr(feature = "arbitrary", derive(Arbitrary))]
-#[allow(missing_docs)] // The names are self evident
 pub enum ShaderStage {
+    /// Render pipeline vertex shader.
     Vertex,
-    Fragment,
-    Compute,
-    // Mesh shader stages
+
+    /// Render pipeline task shader.
     Task,
+
+    /// Render pipeline mesh shader.
     Mesh,
+
+    /// Render pipeline fragment shader.
+    Fragment,
+
+    /// Compute pipeline shader.
+    Compute,
 }
 
 impl ShaderStage {
@@ -964,6 +971,9 @@ pub enum Binding {
 
     /// Indexed location.
     ///
+    /// This is a value passed to a [`Fragment`] shader from a [`Vertex`] or
+    /// [`Mesh`] shader.
+    ///
     /// Values passed from the [`Vertex`] stage to the [`Fragment`] stage must
     /// have their `interpolation` defaulted (i.e. not `None`) by the front end
     /// as appropriate for that language.
@@ -977,6 +987,7 @@ pub enum Binding {
     /// interpolation must be `Flat`.
     ///
     /// [`Vertex`]: crate::ShaderStage::Vertex
+    /// [`Mesh`]: crate::ShaderStage::Mesh
     /// [`Fragment`]: crate::ShaderStage::Fragment
     Location {
         location: u32,
@@ -1751,10 +1762,12 @@ pub enum Expression {
         query: Handle<Expression>,
         committed: bool,
     },
+
     /// Result of a [`SubgroupBallot`] statement.
     ///
     /// [`SubgroupBallot`]: Statement::SubgroupBallot
     SubgroupBallotResult,
+
     /// Result of a [`SubgroupCollectiveOperation`] or [`SubgroupGather`] statement.
     ///
     /// [`SubgroupCollectiveOperation`]: Statement::SubgroupCollectiveOperation
@@ -2343,7 +2356,9 @@ pub struct EntryPoint {
     pub workgroup_size_overrides: Option<[Option<Handle<Expression>>; 3]>,
     /// The entrance function.
     pub function: Function,
-    /// The information relating to a mesh shader
+    /// Information for [`Mesh`] shaders.
+    ///
+    /// [`Mesh`]: ShaderStage::Mesh
     pub mesh_info: Option<MeshStageInfo>,
     /// The unique global variable used as a task payload from task shader to mesh shader
     pub task_payload: Option<Handle<GlobalVariable>>,

naga/src/valid/analyzer.rs

@@ -1151,7 +1151,7 @@ impl FunctionInfo {
                         let _ = self.add_ref(index);
                         let _ = self.add_ref(value);
                         let ty =
-                            self.expressions[value.index()].ty.clone().handle().ok_or(
+                            self.expressions[value.index()].ty.handle().ok_or(
                                 FunctionError::InvalidMeshShaderOutputType(value).with_span(),
                             )?;
 
@@ -1210,6 +1210,7 @@ impl FunctionInfo {
         Ok(combined_uniformity)
     }
 
+    /// Note that this function supplies vertex 
     fn try_update_mesh_vertex_type(
         &mut self,
         ty: Handle<crate::Type>,
@@ -1244,14 +1245,15 @@ impl FunctionInfo {
         Ok(())
     }
 
+    /// Update this function's mesh shader info, given that it calls `callee`.
     fn try_update_mesh_info(
         &mut self,
-        other: &FunctionMeshShaderInfo,
+        callee: &FunctionMeshShaderInfo,
     ) -> Result<(), WithSpan<FunctionError>> {
-        if let &Some(ref other_vertex) = &other.vertex_type {
+        if let &Some(ref other_vertex) = &callee.vertex_type {
             self.try_update_mesh_vertex_type(other_vertex.0, other_vertex.1)?;
         }
-        if let &Some(ref other_primitive) = &other.vertex_type {
+        if let &Some(ref other_primitive) = &callee.vertex_type {
             self.try_update_mesh_primitive_type(other_primitive.0, other_primitive.1)?;
         }
         Ok(())

naga/src/valid/interface.rs

@@ -167,10 +167,12 @@ fn storage_usage(access: crate::StorageAccess) -> GlobalUse {
     storage_usage
 }
 
+/// An output from a mesh shader.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 enum MeshOutputType {
     None,
     VertexOutput,
+    /// 
     PrimitiveOutput,
 }
 
@@ -856,13 +858,14 @@ impl super::Validator {
             {
                 return Err(EntryPointError::UnexpectedMeshShaderEntryResult.with_span());
             }
-            // Cannot have any other built-ins or @location outputs as those are per-vertex or per-primitive
-            if ep.stage == crate::ShaderStage::Task
-                && (!result_built_ins.contains(&crate::BuiltIn::MeshTaskSize)
-                    || result_built_ins.len() != 1
-                    || !self.location_mask.is_empty())
-            {
-                return Err(EntryPointError::WrongTaskShaderEntryResult.with_span());
+            // Task shaders must have a single `MeshTaskSize` output, and nothing else.
+            if ep.stage == crate::ShaderStage::Task {
+                let ok = result_built_ins.contains(&crate::BuiltIn::MeshTaskSize)
+                     && result_built_ins.len() == 1
+                     && self.location_mask.is_empty();
+                if !ok {
+                    return Err(EntryPointError::WrongTaskShaderEntryResult.with_span());
+                }
             }
             if !self.blend_src_mask.is_empty() {
                 info.dual_source_blending = true;
@@ -960,8 +963,10 @@ impl super::Validator {
             }
         }
 
+        // If this is a `Mesh` entry point, check its interface.
         if let &Some(ref mesh_info) = &ep.mesh_info {
-            // Technically it is allowed to not output anything
+            // Mesh shaders don't return any value. All their results are supplied through
+            // [`SetVertex`] and [`SetPrimitive`] calls.
             // TODO: check that only the allowed builtins are used here
             if let Some(used_vertex_type) = info.mesh_shader_info.vertex_type {
                 if used_vertex_type.0 != mesh_info.vertex_output_type {

wgpu/src/api/render_pass.rs

@@ -231,9 +231,9 @@ impl RenderPass<'_> {
         self.inner.draw_indexed(indices, base_vertex, instances);
     }
 
-    /// Draws using a mesh shader pipeline.
+    /// Draws using a mesh pipeline.
     ///
-    /// The current pipeline must be a mesh shader pipeline.
+    /// The current pipeline must be a mesh pipeline.
     ///
     /// If the current pipeline has a task shader, run it with an workgroup for
     /// every `vec3<u32>(i, j, k)` where `i`, `j`, and `k` are between `0` and
@@ -290,7 +290,7 @@ impl RenderPass<'_> {
             .draw_indexed_indirect(&indirect_buffer.inner, indirect_offset);
     }
 
-    /// Draws using a mesh shader pipeline,
+    /// Draws using a mesh pipeline,
     /// based on the contents of the `indirect_buffer`
     ///
     /// This is like calling [`RenderPass::draw_mesh_tasks`] but the contents of the call are specified in the `indirect_buffer`.

wgpu/src/api/render_pipeline.rs

@@ -152,13 +152,15 @@ static_assertions::assert_impl_all!(FragmentState<'_>: Send, Sync);
 pub struct TaskState<'a> {
     /// The compiled shader module for this stage.
     pub module: &'a ShaderModule,
-    /// The name of the entry point in the compiled shader to use.
+
+    /// The name of the task shader entry point in the shader module to use.
     ///
-    /// If [`Some`], there must be a vertex-stage shader entry point with this name in `module`.
-    /// Otherwise, expect exactly one vertex-stage entry point in `module`, which will be
-    /// selected.
+    /// If [`Some`], there must be a task shader entry point with the given name
+    /// in `module`. Otherwise, there must be exactly one task shader entry
+    /// point in `module`, which will be selected.
     pub entry_point: Option<&'a str>,
-    /// Advanced options for when this pipeline is compiled
+
+    /// Advanced options for when this pipeline is compiled.
     ///
     /// This implements `Default`, and for most users can be set to `Default::default()`
     pub compilation_options: PipelineCompilationOptions<'a>,
@@ -299,8 +301,15 @@ pub struct MeshPipelineDescriptor<'a> {
     ///
     /// [default layout]: https://www.w3.org/TR/webgpu/#default-pipeline-layout
     pub layout: Option<&'a PipelineLayout>,
-    /// The compiled task stage and its entry point.
+
+    /// The mesh pipeline's task shader.
+    ///
+    /// If this is `None`, the mesh pipeline has no task shader. Executing a
+    /// mesh drawing command simply dispatches a grid of mesh shaders directly.
+    ///
+    /// [`draw_mesh_tasks`]: RenderPass::draw_mesh_tasks
     pub task: Option<TaskState<'a>>,
+
     /// The compiled mesh stage and its entry point
     pub mesh: MeshState<'a>,
     /// The properties of the pipeline at the primitive assembly and rasterization level.