[llvm-branch-commits] [clang] [llvm] [HLSL][RootSignature] Add parsing of Register in params for RootParams (PR #140148)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-clang

Author: Finn Plummer (inbelic)

<details>
<summary>Changes</summary>

- defines the `parseRootParamParams` infastructure for parsing the params of a RootParam
- defines the register type to illustrate use
- add tests to demonstrate functionality

Part 2 of https://github.com/llvm/llvm-project/issues/126577

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


4 Files Affected:

- (modified) clang/include/clang/Parse/ParseHLSLRootSignature.h (+6) 
- (modified) clang/lib/Parse/ParseHLSLRootSignature.cpp (+42) 
- (modified) clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp (+9-4) 
- (modified) llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h (+1) 


``````````diff
diff --git a/clang/include/clang/Parse/ParseHLSLRootSignature.h b/clang/include/clang/Parse/ParseHLSLRootSignature.h
index b96416ab6c6c4..92b8e7389426a 100644
--- a/clang/include/clang/Parse/ParseHLSLRootSignature.h
+++ b/clang/include/clang/Parse/ParseHLSLRootSignature.h
@@ -89,6 +89,12 @@ class RootSignatureParser {
   };
   std::optional<ParsedConstantParams> parseRootConstantParams();
 
+  struct ParsedRootParamParams {
+    std::optional<llvm::hlsl::rootsig::Register> Reg;
+  };
+  std::optional<ParsedRootParamParams>
+  parseRootParamParams(RootSignatureToken::Kind RegType);
+
   struct ParsedClauseParams {
     std::optional<llvm::hlsl::rootsig::Register> Reg;
     std::optional<uint32_t> NumDescriptors;
diff --git a/clang/lib/Parse/ParseHLSLRootSignature.cpp b/clang/lib/Parse/ParseHLSLRootSignature.cpp
index 3ed60442eaa14..de7c2452b3fd4 100644
--- a/clang/lib/Parse/ParseHLSLRootSignature.cpp
+++ b/clang/lib/Parse/ParseHLSLRootSignature.cpp
@@ -176,20 +176,37 @@ std::optional<RootParam> RootSignatureParser::parseRootParam() {
     return std::nullopt;
 
   RootParam Param;
+  TokenKind ExpectedReg;
   switch (ParamKind) {
   default:
     llvm_unreachable("Switch for consumed token was not provided");
   case TokenKind::kw_CBV:
     Param.Type = ParamType::CBuffer;
+    ExpectedReg = TokenKind::bReg;
     break;
   case TokenKind::kw_SRV:
     Param.Type = ParamType::SRV;
+    ExpectedReg = TokenKind::tReg;
     break;
   case TokenKind::kw_UAV:
     Param.Type = ParamType::UAV;
+    ExpectedReg = TokenKind::uReg;
     break;
   }
 
+  auto Params = parseRootParamParams(ExpectedReg);
+  if (!Params.has_value())
+    return std::nullopt;
+
+  // Check mandatory parameters were provided
+  if (!Params->Reg.has_value()) {
+    getDiags().Report(CurToken.TokLoc, diag::err_hlsl_rootsig_missing_param)
+        << ExpectedReg;
+    return std::nullopt;
+  }
+
+  Param.Reg = Params->Reg.value();
+
   if (consumeExpectedToken(TokenKind::pu_r_paren,
                            diag::err_hlsl_unexpected_end_of_params,
                            /*param of=*/TokenKind::kw_RootConstants))
@@ -398,6 +415,31 @@ RootSignatureParser::parseRootConstantParams() {
   return Params;
 }
 
+std::optional<RootSignatureParser::ParsedRootParamParams>
+RootSignatureParser::parseRootParamParams(TokenKind RegType) {
+  assert(CurToken.TokKind == TokenKind::pu_l_paren &&
+         "Expects to only be invoked starting at given token");
+
+  ParsedRootParamParams Params;
+  do {
+    // ( `b` | `t` | `u`) POS_INT
+    if (tryConsumeExpectedToken(RegType)) {
+      if (Params.Reg.has_value()) {
+        getDiags().Report(CurToken.TokLoc, diag::err_hlsl_rootsig_repeat_param)
+            << CurToken.TokKind;
+        return std::nullopt;
+      }
+      auto Reg = parseRegister();
+      if (!Reg.has_value())
+        return std::nullopt;
+      Params.Reg = Reg;
+    }
+
+  } while (tryConsumeExpectedToken(TokenKind::pu_comma));
+
+  return Params;
+}
+
 std::optional<RootSignatureParser::ParsedClauseParams>
 RootSignatureParser::parseDescriptorTableClauseParams(TokenKind RegType) {
   assert(CurToken.TokKind == TokenKind::pu_l_paren &&
diff --git a/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp b/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
index 876dcd17f0389..0f32523493a53 100644
--- a/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
+++ b/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
@@ -346,9 +346,9 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseRootFlagsTest) {
 
 TEST_F(ParseHLSLRootSignatureTest, ValidParseRootParamsTest) {
   const llvm::StringLiteral Source = R"cc(
-    CBV(),
-    SRV(),
-    UAV()
+    CBV(b0),
+    SRV(t42),
+    UAV(u34893247)
   )cc";
 
   TrivialModuleLoader ModLoader;
@@ -368,15 +368,20 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseRootParamsTest) {
 
   RootElement Elem = Elements[0];
   ASSERT_TRUE(std::holds_alternative<RootParam>(Elem));
-  ASSERT_EQ(std::get<RootParam>(Elem).Type, ParamType::CBuffer);
+  ASSERT_EQ(std::get<RootParam>(Elem).Reg.ViewType, RegisterType::BReg);
+  ASSERT_EQ(std::get<RootParam>(Elem).Reg.Number, 0u);
 
   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);
 
   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_TRUE(Consumer->isSatisfied());
 }
diff --git a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
index 2bcae22c6cfb2..5b2bb59fe9c02 100644
--- a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
+++ b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
@@ -88,6 +88,7 @@ struct RootConstants {
 using ParamType = llvm::dxil::ResourceClass;
 struct RootParam {
   ParamType Type;
+  Register Reg;
 };
 
 // Models the end of a descriptor table and stores its visibility

``````````

</details>


https://github.com/llvm/llvm-project/pull/140148