add support for custom parameter parsing - DescriptorRangeOffset

Author: inbelic

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -122,6 +122,7 @@ class RootSignatureParser {
   // Common parsing helpers
   bool ParseRegister(rs::Register *Reg);
   bool ParseUInt(uint32_t *X);
+  bool ParseDescriptorRangeOffset(rs::DescriptorRangeOffset *X);
 
   // Increment the token iterator if we have not reached the end.
   // Return value denotes if we were already at the last token.

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -329,6 +329,7 @@ bool RootSignatureParser::ParseDescriptorTableClause() {
   llvm::SmallDenseMap<TokenKind, rs::ParamType> RefMap = {
       {TokenKind::kw_numDescriptors, &Clause.NumDescriptors},
       {TokenKind::kw_space, &Clause.Space},
+      {TokenKind::kw_offset, &Clause.Offset},
   };
   if (ParseOptionalParams({RefMap}))
     return true;
@@ -352,6 +353,9 @@ bool RootSignatureParser::ParseParam(ParamType Ref) {
 
   bool Error;
   std::visit(OverloadedMethods{[&](uint32_t *X) { Error = ParseUInt(X); },
+                               [&](DescriptorRangeOffset *X) {
+                                 Error = ParseDescriptorRangeOffset(X);
+                               },
   }, Ref);
 
   return Error;
@@ -385,6 +389,21 @@ bool RootSignatureParser::ParseOptionalParams(
   return false;
 }
 
+bool RootSignatureParser::ParseDescriptorRangeOffset(DescriptorRangeOffset *X) {
+  if (ConsumeExpectedToken(
+          {TokenKind::int_literal, TokenKind::en_DescriptorRangeOffsetAppend}))
+    return true;
+
+  // Edge case for the offset enum -> static value
+  if (CurTok->Kind == TokenKind::en_DescriptorRangeOffsetAppend) {
+    *X = DescriptorTableOffsetAppend;
+    return false;
+  }
+
+  *X = DescriptorRangeOffset(CurTok->NumLiteral.getInt().getExtValue());
+  return false;
+}
+
 bool RootSignatureParser::ParseUInt(uint32_t *X) {
   if (ConsumeExpectedToken(TokenKind::int_literal))
     return true;

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -309,8 +309,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   const llvm::StringLiteral Source = R"cc(
     DescriptorTable(
       CBV(b0),
-      SRV(t42, space = 3, numDescriptors = 4),
-      Sampler(s987, space = 2),
+      SRV(t42, space = 3, offset = 32, numDescriptors = 4),
+      Sampler(s987, space = 2, offset = DESCRIPTOR_RANGE_OFFSET_APPEND),
       UAV(u987234)
     ),
     DescriptorTable()
@@ -340,6 +340,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number, (uint32_t)0);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, (uint32_t)1);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, (uint32_t)0);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
+            DescriptorRangeOffset(DescriptorTableOffsetAppend));
 
   Elem = Elements[1];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -350,6 +352,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
             (uint32_t)42);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, (uint32_t)4);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, (uint32_t)3);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
+            DescriptorRangeOffset(32));
 
   Elem = Elements[2];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -360,6 +364,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
             (uint32_t)987);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, (uint32_t)1);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, (uint32_t)2);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
+            DescriptorRangeOffset(DescriptorTableOffsetAppend));
 
   Elem = Elements[3];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -370,6 +376,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
             (uint32_t)987234);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, (uint32_t)1);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, (uint32_t)0);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
+            DescriptorRangeOffset(DescriptorTableOffsetAppend));
 
   Elem = Elements[4];
   ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));

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

@@ -21,6 +21,10 @@ namespace llvm {
 namespace hlsl {
 namespace root_signature {
 
+// Definition of the various enumerations and flags
+
+enum class DescriptorRangeOffset : uint32_t;
+
 // Definitions of the in-memory data layout structures
 
 // Models the different registers: bReg | tReg | uReg | sReg
@@ -35,21 +39,24 @@ struct DescriptorTable {
   uint32_t NumClauses = 0; // The number of clauses in the table
 };
 
+static const DescriptorRangeOffset DescriptorTableOffsetAppend =
+    DescriptorRangeOffset(0xffffffff);
 // Models DTClause : CBV | SRV | UAV | Sampler, by collecting like parameters
 using ClauseType = llvm::dxil::ResourceClass;
 struct DescriptorTableClause {
   ClauseType Type;
   Register Register;
   uint32_t NumDescriptors = 1;
   uint32_t Space = 0;
+  DescriptorRangeOffset Offset = DescriptorTableOffsetAppend;
 };
 
 // Models RootElement : DescriptorTable | DescriptorTableClause
 using RootElement = std::variant<DescriptorTable, DescriptorTableClause>;
 
 // Models a reference to all assignment parameter types that any RootElement
 // may have. Things of the form: Keyword = Param
-using ParamType = std::variant<uint32_t *>;
+using ParamType = std::variant<uint32_t *, DescriptorRangeOffset *>;
 } // namespace root_signature
 } // namespace hlsl
 } // namespace llvm