[HLSL] Further improve to numthreads diagnostics

This adds diagnostics for conflicting attributes on the same
declarataion, conflicting attributes on a forward and final
declaration, and defines a more narrowly scoped HLSLEntry attribute
target.

Big shout out to @aaron.ballman for the great feedback and review on
this!

6 files changed, 84 insertions, 14 deletions

diff --git a/clang/include/clang/Basic/Attr.td b/clang/include/clang/Basic/Attr.td
index 97b1027..4789493 100644
--- a/clang/include/clang/Basic/Attr.td
+++ b/clang/include/clang/Basic/Attr.td
@@ -126,8 +126,10 @@ def FunctionTmpl
                                  FunctionDecl::TK_FunctionTemplate}],
                     "function templates">;
 
-def GlobalFunction
-    : SubsetSubject<Function, [{S->isGlobal()}], "global functions">;
+def HLSLEntry
+    : SubsetSubject<Function,
+                    [{S->isExternallyVisible() && !isa<CXXMethodDecl>(S)}],
+                    "global functions">;
 
 def ClassTmpl : SubsetSubject<CXXRecord, [{S->getDescribedClassTemplate()}],
                               "class templates">;
@@ -3946,7 +3948,7 @@ def Error : InheritableAttr {
 def HLSLNumThreads: InheritableAttr {
   let Spellings = [Microsoft<"numthreads">];
   let Args = [IntArgument<"X">, IntArgument<"Y">, IntArgument<"Z">];
-  let Subjects = SubjectList<[GlobalFunction]>;
+  let Subjects = SubjectList<[HLSLEntry]>;
   let LangOpts = [HLSL];
   let Documentation = [NumThreadsDocs];
 }
diff --git a/clang/include/clang/Basic/DiagnosticSemaKinds.td b/clang/include/clang/Basic/DiagnosticSemaKinds.td
index a272cb7..aec172c 100644
--- a/clang/include/clang/Basic/DiagnosticSemaKinds.td
+++ b/clang/include/clang/Basic/DiagnosticSemaKinds.td
@@ -11570,6 +11570,7 @@ def err_hlsl_attr_unsupported_in_stage : Error<"attribute %0 is unsupported in %
 
 def err_hlsl_numthreads_argument_oor : Error<"argument '%select{X|Y|Z}0' to numthreads attribute cannot exceed %1">;
 def err_hlsl_numthreads_invalid : Error<"total number of threads cannot exceed %0">;
+def err_hlsl_attribute_param_mismatch : Error<"%0 attribute parameters do not match the previous declaration">;
 
 } // end of sema component.
 
diff --git a/clang/include/clang/Sema/Sema.h b/clang/include/clang/Sema/Sema.h
index 6523c30..c0ad55d 100644
--- a/clang/include/clang/Sema/Sema.h
+++ b/clang/include/clang/Sema/Sema.h
@@ -3471,6 +3471,9 @@ public:
   EnforceTCBLeafAttr *mergeEnforceTCBLeafAttr(Decl *D,
                                               const EnforceTCBLeafAttr &AL);
   BTFDeclTagAttr *mergeBTFDeclTagAttr(Decl *D, const BTFDeclTagAttr &AL);
+  HLSLNumThreadsAttr *mergeHLSLNumThreadsAttr(Decl *D,
+                                              const AttributeCommonInfo &AL,
+                                              int X, int Y, int Z);
 
   void mergeDeclAttributes(NamedDecl *New, Decl *Old,
                            AvailabilityMergeKind AMK = AMK_Redeclaration);
diff --git a/clang/lib/Sema/SemaDecl.cpp b/clang/lib/Sema/SemaDecl.cpp
index b913f80..1e25346 100644
--- a/clang/lib/Sema/SemaDecl.cpp
+++ b/clang/lib/Sema/SemaDecl.cpp
@@ -2770,6 +2770,9 @@ static bool mergeDeclAttribute(Sema &S, NamedDecl *D,
     NewAttr = S.mergeEnforceTCBLeafAttr(D, *TCBLA);
   else if (const auto *BTFA = dyn_cast<BTFDeclTagAttr>(Attr))
     NewAttr = S.mergeBTFDeclTagAttr(D, *BTFA);
+  else if (const auto *NT = dyn_cast<HLSLNumThreadsAttr>(Attr))
+    NewAttr =
+        S.mergeHLSLNumThreadsAttr(D, *NT, NT->getX(), NT->getY(), NT->getZ());
   else if (Attr->shouldInheritEvenIfAlreadyPresent() || !DeclHasAttr(D, Attr))
     NewAttr = cast<InheritableAttr>(Attr->clone(S.Context));
 
diff --git a/clang/lib/Sema/SemaDeclAttr.cpp b/clang/lib/Sema/SemaDeclAttr.cpp
index 87e1663..4b5201d 100644
--- a/clang/lib/Sema/SemaDeclAttr.cpp
+++ b/clang/lib/Sema/SemaDeclAttr.cpp
@@ -6892,7 +6892,22 @@ static void handleHLSLNumThreadsAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
     return;
   }
 
