Mesh shader HLSL writer

CHANGELOG.md

@@ -47,6 +47,12 @@ Bottom level categories:
 - Added support for obtaining `AdapterInfo` from `Device`. By @sagudev in [#8807](https://github.com/gfx-rs/wgpu/pull/8807).
 - Added `Limits::or_worse_values_from`. By @atlv24 in [#8870](https://github.com/gfx-rs/wgpu/pull/8870).
 
+### General
+
+#### DX12
+
+- Full support for mesh shaders in HLSL/DX12. By @inner-daemons in #8752.
+
 ### Bug Fixes
 
 #### General
@@ -60,6 +66,9 @@ Bottom level categories:
 - The validator checks that override-sized arrays have a positive size, if overrides have been resolved. By @andyleiserson in [#8822](https://github.com/gfx-rs/wgpu/pull/8822).
 - Fix some cases where f16 constants were not working. By @andyleiserson in [#8816](https://github.com/gfx-rs/wgpu/pull/8816).
 
+#### 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)

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

@@ -260,11 +260,13 @@ impl ExampleContext {
     async fn init_async<E: Example>(surface: &mut SurfaceWrapper, window: Arc<Window>) -> Self {
         log::info!("Initializing wgpu...");
 
-        let instance_descriptor = wgpu::InstanceDescriptor::from_env_or_default()
-            .with_display_handle(Box::new(
-                // TODO: Use event_loop.owned_display_handle() with winit 0.30
-                window.clone(),
-            ));
+        let mut instance_descriptor = wgpu::InstanceDescriptor::default();
+        // Use static DXC by default so we can utilize the latest features
+        instance_descriptor.backend_options.dx12.shader_compiler = wgpu::Dx12Compiler::StaticDxc;
+        let instance_descriptor = instance_descriptor.with_env().with_display_handle(Box::new(
+            // TODO: Use event_loop.owned_display_handle() with winit 0.30
+            window.clone(),
+        ));
         let instance = wgpu::Instance::new(instance_descriptor);
         surface.pre_adapter(&instance, window);
         let adapter = get_adapter_with_capabilities_or_from_env(

examples/features/src/mesh_shader/mod.rs

@@ -1,40 +1,14 @@
 // 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()),
-    })
-}
-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();
+    // Workgroup memory zero initialization can be expensive for mesh shaders
     unsafe {
-        device.create_shader_module_passthrough(wgpu::ShaderModuleDescriptorPassthrough {
-            entry_point: entry.to_owned(),
-            label: None,
-            num_workgroups: (1, 1, 1),
-            dxil: Some(std::borrow::Cow::Owned(file)),
-            ..Default::default()
-        })
+        device.create_shader_module_trusted(
+            wgpu::ShaderModuleDescriptor {
+                label: None,
+                source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()),
+            },
+            wgpu::ShaderRuntimeChecks::unchecked(),
+        )
     }
 }
 
@@ -61,22 +35,14 @@ impl crate::framework::Example for Example {
         _queue: &wgpu::Queue,
     ) -> Self {
         let (ts, ms, fs, ts_name, ms_name, fs_name) = match adapter.get_info().backend {
-            wgpu::Backend::Vulkan => (
+            wgpu::Backend::Vulkan | wgpu::Backend::Dx12 => (
                 compile_wgsl(device),
                 compile_wgsl(device),
                 compile_wgsl(device),
                 "ts_main",
                 "ms_main",
                 "fs_main",
             ),
-            wgpu::Backend::Dx12 => (
-                compile_hlsl(device, "Task", "as"),
-                compile_hlsl(device, "Mesh", "ms"),
-                compile_hlsl(device, "Frag", "ps"),
-                "main",
-                "main",
-                "main",
-            ),
             wgpu::Backend::Metal => (
                 compile_msl(device, "taskShader"),
                 compile_msl(device, "meshShader"),

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,27 @@ 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[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[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.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);
+        return;
+    }
 }
 
 @fragment

naga-test/src/lib.rs

@@ -164,6 +164,10 @@ impl SpirvOutParameters {
             ray_query_initialization_tracking: true,
             debug_info,
             use_storage_input_output_16: self.use_storage_input_output_16,
+            task_runtime_limits: Some(naga::back::TaskRuntimeLimits {
+                max_mesh_workgroups_per_dim: 256,
+                max_mesh_workgroups_total: 1024,
+            }),
         }
     }
 }

naga/hlsl-snapshots/src/lib.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/src/back/hlsl/conv.rs

@@ -149,8 +149,8 @@ impl crate::StorageFormat {
 }
 
 impl crate::BuiltIn {
-    pub(super) fn to_hlsl_str(self) -> Result<&'static str, Error> {
-        Ok(match self {
+    pub(super) fn to_hlsl_str(self) -> Result<Option<&'static str>, Error> {
+        Ok(Some(match self {
             Self::Position { .. } => "SV_Position",
             // vertex
             Self::ClipDistance => "SV_ClipDistance",
@@ -186,13 +186,15 @@ 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),
+        }))
     }
 }