add support for parsing Flag parameters

inbelic · inbelic · commit 6775f3da0379 · 2025-01-28T22:10:44.000Z
- use DescriptorRangeFlags to demonstrate valid functionality
diff --git a/clang/include/clang/Parse/ParseHLSLRootSignature.h b/clang/include/clang/Parse/ParseHLSLRootSignature.h
@@ -125,9 +125,13 @@ class RootSignatureParser {
   bool ParseDescriptorRangeOffset(rs::DescriptorRangeOffset *X);
 
   // Various flags/enum parsing helpers
-  template <typename EnumType>
+  template <bool AllowZero = false, typename EnumType>
   bool ParseEnum(llvm::SmallDenseMap<TokenKind, EnumType> EnumMap,
                  EnumType *Enum);
+  template <typename FlagType>
+  bool ParseFlags(llvm::SmallDenseMap<TokenKind, FlagType> EnumMap,
+                  FlagType *Enum);
+  bool ParseDescriptorRangeFlags(rs::DescriptorRangeFlags *Enum);
   bool ParseShaderVisibility(rs::ShaderVisibility *Enum);
 
   // Increment the token iterator if we have not reached the end.
diff --git a/clang/lib/Parse/ParseHLSLRootSignature.cpp b/clang/lib/Parse/ParseHLSLRootSignature.cpp
@@ -330,6 +330,8 @@ bool RootSignatureParser::ParseDescriptorTableClause() {
     llvm_unreachable("Switch for an expected token was not provided");
     return true;
   }
+  Clause.SetDefaultFlags();
+
   if (ConsumeExpectedToken(TokenKind::pu_l_paren))
     return true;
 
@@ -345,6 +347,7 @@ bool RootSignatureParser::ParseDescriptorTableClause() {
       {TokenKind::kw_numDescriptors, &Clause.NumDescriptors},
       {TokenKind::kw_space, &Clause.Space},
       {TokenKind::kw_offset, &Clause.Offset},
+      {TokenKind::kw_flags, &Clause.Flags},
   };
   if (ParseOptionalParams({RefMap}))
     return true;
@@ -371,6 +374,9 @@ bool RootSignatureParser::ParseParam(ParamType Ref) {
                                [&](DescriptorRangeOffset *X) {
                                  Error = ParseDescriptorRangeOffset(X);
                                },
+                               [&](DescriptorRangeFlags *Flags) {
+                                 Error = ParseDescriptorRangeFlags(Flags);
+                               },
                                [&](ShaderVisibility *Enum) {
                                  Error = ParseShaderVisibility(Enum);
                                },
@@ -454,17 +460,29 @@ bool RootSignatureParser::ParseRegister(Register *Register) {
   return false;
 }
 
-template <typename EnumType>
+template <bool AllowZero, typename EnumType>
 bool RootSignatureParser::ParseEnum(
     llvm::SmallDenseMap<TokenKind, EnumType> EnumMap, EnumType *Enum) {
   SmallVector<TokenKind> EnumToks;
+  if (AllowZero)
+    EnumToks.push_back(TokenKind::int_literal); //  '0' is a valid flag value
   for (auto EnumPair : EnumMap)
     EnumToks.push_back(EnumPair.first);
 
   // If invoked we expect to have an enum
   if (ConsumeExpectedToken(EnumToks))
     return true;
 
+  // Handle the edge case when '0' is used to specify None
+  if (CurTok->Kind == TokenKind::int_literal) {
+    if (CurTok->NumLiteral.getInt() != 0) {
+      return true;
+    }
+    // Set enum to None equivalent
+    *Enum = EnumType(0);
+    return false;
+  }
+
   // Effectively a switch statement on the token kinds
   for (auto EnumPair : EnumMap)
     if (CurTok->Kind == EnumPair.first) {
@@ -476,6 +494,36 @@ bool RootSignatureParser::ParseEnum(
   return true;
 }
 
+template <typename FlagType>
+bool RootSignatureParser::ParseFlags(
+    llvm::SmallDenseMap<TokenKind, FlagType> FlagMap, FlagType *Flags) {
+  // Override the default value to 0 so that we can correctly 'or' the values
+  *Flags = FlagType(0);
+
+  do {
+    FlagType Flag;
+    if (ParseEnum<true>(FlagMap, &Flag))
+      return true;
+    // Store the 'or'
+    *Flags |= Flag;
+
+  } while (!TryConsumeExpectedToken(TokenKind::pu_or));
+
+  return false;
+}
+
+bool RootSignatureParser::ParseDescriptorRangeFlags(
+    DescriptorRangeFlags *Flags) {
+  // Define the possible flag kinds
+  llvm::SmallDenseMap<TokenKind, DescriptorRangeFlags> FlagMap = {
+#define DESCRIPTOR_RANGE_FLAG_ENUM(NAME, LIT, ON)                              \
+  {TokenKind::en_##NAME, DescriptorRangeFlags::NAME},
+#include "clang/Parse/HLSLRootSignatureTokenKinds.def"
+  };
+
+  return ParseFlags(FlagMap, Flags);
+}
+
 bool RootSignatureParser::ParseShaderVisibility(ShaderVisibility *Enum) {
   // Define the possible flag kinds
   llvm::SmallDenseMap<TokenKind, ShaderVisibility> EnumMap = {
diff --git a/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp b/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
@@ -310,9 +310,13 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
     DescriptorTable(
       visibility = SHADER_VISIBILITY_PIXEL,
       CBV(b0),
-      SRV(t42, space = 3, offset = 32, numDescriptors = 4),
+      SRV(t42, space = 3, offset = 32, numDescriptors = 4, flags = 0),
       Sampler(s987, space = 2, offset = DESCRIPTOR_RANGE_OFFSET_APPEND),
-      UAV(u987234)
+      UAV(u987234,
+        flags = Descriptors_Volatile | Data_Volatile
+                      | Data_Static_While_Set_At_Execute | Data_Static
+                      | Descriptors_Static_Keeping_Buffer_Bounds_Checks
+      )
     ),
     DescriptorTable()
   )cc";
@@ -343,6 +347,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, (uint32_t)0);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
             DescriptorRangeOffset(DescriptorTableOffsetAppend));
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Flags,
+            DescriptorRangeFlags::DataStaticWhileSetAtExecute);
 
   Elem = Elements[1];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -355,6 +361,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, (uint32_t)3);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
             DescriptorRangeOffset(32));
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Flags,
+            DescriptorRangeFlags::None);
 
   Elem = Elements[2];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -367,6 +375,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, (uint32_t)2);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
             DescriptorRangeOffset(DescriptorTableOffsetAppend));
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Flags,
+            DescriptorRangeFlags::None);
 
   Elem = Elements[3];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -379,6 +389,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Space, (uint32_t)0);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Offset,
             DescriptorRangeOffset(DescriptorTableOffsetAppend));
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Flags,
+            DescriptorRangeFlags::ValidFlags);
 
   Elem = Elements[4];
   ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));
