mesh shaders: added per_primitive_ext output attribute

Author: Firestar99

crates/rustc_codegen_spirv/src/attr.rs

@@ -92,6 +92,7 @@ pub enum SpirvAttribute {
     DescriptorSet(u32),
     Binding(u32),
     Flat,
+    PerPrimitiveExt,
     Invariant,
     InputAttachmentIndex(u32),
     SpecConstant(SpecConstant),
@@ -128,6 +129,7 @@ pub struct AggregatedSpirvAttributes {
     pub binding: Option<Spanned<u32>>,
     pub flat: Option<Spanned<()>>,
     pub invariant: Option<Spanned<()>>,
+    pub per_primitive_ext: Option<Spanned<()>>,
     pub input_attachment_index: Option<Spanned<u32>>,
     pub spec_constant: Option<Spanned<SpecConstant>>,
 
@@ -213,6 +215,12 @@ impl AggregatedSpirvAttributes {
             Binding(value) => try_insert(&mut self.binding, value, span, "#[spirv(binding)]"),
             Flat => try_insert(&mut self.flat, (), span, "#[spirv(flat)]"),
             Invariant => try_insert(&mut self.invariant, (), span, "#[spirv(invariant)]"),
+            PerPrimitiveExt => try_insert(
+                &mut self.per_primitive_ext,
+                (),
+                span,
+                "#[spirv(per_primitive_ext)]",
+            ),
             InputAttachmentIndex(value) => try_insert(
                 &mut self.input_attachment_index,
                 value,
@@ -314,6 +322,7 @@ impl CheckSpirvAttrVisitor<'_> {
                 | SpirvAttribute::Binding(_)
                 | SpirvAttribute::Flat
                 | SpirvAttribute::Invariant
+                | SpirvAttribute::PerPrimitiveExt
                 | SpirvAttribute::InputAttachmentIndex(_)
                 | SpirvAttribute::SpecConstant(_) => match target {
                     Target::Param => {

crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

@@ -732,6 +732,27 @@ impl<'tcx> CodegenCx<'tcx> {
             self.emit_global()
                 .decorate(var_id.unwrap(), Decoration::Invariant, std::iter::empty());
         }
+        if let Some(per_primitive_ext) = attrs.per_primitive_ext {
+            if storage_class != Ok(StorageClass::Output) {
+                self.tcx.dcx().span_fatal(
+                    per_primitive_ext.span,
+                    "`#[spirv(per_primitive_ext)]` is only valid on Output variables",
+                );
+            }
+            if !(execution_model == ExecutionModel::MeshEXT
+                || execution_model == ExecutionModel::MeshNV)
+            {
+                self.tcx.dcx().span_fatal(
+                    per_primitive_ext.span,
+                    "`#[spirv(per_primitive_ext)]` is only valid in mesh shaders",
+                );
+            }
+            self.emit_global().decorate(
+                var_id.unwrap(),
+                Decoration::PerPrimitiveEXT,
+                std::iter::empty(),
+            );
+        }
 
         let is_subpass_input = match self.lookup_type(value_spirv_type) {
             SpirvType::Image {

crates/rustc_codegen_spirv/src/symbols.rs

@@ -355,6 +355,7 @@ impl Symbols {
             ("block", SpirvAttribute::Block),
             ("flat", SpirvAttribute::Flat),
             ("invariant", SpirvAttribute::Invariant),
+            ("per_primitive_ext", SpirvAttribute::PerPrimitiveExt),
             (
                 "sampled_image",
                 SpirvAttribute::IntrinsicType(IntrinsicType::SampledImage),

tests/ui/arch/mesh_shader_per_primitive.rs

@@ -0,0 +1,34 @@
+// build-pass
+// only-vulkan1.2
+// compile-flags: -Ctarget-feature=+MeshShadingEXT,+ext:SPV_EXT_mesh_shader
+
+use spirv_std::arch::set_mesh_outputs_ext;
+use spirv_std::glam::{UVec3, Vec4};
+use spirv_std::spirv;
+
+#[spirv(mesh_ext(
+    threads(1),
+    output_vertices = 3,
+    output_primitives_ext = 1,
+    output_triangles_ext
+))]
+pub fn main(
+    #[spirv(position)] positions: &mut [Vec4; 3],
+    out_per_vertex: &mut [u32; 3],
+    #[spirv(per_primitive_ext)] out_per_primitive: &mut [u32; 1],
+    #[spirv(primitive_triangle_indices_ext)] indices: &mut [UVec3; 1],
+) {
+    unsafe {
+        set_mesh_outputs_ext(3, 1);
+    }
+
+    positions[0] = Vec4::new(-0.5, 0.5, 0.0, 1.0);
+    positions[1] = Vec4::new(0.5, 0.5, 0.0, 1.0);
+    positions[2] = Vec4::new(0.0, -0.5, 0.0, 1.0);
+    out_per_vertex[0] = 0;
+    out_per_vertex[1] = 1;
+    out_per_vertex[2] = 2;
+
+    indices[0] = UVec3::new(0, 1, 2);
+    out_per_primitive[0] = 42;
+}