[HLSL][RootSignature] Add parsing of remaining Descriptor Table params (llvm#137038)

inbelic · web-flow · commit 9b74dce4a326 · 2025-04-25T14:38:02.000-07:00
- defines the special values for `DESCRIPTOR_RANGE_OFFSET_APPEND` and `unbounded` for the `offset` and `numDescriptors` parameters respectively - adds these parmaters to the `DescriptorClause` struct and the params struct - plugs in parsing of `numDescriptors` and `offset` into `parseDescriptorTableClauseParams` - defines the `unbounded` enum keyword for the lexer to expose to the parser - adds corresponding unit tests Part 5 of llvm#126569
diff --git a/clang/include/clang/Lex/HLSLRootSignatureTokenKinds.def b/clang/include/clang/Lex/HLSLRootSignatureTokenKinds.def
@@ -27,6 +27,9 @@
 #endif
 
 // Defines the various types of enum
+#ifndef UNBOUNDED_ENUM
+#define UNBOUNDED_ENUM(NAME, LIT) ENUM(NAME, LIT)
+#endif
 #ifndef DESCRIPTOR_RANGE_OFFSET_ENUM
 #define DESCRIPTOR_RANGE_OFFSET_ENUM(NAME, LIT) ENUM(NAME, LIT)
 #endif
@@ -87,6 +90,9 @@ KEYWORD(flags)
 KEYWORD(numDescriptors)
 KEYWORD(offset)
 
+// Unbounded Enum:
+UNBOUNDED_ENUM(unbounded, "unbounded")
+
 // Descriptor Range Offset Enum:
 DESCRIPTOR_RANGE_OFFSET_ENUM(DescriptorRangeOffsetAppend, "DESCRIPTOR_RANGE_OFFSET_APPEND")
 
@@ -118,6 +124,7 @@ SHADER_VISIBILITY_ENUM(Mesh, "SHADER_VISIBILITY_MESH")
 #undef DESCRIPTOR_RANGE_FLAG_ENUM_ON
 #undef ROOT_DESCRIPTOR_FLAG_ENUM
 #undef DESCRIPTOR_RANGE_OFFSET_ENUM
+#undef UNBOUNDED_ENUM
 #undef ENUM
 #undef KEYWORD
 #undef PUNCTUATOR
diff --git a/clang/include/clang/Parse/ParseHLSLRootSignature.h b/clang/include/clang/Parse/ParseHLSLRootSignature.h
@@ -80,7 +80,9 @@ class RootSignatureParser {
   /// state of parsed params
   struct ParsedClauseParams {
     std::optional<llvm::hlsl::rootsig::Register> Reg;
+    std::optional<uint32_t> NumDescriptors;
     std::optional<uint32_t> Space;
+    std::optional<uint32_t> Offset;
     std::optional<llvm::hlsl::rootsig::DescriptorRangeFlags> Flags;
   };
   std::optional<ParsedClauseParams>
diff --git a/clang/lib/Parse/ParseHLSLRootSignature.cpp b/clang/lib/Parse/ParseHLSLRootSignature.cpp
@@ -145,9 +145,15 @@ RootSignatureParser::parseDescriptorTableClause() {
   Clause.Reg = Params->Reg.value();
 
   // Fill in optional values
+  if (Params->NumDescriptors.has_value())
+    Clause.NumDescriptors = Params->NumDescriptors.value();
+
   if (Params->Space.has_value())
     Clause.Space = Params->Space.value();
 
+  if (Params->Offset.has_value())
+    Clause.Offset = Params->Offset.value();
+
   if (Params->Flags.has_value())
     Clause.Flags = Params->Flags.value();
 
@@ -182,6 +188,29 @@ RootSignatureParser::parseDescriptorTableClauseParams(TokenKind RegType) {
       Params.Reg = Reg;
     }
 
+    // `numDescriptors` `=` POS_INT | unbounded
+    if (tryConsumeExpectedToken(TokenKind::kw_numDescriptors)) {
+      if (Params.NumDescriptors.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;
+
+      std::optional<uint32_t> NumDescriptors;
+      if (tryConsumeExpectedToken(TokenKind::en_unbounded))
+        NumDescriptors = NumDescriptorsUnbounded;
+      else {
+        NumDescriptors = parseUIntParam();
+        if (!NumDescriptors.has_value())
+          return std::nullopt;
+      }
+
+      Params.NumDescriptors = NumDescriptors;
+    }
+
     // `space` `=` POS_INT
     if (tryConsumeExpectedToken(TokenKind::kw_space)) {
       if (Params.Space.has_value()) {
@@ -199,6 +228,29 @@ RootSignatureParser::parseDescriptorTableClauseParams(TokenKind RegType) {
       Params.Space = Space;
     }
 
+    // `offset` `=` POS_INT | DESCRIPTOR_RANGE_OFFSET_APPEND
+    if (tryConsumeExpectedToken(TokenKind::kw_offset)) {
+      if (Params.Offset.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;
+
+      std::optional<uint32_t> Offset;
+      if (tryConsumeExpectedToken(TokenKind::en_DescriptorRangeOffsetAppend))
+        Offset = DescriptorTableOffsetAppend;
+      else {
+        Offset = parseUIntParam();
+        if (!Offset.has_value())
+          return std::nullopt;
+      }
+
+      Params.Offset = Offset;
+    }
+
     // `flags` `=` DESCRIPTOR_RANGE_FLAGS
     if (tryConsumeExpectedToken(TokenKind::kw_flags)) {
       if (Params.Flags.has_value()) {
diff --git a/clang/unittests/Lex/LexHLSLRootSignatureTest.cpp b/clang/unittests/Lex/LexHLSLRootSignatureTest.cpp
@@ -93,6 +93,7 @@ TEST_F(LexHLSLRootSignatureTest, ValidLexAllTokensTest) {
     space visibility flags
     numDescriptors offset
 
+    unbounded
     DESCRIPTOR_RANGE_OFFSET_APPEND
 
     DATA_VOLATILE
diff --git a/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp b/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
@@ -130,10 +130,10 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   const llvm::StringLiteral Source = R"cc(
     DescriptorTable(
       CBV(b0),
-      SRV(space = 3, t42, flags = 0),
+      SRV(space = 3, offset = 32, t42, flags = 0, numDescriptors = 4),
       visibility = SHADER_VISIBILITY_PIXEL,
-      Sampler(s987, space = +2),
-      UAV(u4294967294,
+      Sampler(s987, space = +2, offset = DESCRIPTOR_RANGE_OFFSET_APPEND),
+      UAV(u4294967294, numDescriptors = unbounded,
         flags = Descriptors_Volatile | Data_Volatile
                       | Data_Static_While_Set_At_Execute | Data_Static
                       | Descriptors_Static_Keeping_Buffer_Bounds_Checks
@@ -162,7 +162,10 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Reg.ViewType,
             RegisterType::BReg);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Reg.Number, 0u);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, 1u);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, 0u);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
+            DescriptorTableOffsetAppend);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Flags,
             DescriptorRangeFlags::DataStaticWhileSetAtExecute);
 
@@ -172,7 +175,9 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Reg.ViewType,
             RegisterType::TReg);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Reg.Number, 42u);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, 4u);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, 3u);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset, 32u);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Flags,
             DescriptorRangeFlags::None);
 
@@ -182,7 +187,10 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Reg.ViewType,
             RegisterType::SReg);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Reg.Number, 987u);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, 1u);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, 2u);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
+            DescriptorTableOffsetAppend);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Flags,
             DescriptorRangeFlags::None);
 
@@ -192,7 +200,11 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Reg.ViewType,
             RegisterType::UReg);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Reg.Number, 4294967294u);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors,
+            NumDescriptorsUnbounded);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, 0u);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
+            DescriptorTableOffsetAppend);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Flags,
             DescriptorRangeFlags::ValidFlags);
 
diff --git a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
@@ -60,12 +60,16 @@ struct DescriptorTable {
   uint32_t NumClauses = 0; // The number of clauses in the table
 };
 
+static const uint32_t NumDescriptorsUnbounded = 0xffffffff;
+static const uint32_t DescriptorTableOffsetAppend = 0xffffffff;
 // Models DTClause : CBV | SRV | UAV | Sampler, by collecting like parameters
 using ClauseType = llvm::dxil::ResourceClass;
 struct DescriptorTableClause {
   ClauseType Type;
   Register Reg;
+  uint32_t NumDescriptors = 1;
   uint32_t Space = 0;
+  uint32_t Offset = DescriptorTableOffsetAppend;
   DescriptorRangeFlags Flags;
 
   void setDefaultFlags() {