diff --git a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
@@ -21,10 +21,43 @@ namespace llvm {
 namespace hlsl {
 namespace root_signature {
 
+#define RS_DEFINE_ENUM_CLASS_FLAGS_OPERATORS(Class)                            \
+  inline Class operator|(Class a, Class b) {                                   \
+    return static_cast<Class>(llvm::to_underlying(a) |                         \
+                              llvm::to_underlying(b));                         \
+  }                                                                            \
+  inline Class operator&(Class a, Class b) {                                   \
+    return static_cast<Class>(llvm::to_underlying(a) &                         \
+                              llvm::to_underlying(b));                         \
+  }                                                                            \
+  inline Class operator~(Class a) {                                            \
+    return static_cast<Class>(~llvm::to_underlying(a));                        \
+  }                                                                            \
+  inline Class &operator|=(Class &a, Class b) {                                \
+    a = a | b;                                                                 \
+    return a;                                                                  \
+  }                                                                            \
+  inline Class &operator&=(Class &a, Class b) {                                \
+    a = a & b;                                                                 \
+    return a;                                                                  \
+  }
+
 // Definition of the various enumerations and flags
 
 enum class DescriptorRangeOffset : uint32_t;
 
+enum class DescriptorRangeFlags : unsigned {
+  None = 0,
+  DescriptorsVolatile = 0x1,
+  DataVolatile = 0x2,
+  DataStaticWhileSetAtExecute = 0x4,
+  DataStatic = 0x8,
+  DescriptorsStaticKeepingBufferBoundsChecks = 0x10000,
+  ValidFlags = 0x1000f,
+  ValidSamplerFlags = DescriptorsVolatile,
+};
+RS_DEFINE_ENUM_CLASS_FLAGS_OPERATORS(DescriptorRangeFlags)
+
 enum class ShaderVisibility {
   All = 0,
   Vertex = 1,
@@ -61,6 +94,24 @@ struct DescriptorTableClause {
   uint32_t NumDescriptors = 1;
   uint32_t Space = 0;
   DescriptorRangeOffset Offset = DescriptorTableOffsetAppend;
+  DescriptorRangeFlags Flags;
+
+  void SetDefaultFlags() {
+    switch (Type) {
+    case ClauseType::CBuffer:
+      Flags = DescriptorRangeFlags::DataStaticWhileSetAtExecute;
+      break;
+    case ClauseType::SRV:
+      Flags = DescriptorRangeFlags::DataStaticWhileSetAtExecute;
+      break;
+    case ClauseType::UAV:
+      Flags = DescriptorRangeFlags::DataVolatile;
+      break;
+    case ClauseType::Sampler:
+      Flags = DescriptorRangeFlags::None;
+      break;
+    }
+  }
 };
 
 // Models RootElement : DescriptorTable | DescriptorTableClause
@@ -69,7 +120,8 @@ 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 *, DescriptorRangeOffset *,
-                               ShaderVisibility *>;
+                               DescriptorRangeFlags *, ShaderVisibility *>;
+
 } // namespace root_signature
 } // namespace hlsl
 } // namespace llvm
