[HLSL] Add utility to convert environment to stage

We had a bunch of places in the code where we were translating triple
environment enum cases to shader stage enum cases. The order of these
enums needs to be kept in sync for the translation to be simple, but we
were not properly handling out-of-bounds cases.

In normal compilation out-of-bounds cases shouldn't be possible because
the driver errors if you don't have a valid shader environment set, but
in clang tooling that error doesn't get treated as fatal and parsing
continues. This can result in crashes in clang tooling for out-of-range
shader stages.

To address this, this patch adds a constexpr method to handle the
conversion which handles out-of-range values by converting them to
`Invalid`.

Since this new method is a constexpr, the tests for this are a group of
static_asserts in the implementation file that verifies the correct
conversion for each valid enum case and validates that other cases are
converted to `Invalid`.

Reviewed By: bogner

Differential Revision: https://reviews.llvm.org/D135595

diff --git a/clang/include/clang/Basic/HLSLRuntime.h b/clang/include/clang/Basic/HLSLRuntime.h
index bc3a064..d3d2306 100644 (file)
--- a/clang/include/clang/Basic/HLSLRuntime.h
+++ b/clang/include/clang/Basic/HLSLRuntime.h
@@ -7,14 +7,14 @@
 //===----------------------------------------------------------------------===//
 //
 /// \file
-/// Defines the clang::IdentifierInfo, clang::IdentifierTable, and
-/// clang::Selector interfaces.
+/// Defines helper utilities for supporting the HLSL runtime environment.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef CLANG_BASIC_HLSLRUNTIME_H
 #define CLANG_BASIC_HLSLRUNTIME_H
 
+#include "clang/Basic/LangOptions.h"
 #include <cstdint>
 
 namespace clang {
@@ -28,6 +28,46 @@ enum class ResourceClass : uint8_t {
   NumClasses
 };
 
+constexpr ShaderStage
+getStageFromEnvironment(const llvm::Triple::EnvironmentType &E) {
+  uint32_t Pipeline =
+      static_cast<uint32_t>(E) - static_cast<uint32_t>(llvm::Triple::Pixel);
+
+  if (Pipeline > (uint32_t)ShaderStage::Invalid)
+    return ShaderStage::Invalid;
+  return static_cast<ShaderStage>(Pipeline);
+}
+
+#define ENUM_COMPARE_ASSERT(Value)                                             \
+  static_assert(                                                               \
+      getStageFromEnvironment(llvm::Triple::Value) == ShaderStage::Value,      \
+      "Mismatch between llvm::Triple and clang::ShaderStage for " #Value);
+
+ENUM_COMPARE_ASSERT(Pixel)
+ENUM_COMPARE_ASSERT(Vertex)
+ENUM_COMPARE_ASSERT(Geometry)
+ENUM_COMPARE_ASSERT(Hull)
+ENUM_COMPARE_ASSERT(Domain)
+ENUM_COMPARE_ASSERT(Compute)
+ENUM_COMPARE_ASSERT(Library)
+ENUM_COMPARE_ASSERT(RayGeneration)
+ENUM_COMPARE_ASSERT(Intersection)
+ENUM_COMPARE_ASSERT(AnyHit)
+ENUM_COMPARE_ASSERT(ClosestHit)
+ENUM_COMPARE_ASSERT(Miss)
+ENUM_COMPARE_ASSERT(Callable)
+ENUM_COMPARE_ASSERT(Mesh)
+ENUM_COMPARE_ASSERT(Amplification)
+
+static_assert(getStageFromEnvironment(llvm::Triple::UnknownEnvironment) ==
+                  ShaderStage::Invalid,
+              "Mismatch between llvm::Triple and "
+              "clang::ShaderStage for Invalid");
+static_assert(getStageFromEnvironment(llvm::Triple::MSVC) ==
+                  ShaderStage::Invalid,
+              "Mismatch between llvm::Triple and "
+              "clang::ShaderStage for Invalid");
+
 } // namespace hlsl
 } // namespace clang
 

diff --git a/clang/include/clang/Basic/LangOptions.h b/clang/include/clang/Basic/LangOptions.h
index a9ed2e2..4cac4c2 100644 (file)
--- a/clang/include/clang/Basic/LangOptions.h
+++ b/clang/include/clang/Basic/LangOptions.h
@@ -54,6 +54,9 @@ protected:
 enum class MSVtorDispMode { Never, ForVBaseOverride, ForVFTable };
 
 /// Shader programs run in specific pipeline stages.
