mesh shader ext: fix task shader payload passing

Author: Firestar99

crates/spirv-std/src/arch/mesh_shading.rs

@@ -64,3 +64,46 @@ pub unsafe fn emit_mesh_tasks_ext(group_count_x: u32, group_count_y: u32, group_
         options(noreturn),
     }
 }
+
+/// Defines the grid size of subsequent mesh shader workgroups to generate
+/// upon completion of the task shader workgroup.
+///
+/// 'Group Count X Y Z' must each be a 32-bit unsigned integer value.
+/// They configure the number of local workgroups in each respective dimensions
+/// for the launch of child mesh tasks. See Vulkan API specification for more detail.
+///
+/// 'Payload' is an optional pointer to the payload structure to pass to the generated mesh shader invocations.
+/// 'Payload' must be the result of an *OpVariable* with a storage class of *TaskPayloadWorkgroupEXT*.
+///
+/// The arguments are taken from the first invocation in each workgroup.
+/// Any invocation must execute this instruction exactly once and under uniform
+/// control flow.
+/// This instruction also serves as an *OpControlBarrier* instruction, and also
+/// performs and adheres to the description and semantics of an *OpControlBarrier*
+/// instruction with the 'Execution' and 'Memory' operands set to *Workgroup* and
+/// the 'Semantics' operand set to a combination of *WorkgroupMemory* and
+/// *AcquireRelease*.
+/// Ceases all further processing: Only instructions executed before
+/// *OpEmitMeshTasksEXT* have observable side effects.
+///
+/// This instruction must be the last instruction in a block.
+///
+/// This instruction is only valid in the *TaskEXT* Execution Model.
+#[spirv_std_macros::gpu_only]
+#[doc(alias = "OpEmitMeshTasksEXT")]
+#[inline]
+pub unsafe fn emit_mesh_tasks_ext_payload<T>(
+    group_count_x: u32,
+    group_count_y: u32,
+    group_count_z: u32,
+    payload: &mut T,
+) -> ! {
+    asm! {
+        "OpEmitMeshTasksEXT {group_count_x} {group_count_y} {group_count_z} {payload}",
+        group_count_x = in(reg) group_count_x,
+        group_count_y = in(reg) group_count_y,
+        group_count_z = in(reg) group_count_z,
+        payload = in(reg) payload,
+        options(noreturn),
+    }
+}

tests/ui/arch/task_shader_payload.rs

@@ -2,7 +2,7 @@
 // only-vulkan1.2
 // compile-flags: -Ctarget-feature=+MeshShadingEXT,+ext:SPV_EXT_mesh_shader
 
-use spirv_std::arch::emit_mesh_tasks_ext;
+use spirv_std::arch::emit_mesh_tasks_ext_payload;
 use spirv_std::spirv;
 
 pub struct Payload {
@@ -16,6 +16,6 @@ pub fn main(#[spirv(task_payload_workgroup_ext)] payload: &mut Payload) {
     payload.second = 2;
 
     unsafe {
-        emit_mesh_tasks_ext(3, 4, 5);
+        emit_mesh_tasks_ext_payload(3, 4, 5, payload);
     }
 }