[HLSL][RootSignature] Add parsing of ShaderVisibility to DescriptorTable

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -85,9 +85,13 @@ class RootSignatureParser {
   std::optional<ParsedClauseParams>
   parseDescriptorTableClauseParams(RootSignatureToken::Kind RegType);
 
+  // Common parsing methods
   std::optional<uint32_t> parseUIntParam();
   std::optional<llvm::hlsl::rootsig::Register> parseRegister();
 
+  /// Parsing methods of various enums
+  std::optional<llvm::hlsl::rootsig::ShaderVisibility> parseShaderVisibility();
+
   /// Use NumericLiteralParser to convert CurToken.NumSpelling into a unsigned
   /// 32-bit integer
   std::optional<uint32_t> handleUIntLiteral();

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -52,6 +52,7 @@ std::optional<DescriptorTable> RootSignatureParser::parseDescriptorTable() {
     return std::nullopt;
 
   DescriptorTable Table;
+  std::optional<ShaderVisibility> Visibility;
 
   // Iterate as many Clauses as possible
   do {
@@ -63,8 +64,27 @@ std::optional<DescriptorTable> RootSignatureParser::parseDescriptorTable() {
       Elements.push_back(*Clause);
       Table.NumClauses++;
     }
+
+    if (tryConsumeExpectedToken(TokenKind::kw_visibility)) {
+      if (Visibility.has_value()) {
+        getDiags().Report(CurToken.TokLoc, diag::err_hlsl_rootsig_repeat_param)
+            << CurToken.TokKind;
+        return std::nullopt;
+      }
+
+      if (consumeExpectedToken(TokenKind::pu_equal))
+        return std::nullopt;
+
+      Visibility = parseShaderVisibility();
+      if (!Visibility.has_value())
+        return std::nullopt;
+    }
   } while (tryConsumeExpectedToken(TokenKind::pu_comma));
 
+  // Fill in optional visibility
+  if (Visibility.has_value())
+    Table.Visibility = Visibility.value();
+
   if (consumeExpectedToken(TokenKind::pu_r_paren,
                            diag::err_hlsl_unexpected_end_of_params,
                            /*param of=*/TokenKind::kw_DescriptorTable))
@@ -222,6 +242,32 @@ std::optional<Register> RootSignatureParser::parseRegister() {
   return Reg;
 }
 
+std::optional<llvm::hlsl::rootsig::ShaderVisibility>
+RootSignatureParser::parseShaderVisibility() {
+  assert(CurToken.TokKind == TokenKind::pu_equal &&
+         "Expects to only be invoked starting at given keyword");
+
+  TokenKind Expected[] = {
+#define SHADER_VISIBILITY_ENUM(NAME, LIT) TokenKind::en_##NAME,
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+  };
+
+  if (!tryConsumeExpectedToken(Expected))
+    return std::nullopt;
+
+  switch (CurToken.TokKind) {
+#define SHADER_VISIBILITY_ENUM(NAME, LIT)                                      \
+  case TokenKind::en_##NAME:                                                   \
+    return ShaderVisibility::NAME;                                             \
+    break;
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+  default:
+    llvm_unreachable("Switch for consumed enum token was not provided");
+  }
+
+  return std::nullopt;
+}
+
 std::optional<uint32_t> RootSignatureParser::handleUIntLiteral() {
   // Parse the numeric value and do semantic checks on its specification
   clang::NumericLiteralParser Literal(CurToken.NumSpelling, CurToken.TokLoc,

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -131,6 +131,7 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
     DescriptorTable(
       CBV(b0),
       SRV(space = 3, t42),
+      visibility = SHADER_VISIBILITY_PIXEL,
       Sampler(s987, space = +2),
       UAV(u4294967294)
     ),
@@ -186,11 +187,14 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   Elem = Elements[4];
   ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));
   ASSERT_EQ(std::get<DescriptorTable>(Elem).NumClauses, (uint32_t)4);
+  ASSERT_EQ(std::get<DescriptorTable>(Elem).Visibility,
+            ShaderVisibility::Pixel);
 
   // Empty Descriptor Table
   Elem = Elements[5];
   ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));
   ASSERT_EQ(std::get<DescriptorTable>(Elem).NumClauses, 0u);
+  ASSERT_EQ(std::get<DescriptorTable>(Elem).Visibility, ShaderVisibility::All);
 
   ASSERT_TRUE(Consumer->isSatisfied());
 }

llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h

@@ -21,6 +21,19 @@ namespace llvm {
 namespace hlsl {
 namespace rootsig {
 
+// Definition of the various enumerations and flags
+
+enum class ShaderVisibility {
+  All = 0,
+  Vertex = 1,
+  Hull = 2,
+  Domain = 3,
+  Geometry = 4,
+  Pixel = 5,
+  Amplification = 6,
+  Mesh = 7,
+};
+
 // Definitions of the in-memory data layout structures
 
 // Models the different registers: bReg | tReg | uReg | sReg
@@ -32,6 +45,7 @@ struct Register {
 
 // Models the end of a descriptor table and stores its visibility
 struct DescriptorTable {
+  ShaderVisibility Visibility = ShaderVisibility::All;
   uint32_t NumClauses = 0; // The number of clauses in the table
 };