[HLSL][RootSignature] Add parsing of optional parameters for RootDescriptor

[inbelic]

• define in-memory representation of optional non-flag parameters to RootDescriptor
• fill in data to parse these params in parseRootDescriptorParams
• add unit tests

Part 3 of #126577

[llvmbot]

@llvm/pr-subscribers-hlsl

@llvm/pr-subscribers-clang

Author: Finn Plummer (inbelic)

Changes

• define in-memory representation of optional non-flag parameters to RootParam
• fill in data to parse these params in parseRootParamParams
• add unit tests

Part 3 of #126577

---

Full diff: https://github.com/llvm/llvm-project/pull/140151.diff

4 Files Affected:

• (modified) clang/include/clang/Parse/ParseHLSLRootSignature.h (+2)
• (modified) clang/lib/Parse/ParseHLSLRootSignature.cpp (+40)
• (modified) clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp (+8-2)
• (modified) llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h (+2)

diff --git a/clang/include/clang/Parse/ParseHLSLRootSignature.h b/clang/include/clang/Parse/ParseHLSLRootSignature.h
index 92b8e7389426a..436d217cec5b1 100644
--- a/clang/include/clang/Parse/ParseHLSLRootSignature.h
+++ b/clang/include/clang/Parse/ParseHLSLRootSignature.h
@@ -91,6 +91,8 @@ class RootSignatureParser {
 
   struct ParsedRootParamParams {
     std::optional<llvm::hlsl::rootsig::Register> Reg;
+    std::optional<uint32_t> Space;
+    std::optional<llvm::hlsl::rootsig::ShaderVisibility> Visibility;
   };
   std::optional<ParsedRootParamParams>
   parseRootParamParams(RootSignatureToken::Kind RegType);
diff --git a/clang/lib/Parse/ParseHLSLRootSignature.cpp b/clang/lib/Parse/ParseHLSLRootSignature.cpp
index de7c2452b3fd4..edb61f29f10d7 100644
--- a/clang/lib/Parse/ParseHLSLRootSignature.cpp
+++ b/clang/lib/Parse/ParseHLSLRootSignature.cpp
@@ -207,6 +207,13 @@ std::optional<RootParam> RootSignatureParser::parseRootParam() {
 
   Param.Reg = Params->Reg.value();
 
+  // Fill in optional values
+  if (Params->Space.has_value())
+    Param.Space = Params->Space.value();
+
+  if (Params->Visibility.has_value())
+    Param.Visibility = Params->Visibility.value();
+
   if (consumeExpectedToken(TokenKind::pu_r_paren,
                            diag::err_hlsl_unexpected_end_of_params,
                            /*param of=*/TokenKind::kw_RootConstants))
@@ -435,6 +442,39 @@ RootSignatureParser::parseRootParamParams(TokenKind RegType) {
       Params.Reg = Reg;
     }
 
+    // `space` `=` POS_INT
+    if (tryConsumeExpectedToken(TokenKind::kw_space)) {
+      if (Params.Space.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;
+
+      auto Space = parseUIntParam();
+      if (!Space.has_value())
+        return std::nullopt;
+      Params.Space = Space;
+    }
+
+    // `visibility` `=` SHADER_VISIBILITY
+    if (tryConsumeExpectedToken(TokenKind::kw_visibility)) {
+      if (Params.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;
+
+      auto Visibility = parseShaderVisibility();
+      if (!Visibility.has_value())
+        return std::nullopt;
+      Params.Visibility = Visibility;
+    }
   } while (tryConsumeExpectedToken(TokenKind::pu_comma));
 
   return Params;
diff --git a/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp b/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
index 0f32523493a53..02bf38dcb110f 100644
--- a/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
+++ b/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
@@ -347,8 +347,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseRootFlagsTest) {
 TEST_F(ParseHLSLRootSignatureTest, ValidParseRootParamsTest) {
   const llvm::StringLiteral Source = R"cc(
     CBV(b0),
-    SRV(t42),
-    UAV(u34893247)
+    SRV(space = 4, t42, visibility = SHADER_VISIBILITY_GEOMETRY),
+    UAV(visibility = SHADER_VISIBILITY_HULL, u34893247)
   )cc";
 
   TrivialModuleLoader ModLoader;
@@ -370,18 +370,24 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseRootParamsTest) {
   ASSERT_TRUE(std::holds_alternative<RootParam>(Elem));
   ASSERT_EQ(std::get<RootParam>(Elem).Reg.ViewType, RegisterType::BReg);
   ASSERT_EQ(std::get<RootParam>(Elem).Reg.Number, 0u);
+  ASSERT_EQ(std::get<RootParam>(Elem).Space, 0u);
+  ASSERT_EQ(std::get<RootParam>(Elem).Visibility, ShaderVisibility::All);
 
   Elem = Elements[1];
   ASSERT_TRUE(std::holds_alternative<RootParam>(Elem));
   ASSERT_EQ(std::get<RootParam>(Elem).Type, ParamType::SRV);
   ASSERT_EQ(std::get<RootParam>(Elem).Reg.ViewType, RegisterType::TReg);
   ASSERT_EQ(std::get<RootParam>(Elem).Reg.Number, 42u);
+  ASSERT_EQ(std::get<RootParam>(Elem).Space, 4u);
+  ASSERT_EQ(std::get<RootParam>(Elem).Visibility, ShaderVisibility::Geometry);
 
   Elem = Elements[2];
   ASSERT_TRUE(std::holds_alternative<RootParam>(Elem));
   ASSERT_EQ(std::get<RootParam>(Elem).Type, ParamType::UAV);
   ASSERT_EQ(std::get<RootParam>(Elem).Reg.ViewType, RegisterType::UReg);
   ASSERT_EQ(std::get<RootParam>(Elem).Reg.Number, 34893247u);
+  ASSERT_EQ(std::get<RootParam>(Elem).Space, 0u);
+  ASSERT_EQ(std::get<RootParam>(Elem).Visibility, ShaderVisibility::Hull);
 
   ASSERT_TRUE(Consumer->isSatisfied());
 }
diff --git a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
index 5b2bb59fe9c02..7aa55215abae3 100644
--- a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
+++ b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
@@ -89,6 +89,8 @@ using ParamType = llvm::dxil::ResourceClass;
 struct RootParam {
   ParamType Type;
   Register Reg;
+  uint32_t Space = 0;
+  ShaderVisibility Visibility = ShaderVisibility::All;
 };
 
 // Models the end of a descriptor table and stores its visibility

[alsepkow]

Reviewed but don't have approval permissions yet. LGTM!