[HLSL] Add utility to convert environment to stage

[beanz]

Summary

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.

[python3kgae]

Backend needs the same thing.
Is it possible to move this to llvm and share it between frontend and backend?

[beanz]

In D135595#3847056, @python3kgae wrote:

Backend needs the same thing.
Is it possible to move this to llvm and share it between frontend and backend?

This translates the triple to a clang-defined enum... so strictly speaking no, that isn't possible.

Can you elaborate on the context where the backend needs this? I suspect in the backend we need a similar but not quite the same utility. Specifically, I think we need to translate the environment value to an unsigned integer. For all cases where an enum makes sense we should be using the Environment value in the backend.

[beanz]

Worth noting, we do have a similar set of static_asserts in Triple.cpp to validate the ordering of enum cases and that the subtraction results in the appropriate values:

https://github.com/llvm/llvm-project/blob/main/llvm/lib/Support/Triple.cpp#L1942

[python3kgae]

In D135595#3847158, @beanz wrote:

In D135595#3847056, @python3kgae wrote:

Backend needs the same thing.
Is it possible to move this to llvm and share it between frontend and backend?

This translates the triple to a clang-defined enum... so strictly speaking no, that isn't possible.

Can you elaborate on the context where the backend needs this? I suspect in the backend we need a similar but not quite the same utility. Specifically, I think we need to translate the environment value to an unsigned integer. For all cases where an enum makes sense we should be using the Environment value in the backend.

I think most use cases could be covered by switch and make default go llvm_unreachable.
It would be helpful to have util functions like isGraphics (which return true for graphics shader like vs/ps), isRay( which return true on ray tracing entry functions).

[beanz]

We should absolutely add utility functions when we have uses for them.

[bogner]

Even without the tooling aspect the addition of the static asserts is a strict improvement. looks good, thanks!

clang/include/clang/Basic/HLSLRuntime.h

39

You're not actually introducing the dependency here (it was already there), but neither ShaderStage in LangOptions.h nor the shader stage EnvironmentType in Triple.h mention that they need to be kept in sync with the other. Can you add comments to both that note the relationship?

[beanz]

clang/include/clang/Basic/HLSLRuntime.h

39

Will do!

clang/include/clang/Basic/HLSLRuntime.h

 //===- HLSLRuntime.h - HLSL Runtime -----------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 /// \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 {
 namespace hlsl {
 
 enum class ResourceClass : uint8_t {
   SRV = 0,
   UAV,
   CBuffer,
   Sampler,
   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
 
 #endif // CLANG_BASIC_HLSLRUNTIME_H

clang/include/clang/Basic/LangOptions.h

 #include "clang/Basic/LangOptions.def"
 };
 
 /// In the Microsoft ABI, this controls the placement of virtual displacement
 /// members used to implement virtual inheritance.
 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,
   Geometry,
   Hull,
   Domain,
   Compute,
   Library,

clang/lib/Frontend/InitPreprocessor.cpp

 //===--- InitPreprocessor.cpp - PP initialization code. ---------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the clang::InitializePreprocessor function.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Basic/Version.h"
 #include "clang/Frontend/FrontendDiagnostic.h"
 #include "clang/Frontend/FrontendOptions.h"
 #include "clang/Frontend/Utils.h"
Builder.defineMacro("__SHADER_STAGE_COMPUTE",
                         Twine((uint32_t)ShaderStage::Compute));
     Builder.defineMacro("__SHADER_STAGE_AMPLIFICATION",
                         Twine((uint32_t)ShaderStage::Amplification));
     Builder.defineMacro("__SHADER_STAGE_MESH",
                         Twine((uint32_t)ShaderStage::Mesh));
     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
     if (TI.getTriple().getOS() == llvm::Triple::ShaderModel) {
       VersionTuple Version = TI.getTriple().getOSVersion();
       Builder.defineMacro("__SHADER_TARGET_MAJOR", Twine(Version.getMajor()));
       unsigned Minor = Version.getMinor().value_or(0);
       Builder.defineMacro("__SHADER_TARGET_MINOR", Twine(Minor));

clang/lib/Sema/SemaDecl.cpp

 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/EvaluatedExprVisitor.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/NonTrivialTypeVisitor.h"
 #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"
 #include "clang/Lex/HeaderSearch.h" // TODO: Sema shouldn't depend on Lex
 #include "clang/Lex/Lexer.h" // TODO: Extract static functions to fix layering.
 #include "clang/Lex/ModuleLoader.h" // TODO: Sema shouldn't depend on Lex
 #include "clang/Lex/Preprocessor.h" // Included for isCodeCompletionEnabled()
 #include "clang/Sema/CXXFieldCollector.h"
if (getLangOpts().HLSL) {
     auto &TargetInfo = getASTContext().getTargetInfo();
     // Skip operator overload which not identifier.
     // Also make sure NewFD is in translation-unit scope.
     if (!NewFD->isInvalidDecl() && Name.isIdentifier() &&
         NewFD->getName() == TargetInfo.getTargetOpts().HLSLEntry &&
         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))
           NewFD->addAttr(Attr);
       }
     }
   }
 

clang/lib/Sema/SemaDeclAttr.cpp

 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/Mangle.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #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"
 #include "clang/Basic/TargetBuiltins.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Sema/DeclSpec.h"
 #include "clang/Sema/DelayedDiagnostic.h"
static void handleUuidAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   UuidAttr *UA = S.mergeUuidAttr(D, AL, OrigStrRef, Guid);
   if (UA)
     D->addAttr(UA);
 }
 
 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;
   }
 
   llvm::VersionTuple SMVersion = Target.getOSVersion();
   uint32_t ZMax = 1024;
   uint32_t ThreadMax = 1024;
if (HLSLNumThreadsAttr *NT = D->getAttr<HLSLNumThreadsAttr>()) {
     }
     return nullptr;
   }
   return ::new (Context) HLSLNumThreadsAttr(Context, AL, X, Y, Z);
 }
 
 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;
   }
 
   D->addAttr(::new (S.Context) HLSLSV_GroupIndexAttr(S.Context, AL));
 }
 
 static void handleHLSLShaderAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   StringRef Str;

llvm/include/llvm/ADT/Triple.h

enum EnvironmentType {
     MSVC,
     Itanium,
     Cygnus,
     CoreCLR,
     Simulator, // Simulator variants of other systems, e.g., Apple's iOS
     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,
     Hull,
     Domain,
     Compute,
     Library,
     RayGeneration,