-  D->addAttr(::new (S.Context) HLSLNumThreadsAttr(S.Context, AL, X, Y, Z));
+  HLSLNumThreadsAttr *NewAttr = S.mergeHLSLNumThreadsAttr(D, AL, X, Y, Z);
+  if (NewAttr)
+    D->addAttr(NewAttr);
+}
+
+HLSLNumThreadsAttr *Sema::mergeHLSLNumThreadsAttr(Decl *D,
+                                                  const AttributeCommonInfo &AL,
+                                                  int X, int Y, int Z) {
+  if (HLSLNumThreadsAttr *NT = D->getAttr<HLSLNumThreadsAttr>()) {
+    if (NT->getX() != X || NT->getY() != Y || NT->getZ() != Z) {
+      Diag(NT->getLocation(), diag::err_hlsl_attribute_param_mismatch) << AL;
+      Diag(AL.getLoc(), diag::note_conflicting_attribute);
+    }
+    return nullptr;
+  }
+  return ::new (Context) HLSLNumThreadsAttr(Context, AL, X, Y, Z);
 }
 
 static void handleMSInheritanceAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
diff --git a/clang/test/SemaHLSL/num_threads.hlsl b/clang/test/SemaHLSL/num_threads.hlsl
index cf9e2480..f93e67d 100644
--- a/clang/test/SemaHLSL/num_threads.hlsl
+++ b/clang/test/SemaHLSL/num_threads.hlsl
@@ -12,6 +12,47 @@
 
 #if __SHADER_TARGET_STAGE == __SHADER_STAGE_COMPUTE || __SHADER_TARGET_STAGE == __SHADER_STAGE_MESH || __SHADER_TARGET_STAGE == __SHADER_STAGE_AMPLIFICATION || __SHADER_TARGET_STAGE == __SHADER_STAGE_LIBRARY
 #ifdef FAIL
+
+// expected-warning@+1 {{'numthreads' attribute only applies to global functions}}
+[numthreads(1,1,1)]
+struct Fido {
+  // expected-warning@+1 {{'numthreads' attribute only applies to global functions}}
+  [numthreads(1,1,1)]
+  void wag() {}
+
+  // expected-warning@+1 {{'numthreads' attribute only applies to global functions}}
+  [numthreads(1,1,1)]
+  static void oops() {}
+};
+
+// expected-warning@+1 {{'numthreads' attribute only applies to global functions}}
+[numthreads(1,1,1)]
+static void oops() {}
+
+namespace spec {
+// expected-warning@+1 {{'numthreads' attribute only applies to global functions}}
+[numthreads(1,1,1)]
+static void oops() {}
+}
+
+// expected-error@+1 {{'numthreads' attribute parameters do not match the previous declaration}}
+[numthreads(1,1,1)]
+// expected-note@+1 {{conflicting attribute is here}}
+[numthreads(2,2,1)]
+int doubledUp() {
+  return 1;
+}
+
+// expected-note@+1 {{conflicting attribute is here}}
+[numthreads(1,1,1)]
+int forwardDecl();
+
+// expected-error@+1 {{'numthreads' attribute parameters do not match the previous declaration}}
+[numthreads(2,2,1)]
+int forwardDecl() {
+  return 1;
+}
+
 #if __SHADER_TARGET_MAJOR == 6
 // expected-error@+1 {{'numthreads' attribute requires exactly 3 arguments}}
 [numthreads]
@@ -37,28 +78,33 @@
 [numthreads(1,2,2)]
 // expected-error@+1 {{total number of threads cannot exceed 768}}
 [numthreads(1024,1,1)]
-#endif
-#endif
+#endif // __SHADER_TARGET_MAJOR
+#endif // FAIL
 // CHECK: HLSLNumThreadsAttr 0x{{[0-9a-fA-F]+}} <line:{{[0-9]+}}:2, col:18> 1 2 1
 [numthreads(1,2,1)]
 int entry() {
  return 1;
 }
 
-// expected-warning@+1 {{'numthreads' attribute only applies to global functions}}
-[numthreads(1,1,1)]
-struct Fido {
-  // expected-warning@+1 {{'numthreads' attribute only applies to global functions}}
-  [numthreads(1,1,1)]
-  void wag() {}
-};
+// Because these two attributes match, they should both appear in the AST
+[numthreads(2,2,1)]
+// CHECK: HLSLNumThreadsAttr 0x{{[0-9a-fA-F]+}} <line:90:2, col:18> 2 2 1
+int secondFn();
 
-#else
+[numthreads(2,2,1)]
+// CHECK: HLSLNumThreadsAttr 0x{{[0-9a-fA-F]+}} <line:94:2, col:18> 2 2 1
+int secondFn() {
+  return 1;
+}
+
+
+#else // Vertex and Pixel only beyond here
 // expected-error-re@+1 {{attribute 'numthreads' is unsupported in {{[A-Za-z]+}} shaders, requires Compute, Amplification, Mesh or Library}}
 [numthreads(1,1,1)]
 int main() {
  return 1;
 }
+
 #endif