From d7df74f3832de10347f7a8c4b33a7d6ad8c4b480 Mon Sep 17 00:00:00 2001
From: Firestar99 <31222740+Firestar99@users.noreply.github.com>
Date: Mon, 28 Oct 2024 12:57:07 +0100
Subject: [PATCH] mesh shaders: fix `per_primitive_ext` validation for fragment
 shaders

---
 .../src/codegen_cx/entry.rs                   | 37 ++++++++++++-------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/entry.rs b/crates/rustc_codegen_spirv/src/codegen_cx/entry.rs
index a42328024c..97a2b441c9 100644
--- a/crates/rustc_codegen_spirv/src/codegen_cx/entry.rs
+++ b/crates/rustc_codegen_spirv/src/codegen_cx/entry.rs
@@ -739,20 +739,31 @@ impl<'tcx> CodegenCx<'tcx> {
                 .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",
-                );
+            match execution_model {
+                ExecutionModel::Fragment => {
+                    if storage_class != Ok(StorageClass::Input) {
+                        self.tcx.dcx().span_fatal(
+                            per_primitive_ext.span,
+                            "`#[spirv(per_primitive_ext)]` in fragment shaders is only valid on Input variables",
+                        );
+                    }
+                }
+                ExecutionModel::MeshNV | ExecutionModel::MeshEXT => {
+                    if storage_class != Ok(StorageClass::Output) {
+                        self.tcx.dcx().span_fatal(
+                            per_primitive_ext.span,
+                            "`#[spirv(per_primitive_ext)]` in mesh shaders is only valid on Output variables",
+                        );
+                    }
+                }
+                _ => {
+                    self.tcx.dcx().span_fatal(
+                        per_primitive_ext.span,
+                        "`#[spirv(per_primitive_ext)]` is only valid in fragment or mesh shaders",
+                    );
+                }
             }
+
             self.emit_global().decorate(
                 var_id.unwrap(),
                 Decoration::PerPrimitiveEXT,