Mesh shader HLSL writer (#8752)

Author: inner-daemons

CHANGELOG.md

@@ -52,6 +52,10 @@ Bottom level categories:
 
 - Unconditionally enable `Features::CLIP_DISTANCES`. By @ErichDonGubler in [#9270](https://github.com/gfx-rs/wgpu/pull/9270).
 
+#### DX12
+
+- Added support for mesh shaders in naga's HLSL writer, completing DX12 support for mesh shaders. By @inner-daemons in #8752.
+
 ### Changes
 
 #### General

docs/api-specs/mesh_shading.md

@@ -90,9 +90,16 @@ An example of using mesh shaders to render a single triangle can be seen [here](
 
 > **NOTE**: More limits will be added when support is added to `naga`.
 
-* `Limits::max_task_workgroup_total_count` - the maximum total number of workgroups from a `draw_mesh_tasks` command or similar. The dimensions passed must be less than or equal to this limit when multiplied together.
-* `Limits::max_task_workgroups_per_dimension` - the maximum for each of the 3 workgroup dimensions in a `draw_mesh_tasks` command. Each dimension passed must be less than or equal to this limit.
-* `max_mesh_multiview_count` - The maximum number of views used when multiview rendering with a mesh shader pipeline.
+* `Limits::max_task_mesh_workgroup_total_count` - the maximum total number of workgroups from a `draw_mesh_tasks` command or similar. The dimensions passed must be less than or equal to this limit when multiplied together.
+* `Limits::max_task_mesh_workgroups_per_dimension` - the maximum for each of the 3 workgroup dimensions in a `draw_mesh_tasks` command. Each dimension passed must be less than or equal to this limit.
+* `max_task_invocations_per_workgroup` - The maximum total number of threads in a task shader workgroup, given by `workgroupSize.x * workgroupSize.y * workgroupSize.z`.
+* `max_task_invocations_per_dimension` the maximum value for each of `workgroupSize.x`, `workgroupSize.y` and `workgroupSize.z` in task shader workgroups.
+* `max_mesh_invocations_per_workgroup` - The maximum total number of threads in a mesh shader workgroup, given by `workgroupSize.x * workgroupSize.y * workgroupSize.z`.
+* `max_mesh_invocations_per_dimension` the maximum value for each of `workgroupSize.x`, `workgroupSize.y` and `workgroupSize.z` in mesh shader workgroups.
+* `max_task_payload_size` - the size of an `var<task_payload>` variable, in bytes.
+* `max_mesh_output_vertices` - the maximum number of vertices that a single mesh shader workgroup may output.
+* `max_mesh_output_primitives` - the maximum number of primitives that a single mesh shader workgroup may output.
+* `max_mesh_multiview_count` - the maximum number of views used when multiview rendering with a mesh shader pipeline.
 * `max_mesh_output_layers` - the maximum number of output layers for a mesh shader pipeline.
 
 ## Naga implementation
@@ -132,8 +139,10 @@ A task shader entry point must also have a `@payload(G)` property, where `G` is
 
 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.
 
+The output of a task shader is set to zero if it violates either of the limits `max_task_mesh_workgroup_total_count` or `max_task_mesh_workgroups_per_dimension`.
+
 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.
+and `@builtin(num_workgroups)` matches the task shader workgroup's return value. If this output violates any limits, it may be zeroed or cause undefined behavior, depending on the compilation options. 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 can use compute and subgroup builtin inputs, in addition to `view_index` and `draw_id`.
 

examples/features/src/mesh_shader/mod.rs

@@ -11,37 +11,6 @@ fn compile_wgsl(device: &wgpu::Device) -> wgpu::ShaderModule {
         )
     }
 }
-fn compile_hlsl(device: &wgpu::Device, entry: &str, stage_str: &str) -> wgpu::ShaderModule {
-    let out_path = format!(
-        "{}/src/mesh_shader/shader.{stage_str}.cso",
-        env!("CARGO_MANIFEST_DIR")
-    );
-    let cmd = std::process::Command::new("dxc")
-        .args([
-            "-T",
-            &format!("{stage_str}_6_5"),
-            "-E",
-            entry,
-            &format!("{}/src/mesh_shader/shader.hlsl", env!("CARGO_MANIFEST_DIR")),
-            "-Fo",
-            &out_path,
-        ])
-        .output()
-        .unwrap();
-    if !cmd.status.success() {
-        panic!("DXC failed:\n{}", String::from_utf8(cmd.stderr).unwrap());
-    }
-    let file = std::fs::read(&out_path).unwrap();
-    std::fs::remove_file(out_path).unwrap();
-    unsafe {
-        device.create_shader_module_passthrough(wgpu::ShaderModuleDescriptorPassthrough {
-            label: None,
-            num_workgroups: (1, 1, 1),
-            dxil: Some(std::borrow::Cow::Owned(file)),
-            ..Default::default()
-        })
-    }
-}
 
 fn compile_msl(device: &wgpu::Device) -> wgpu::ShaderModule {
     unsafe {
@@ -67,7 +36,7 @@ fn get_shaders(device: &wgpu::Device, backend: wgpu::Backend) -> Shaders {
     // In the case that the platform does support mesh shaders, the dummy
     // shader is used to avoid requiring PASSTHROUGH_SHADERS.
     match backend {
-        wgpu::Backend::Vulkan => {
+        wgpu::Backend::Vulkan | wgpu::Backend::Dx12 => {
             let compiled = compile_wgsl(device);
             Shaders {
                 ts: compiled.clone(),
@@ -78,14 +47,6 @@ fn get_shaders(device: &wgpu::Device, backend: wgpu::Backend) -> Shaders {
                 fs_name: "fs_main",
             }
         }
-        wgpu::Backend::Dx12 => Shaders {
-            ts: compile_hlsl(device, "Task", "as"),
-            ms: compile_hlsl(device, "Mesh", "ms"),
-            fs: compile_hlsl(device, "Frag", "ps"),
-            ts_name: "main",
-            ms_name: "main",
-            fs_name: "main",
-        },
         wgpu::Backend::Metal => {
             let compiled = compile_msl(device);
             Shaders {

examples/features/src/mesh_shader/shader.hlsl

@@ -575,6 +575,8 @@ fn run() -> anyhow::Result<()> {
 
     params.spv_out.mesh_shader_primitive_indices_clamp = args.validate_mesh_output;
     params.spv_out.task_dispatch_limits = args.task_limits.0;
+    params.hlsl.mesh_shader_primitive_indices_clamp = args.validate_mesh_output;
+    params.hlsl.task_dispatch_limits = args.task_limits.0;
 
     if args.bulk_validate {
         return bulk_validate(&args, &params);

naga-cli/src/bin/naga.rs

@@ -50,6 +50,8 @@ pub struct Config {
     pub vertex: Vec<ConfigItem>,
     pub fragment: Vec<ConfigItem>,
     pub compute: Vec<ConfigItem>,
+    pub task: Vec<ConfigItem>,
+    pub mesh: Vec<ConfigItem>,
 }
 
 impl Config {
@@ -59,6 +61,8 @@ impl Config {
             vertex: Default::default(),
             fragment: Default::default(),
             compute: Default::default(),
+            task: Default::default(),
+            mesh: Default::default(),
         }
     }
 
@@ -85,8 +89,14 @@ impl Config {
             vertex,
             fragment,
             compute,
+            task,
+            mesh,
         } = self;
-        vertex.is_empty() && fragment.is_empty() && compute.is_empty()
+        vertex.is_empty()
+            && fragment.is_empty()
+            && compute.is_empty()
+            && task.is_empty()
+            && mesh.is_empty()
     }
 }
 

naga/hlsl-snapshots/src/lib.rs

@@ -149,8 +149,10 @@ impl crate::StorageFormat {
 }
 
 impl crate::BuiltIn {
-    pub(super) fn to_hlsl_str(self) -> Result<&'static str, Error> {
-        Ok(match self {
+    /// Returns `None` for "virtual" builtins, i.e. mesh shader builtins that are
+    /// used by naga but not recognized by HLSL.
+    pub(super) fn to_hlsl_str(self) -> Result<Option<&'static str>, Error> {
+        Ok(Some(match self {
             Self::Position { .. } => "SV_Position",
             // vertex
             Self::ClipDistances => "SV_ClipDistance",
@@ -186,12 +188,14 @@ impl crate::BuiltIn {
                 return Err(Error::Custom(format!("Unsupported builtin {self:?}")))
             }
             Self::CullPrimitive => "SV_CullPrimitive",
-            Self::PointIndex | Self::LineIndices | Self::TriangleIndices => unimplemented!(),
             Self::MeshTaskSize
             | Self::VertexCount
             | Self::PrimitiveCount
             | Self::Vertices
-            | Self::Primitives => unreachable!(),
+            | Self::Primitives
+            | Self::PointIndex
+            | Self::LineIndices
+            | Self::TriangleIndices => return Ok(None),
             Self::RayInvocationId
             | Self::NumRayInvocations
             | Self::InstanceCustomData
@@ -205,7 +209,7 @@ impl crate::BuiltIn {
             | Self::ObjectToWorld
             | Self::WorldToObject
             | Self::HitKind => unreachable!(),
-        })
+        }))
     }
 }