+/// The order of these values matters, and must be kept in sync with the
+/// Triple Environment enum in llvm::Triple. The ordering is enforced in
+///  static_asserts in Triple.cpp and in clang/Basic/HLSLRuntime.h.
 enum class ShaderStage {
   Pixel = 0,
   Vertex,

diff --git a/clang/lib/Frontend/InitPreprocessor.cpp b/clang/lib/Frontend/InitPreprocessor.cpp
index e9bfab9..2273fb1 100644 (file)
--- a/clang/lib/Frontend/InitPreprocessor.cpp
+++ b/clang/lib/Frontend/InitPreprocessor.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Basic/FileManager.h"
+#include "clang/Basic/HLSLRuntime.h"
 #include "clang/Basic/MacroBuilder.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/SyncScope.h"
@@ -402,8 +403,8 @@ static void InitializeStandardPredefinedMacros(const TargetInfo &TI,
     Builder.defineMacro("__SHADER_STAGE_LIBRARY",
                         Twine((uint32_t)ShaderStage::Library));
     // The current shader stage itself
-    uint32_t StageInteger = (uint32_t)TI.getTriple().getEnvironment() -
-                            (uint32_t)llvm::Triple::Pixel;
+    uint32_t StageInteger = static_cast<uint32_t>(
+        hlsl::getStageFromEnvironment(TI.getTriple().getEnvironment()));
 
     Builder.defineMacro("__SHADER_TARGET_STAGE", Twine(StageInteger));
     // Add target versions

diff --git a/clang/lib/Sema/SemaDecl.cpp b/clang/lib/Sema/SemaDecl.cpp
index 5cc3db6..33be49c 100644 (file)
--- a/clang/lib/Sema/SemaDecl.cpp
+++ b/clang/lib/Sema/SemaDecl.cpp
@@ -27,6 +27,7 @@
 #include "clang/AST/Randstruct.h"
 #include "clang/AST/StmtCXX.h"
 #include "clang/Basic/Builtins.h"
+#include "clang/Basic/HLSLRuntime.h"
 #include "clang/Basic/PartialDiagnostic.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TargetInfo.h"
@@ -10117,10 +10118,11 @@ Sema::ActOnFunctionDeclarator(Scope *S, Declarator &D, DeclContext *DC,
         S->getDepth() == 0) {
       CheckHLSLEntryPoint(NewFD);
       if (!NewFD->isInvalidDecl()) {
-        auto TripleShaderType = TargetInfo.getTriple().getEnvironment();
+        auto Env = TargetInfo.getTriple().getEnvironment();
         AttributeCommonInfo AL(NewFD->getBeginLoc());
-        HLSLShaderAttr::ShaderType ShaderType = (HLSLShaderAttr::ShaderType)(
-            TripleShaderType - (uint32_t)llvm::Triple::Pixel);
+        HLSLShaderAttr::ShaderType ShaderType =
+            static_cast<HLSLShaderAttr::ShaderType>(
+                hlsl::getStageFromEnvironment(Env));
         // To share code with HLSLShaderAttr, add HLSLShaderAttr to entry
         // function.
         if (HLSLShaderAttr *Attr = mergeHLSLShaderAttr(NewFD, AL, ShaderType))

diff --git a/clang/lib/Sema/SemaDeclAttr.cpp b/clang/lib/Sema/SemaDeclAttr.cpp
index 51d22d2..d13bcf4 100644 (file)
--- a/clang/lib/Sema/SemaDeclAttr.cpp
+++ b/clang/lib/Sema/SemaDeclAttr.cpp
@@ -24,6 +24,7 @@
 #include "clang/AST/Type.h"
 #include "clang/Basic/CharInfo.h"
 #include "clang/Basic/DarwinSDKInfo.h"
+#include "clang/Basic/HLSLRuntime.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
@@ -6819,12 +6820,12 @@ static void handleUuidAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
 static void handleHLSLNumThreadsAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   using llvm::Triple;
   Triple Target = S.Context.getTargetInfo().getTriple();
+  auto Env = S.Context.getTargetInfo().getTriple().getEnvironment();
   if (!llvm::is_contained({Triple::Compute, Triple::Mesh, Triple::Amplification,
                            Triple::Library},
-                          Target.getEnvironment())) {
+                          Env)) {
     uint32_t Pipeline =
-        (uint32_t)S.Context.getTargetInfo().getTriple().getEnvironment() -
-        (uint32_t)llvm::Triple::Pixel;
+        static_cast<uint32_t>(hlsl::getStageFromEnvironment(Env));
     S.Diag(AL.getLoc(), diag::err_hlsl_attr_unsupported_in_stage)
         << AL << Pipeline << "Compute, Amplification, Mesh or Library";
     return;
@@ -6891,16 +6892,13 @@ HLSLNumThreadsAttr *Sema::mergeHLSLNumThreadsAttr(Decl *D,
 
 static void handleHLSLSVGroupIndexAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   using llvm::Triple;
-  Triple Target = S.Context.getTargetInfo().getTriple();
-  if (Target.getEnvironment() != Triple::Compute &&
-      Target.getEnvironment() != Triple::Library) {
+  auto Env = S.Context.getTargetInfo().getTriple().getEnvironment();
+  if (Env != Triple::Compute && Env != Triple::Library) {
     // FIXME: it is OK for a compute shader entry and pixel shader entry live in
     // same HLSL file. Issue https://github.com/llvm/llvm-project/issues/57880.
-    uint32_t Pipeline =
-        (uint32_t)S.Context.getTargetInfo().getTriple().getEnvironment() -
-        (uint32_t)llvm::Triple::Pixel;
+    ShaderStage Pipeline = hlsl::getStageFromEnvironment(Env);
     S.Diag(AL.getLoc(), diag::err_hlsl_attr_unsupported_in_stage)
-        << AL << Pipeline << "Compute";
+        << AL << (uint32_t)Pipeline << "Compute";
     return;
   }
 

diff --git a/llvm/include/llvm/ADT/Triple.h b/llvm/include/llvm/ADT/Triple.h
index 6de79c8..dc96398 100644 (file)
--- a/llvm/include/llvm/ADT/Triple.h
+++ b/llvm/include/llvm/ADT/Triple.h
@@ -249,6 +249,9 @@ public:
     MacABI, // Mac Catalyst variant of Apple's iOS deployment target.
     
     // Shader Stages
+    // The order of these values matters, and must be kept in sync with the
+    // language options enum in Clang. The ordering is enforced in
+    // static_asserts in Triple.cpp and in Clang.
     Pixel,
     Vertex,
     Geometry,