add support for custom parameter parsing - DescriptorRangeOffset

Author: inbelic

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -132,6 +132,8 @@ class RootSignatureParser {
   // Common parsing helpers
   bool ParseRegister(llvm::hlsl::rootsig::Register *Reg);
   bool ParseUInt(uint32_t *X);
+  bool
+  ParseDescriptorRangeOffset(llvm::hlsl::rootsig::DescriptorRangeOffset *X);
 
   /// Invoke the lexer to consume a token and update CurToken with the result
   ///

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -324,6 +324,7 @@ bool RootSignatureParser::ParseDescriptorTableClause() {
   llvm::SmallDenseMap<TokenKind, ParamType> RefMap = {
       {TokenKind::kw_numDescriptors, &Clause.NumDescriptors},
       {TokenKind::kw_space, &Clause.Space},
+      {TokenKind::kw_offset, &Clause.Offset},
   };
   if (ParseOptionalParams({RefMap}))
     return true;
@@ -347,6 +348,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;
@@ -380,6 +384,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 (CurToken.Kind == TokenKind::en_DescriptorRangeOffsetAppend) {
+    *X = DescriptorTableOffsetAppend;
+    return false;
+  }
+
+  *X = DescriptorRangeOffset(CurToken.NumLiteral.getInt().getExtValue());
+  return false;
+}
+
 bool RootSignatureParser::ParseUInt(uint32_t *X) {
   // Treat a postively signed integer as though it is unsigned to match DXC
   TryConsumeExpectedToken(TokenKind::pu_plus);

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -339,8 +339,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()
@@ -366,6 +366,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));
@@ -376,6 +378,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));
@@ -386,6 +390,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));
@@ -396,6 +402,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 rootsig {
 
+// 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 rootsig
 } // namespace hlsl