[HLSL] Update Frontend to support version 1.2 of root signature

clang/include/clang/Lex/HLSLRootSignatureTokenKinds.def

@@ -65,6 +65,9 @@
 #ifndef STATIC_BORDER_COLOR_ENUM
 #define STATIC_BORDER_COLOR_ENUM(NAME, LIT) ENUM(NAME, LIT)
 #endif
+#ifndef STATIC_SAMPLER_FLAG_ENUM
+#define STATIC_SAMPLER_FLAG_ENUM(NAME, LIT) ENUM(NAME, LIT)
+#endif
 
 // General Tokens:
 TOK(invalid, "invalid identifier")
@@ -228,6 +231,10 @@ STATIC_BORDER_COLOR_ENUM(OpaqueWhite, "STATIC_BORDER_COLOR_OPAQUE_WHITE")
 STATIC_BORDER_COLOR_ENUM(OpaqueBlackUint, "STATIC_BORDER_COLOR_OPAQUE_BLACK_UINT")
 STATIC_BORDER_COLOR_ENUM(OpaqueWhiteUint, "STATIC_BORDER_COLOR_OPAQUE_WHITE_UINT")
 
+// Root Descriptor Flag Enums:
+STATIC_SAMPLER_FLAG_ENUM(UintBorderColor, "UINT_BORDER_COLOR")
+STATIC_SAMPLER_FLAG_ENUM(NonNormalizedCoordinates, "NON_NORMALIZED_COORDINATES")
+
 #undef STATIC_BORDER_COLOR_ENUM
 #undef COMPARISON_FUNC_ENUM
 #undef TEXTURE_ADDRESS_MODE_ENUM
@@ -237,6 +244,7 @@ STATIC_BORDER_COLOR_ENUM(OpaqueWhiteUint, "STATIC_BORDER_COLOR_OPAQUE_WHITE_UINT
 #undef DESCRIPTOR_RANGE_FLAG_ENUM_OFF
 #undef DESCRIPTOR_RANGE_FLAG_ENUM_ON
 #undef ROOT_DESCRIPTOR_FLAG_ENUM
+#undef STATIC_SAMPLER_FLAG_ENUM
 #undef ROOT_FLAG_ENUM
 #undef DESCRIPTOR_RANGE_OFFSET_ENUM
 #undef UNBOUNDED_ENUM

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -130,6 +130,7 @@ class RootSignatureParser {
     std::optional<float> MaxLOD;
     std::optional<uint32_t> Space;
     std::optional<llvm::dxbc::ShaderVisibility> Visibility;
+    std::optional<llvm::dxbc::StaticSamplerFlags> Flags;
   };
   std::optional<ParsedStaticSamplerParams> parseStaticSamplerParams();
 
@@ -153,6 +154,8 @@ class RootSignatureParser {
   parseRootDescriptorFlags(RootSignatureToken::Kind Context);
   std::optional<llvm::dxbc::DescriptorRangeFlags>
   parseDescriptorRangeFlags(RootSignatureToken::Kind Context);
+  std::optional<llvm::dxbc::StaticSamplerFlags>
+  parseStaticSamplerFlags(RootSignatureToken::Kind Context);
 
   /// Use NumericLiteralParser to convert CurToken.NumSpelling into a unsigned
   /// 32-bit integer

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -485,6 +485,9 @@ std::optional<StaticSampler> RootSignatureParser::parseStaticSampler() {
   if (Params->Visibility.has_value())
     Sampler.Visibility = Params->Visibility.value();
 
+  if (Params->Flags.has_value())
+    Sampler.Flags = Params->Flags.value();
+
   return Sampler;
 }
 
@@ -926,6 +929,20 @@ RootSignatureParser::parseStaticSamplerParams() {
       if (!Visibility.has_value())
         return std::nullopt;
       Params.Visibility = Visibility;
+    } else if (tryConsumeExpectedToken(TokenKind::kw_flags)) {
+      // `flags` `=` UINT_BORDER_COLOR
+      if (Params.Flags.has_value()) {
+        reportDiag(diag::err_hlsl_rootsig_repeat_param) << CurToken.TokKind;
+        return std::nullopt;
+      }
+
+      if (consumeExpectedToken(TokenKind::pu_equal))
+        return std::nullopt;
+
+      auto Flags = parseStaticSamplerFlags(TokenKind::kw_flags);
+      if (!Flags.has_value())
+        return std::nullopt;
+      Params.Flags = Flags;
     } else {
       consumeNextToken(); // let diagnostic be at the start of invalid token
       reportDiag(diag::err_hlsl_invalid_token)
@@ -1255,6 +1272,50 @@ RootSignatureParser::parseDescriptorRangeFlags(TokenKind Context) {
   return Flags;
 }
 
+std::optional<llvm::dxbc::StaticSamplerFlags>
+RootSignatureParser::parseStaticSamplerFlags(TokenKind Context) {
+  assert(CurToken.TokKind == TokenKind::pu_equal &&
+         "Expects to only be invoked starting at given keyword");
+
+  // Handle the edge-case of '0' to specify no flags set
+  if (tryConsumeExpectedToken(TokenKind::int_literal)) {
+    if (!verifyZeroFlag()) {
+      reportDiag(diag::err_hlsl_rootsig_non_zero_flag);
+      return std::nullopt;
+    }
+    return llvm::dxbc::StaticSamplerFlags::None;
+  }
+
+  TokenKind Expected[] = {
+#define STATIC_SAMPLER_FLAG_ENUM(NAME, LIT) TokenKind::en_##NAME,
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+  };
+
+  std::optional<llvm::dxbc::StaticSamplerFlags> Flags;
+
+  do {
+    if (tryConsumeExpectedToken(Expected)) {
+      switch (CurToken.TokKind) {
+#define STATIC_SAMPLER_FLAG_ENUM(NAME, LIT)                                    \
+  case TokenKind::en_##NAME:                                                   \
+    Flags = maybeOrFlag<llvm::dxbc::StaticSamplerFlags>(                       \
+        Flags, llvm::dxbc::StaticSamplerFlags::NAME);                          \
+    break;
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+      default:
+        llvm_unreachable("Switch for consumed enum token was not provided");
+      }
+    } else {
+      consumeNextToken(); // consume token to point at invalid token
+      reportDiag(diag::err_hlsl_invalid_token)
+          << /*value=*/1 << /*value of*/ Context;
+      return std::nullopt;
+    }
+  } while (tryConsumeExpectedToken(TokenKind::pu_or));
+
+  return Flags;
+}
+
 std::optional<uint32_t> RootSignatureParser::handleUIntLiteral() {
   // Parse the numeric value and do semantic checks on its specification
   clang::NumericLiteralParser Literal(

clang/test/CodeGenHLSL/RootSignature.hlsl

@@ -82,8 +82,8 @@ void RootDescriptorsEntry() {}
 // checking minLOD, maxLOD
 // CHECK-SAME: float -1.280000e+02, float 1.280000e+02,
 
-// checking register, space and visibility
-// CHECK-SAME: i32 42, i32 0, i32 0}
+// checking register, space, visibility and flag
+// CHECK-SAME: i32 42, i32 0, i32 0, i32 1}
 
 #define SampleStaticSampler \
   "StaticSampler(s42, " \
@@ -96,6 +96,7 @@ void RootDescriptorsEntry() {}
   " borderColor = STATIC_BORDER_COLOR_OPAQUE_WHITE, " \
   " minLOD = -128.f, maxLOD = 128.f, " \
   " space = 0, visibility = SHADER_VISIBILITY_ALL, " \
+  " flags = UINT_BORDER_COLOR" \
   ")"
 [shader("compute"), RootSignature(SampleStaticSampler)]
 [numthreads(1,1,1)]

clang/unittests/Lex/LexHLSLRootSignatureTest.cpp

@@ -226,6 +226,9 @@ TEST_F(LexHLSLRootSignatureTest, ValidLexAllTokensTest) {
     STATIC_BORDER_COLOR_OPAQUE_WHITE
     STATIC_BORDER_COLOR_OPAQUE_BLACK_UINT
     STATIC_BORDER_COLOR_OPAQUE_WHITE_UINT
+
+    UINT_BORDER_COLOR
+    NON_NORMALIZED_COORDINATES
   )cc";
   hlsl::RootSignatureLexer Lexer(Source);
 

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -263,7 +263,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseStaticSamplerTest) {
       filter = FILTER_MAXIMUM_MIN_POINT_MAG_LINEAR_MIP_POINT,
       maxLOD = 9000, addressU = TEXTURE_ADDRESS_MIRROR,
       comparisonFunc = COMPARISON_NOT_EQUAL,
-      borderColor = STATIC_BORDER_COLOR_OPAQUE_BLACK_UINT
+      borderColor = STATIC_BORDER_COLOR_OPAQUE_BLACK_UINT,
+      flags = 0
     )
   )cc";
 
@@ -336,6 +337,37 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseStaticSamplerTest) {
   ASSERT_TRUE(Consumer->isSatisfied());
 }
 
+TEST_F(ParseHLSLRootSignatureTest, ValidStaticSamplerFlagsTest) {
+  const llvm::StringLiteral Source = R"cc(
+    StaticSampler(s0, flags = UINT_BORDER_COLOR | NON_NORMALIZED_COORDINATES)
+  )cc";
+
+  auto Ctx = createMinimalASTContext();
+  StringLiteral *Signature = wrapSource(Ctx, Source);
+
+  TrivialModuleLoader ModLoader;
+  auto PP = createPP(Source, ModLoader);
+
+  hlsl::RootSignatureParser Parser(RootSignatureVersion::V1_1, Signature, *PP);
+
+  // Test no diagnostics produced
+  Consumer->setNoDiag();
+
+  ASSERT_FALSE(Parser.parse());
+
+  auto Elements = Parser.getElements();
+  ASSERT_EQ(Elements.size(), 1u);
+
+  RootElement Elem = Elements[0].getElement();
+  ASSERT_TRUE(std::holds_alternative<StaticSampler>(Elem));
+  auto ValidStaticSamplerFlags =
+      llvm::dxbc::StaticSamplerFlags::NonNormalizedCoordinates |
+      llvm::dxbc::StaticSamplerFlags::UintBorderColor;
+  ASSERT_EQ(std::get<StaticSampler>(Elem).Flags, ValidStaticSamplerFlags);
+
+  ASSERT_TRUE(Consumer->isSatisfied());
+}
+
 TEST_F(ParseHLSLRootSignatureTest, ValidParseFloatsTest) {
   const llvm::StringLiteral Source = R"cc(
     StaticSampler(s0, mipLODBias = 0),

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -185,6 +185,15 @@ enum class DescriptorRangeFlags : uint32_t {
 
 LLVM_ABI ArrayRef<EnumEntry<DescriptorRangeFlags>> getDescriptorRangeFlags();
 
+#define STATIC_SAMPLER_FLAG(Num, Enum, Flag) Enum = Num,
+enum class StaticSamplerFlags : uint32_t {
+#include "DXContainerConstants.def"
+
+  LLVM_MARK_AS_BITMASK_ENUM(NonNormalizedCoordinates)
+};
+
+LLVM_ABI ArrayRef<EnumEntry<StaticSamplerFlags>> getStaticSamplerFlags();
+
 #define ROOT_PARAMETER(Val, Enum) Enum = Val,
 enum class RootParameterType : uint32_t {
 #include "DXContainerConstants.def"
@@ -813,6 +822,21 @@ struct DescriptorRange {
   }
 };
 } // namespace v2
+namespace v3 {
+struct StaticSampler : public v1::StaticSampler {
+  uint32_t Flags;
+
+  StaticSampler() = default;
+  explicit StaticSampler(v1::StaticSampler &Base)
+      : v1::StaticSampler(Base), Flags(0U) {}
+
+  void swapBytes() {
+    v1::StaticSampler::swapBytes();
+    sys::swapByteOrder(Flags);
+  }
+};
+
+} // namespace v3
 } // namespace RTS0
 
 // D3D_ROOT_SIGNATURE_VERSION

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

@@ -104,6 +104,16 @@ DESCRIPTOR_RANGE_FLAG(0x10000, DescriptorsStaticKeepingBufferBoundsChecks, DESCR
 #undef DESCRIPTOR_RANGE_FLAG
 #endif // DESCRIPTOR_RANGE_FLAG
 
+// STATIC_SAMPLER_FLAG(flag value, name, flag).
+#ifdef STATIC_SAMPLER_FLAG
+
+STATIC_SAMPLER_FLAG(0x0, None, SAMPLER_FLAG_NONE)
+STATIC_SAMPLER_FLAG(0x1, UintBorderColor, SAMPLER_FLAG_UINT_BORDER_COLOR)
+STATIC_SAMPLER_FLAG(0x2, NonNormalizedCoordinates, SAMPLER_FLAG_NON_NORMALIZED_COORDINATES)
+
+#undef STATIC_SAMPLER_FLAG
+#endif // STATIC_SAMPLER_FLAG
+
 #ifdef ROOT_PARAMETER
 
 ROOT_PARAMETER(0, DescriptorTable)

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

@@ -131,6 +131,7 @@ struct StaticSampler {
   float MaxLOD = std::numeric_limits<float>::max();
   uint32_t Space = 0;
   dxbc::ShaderVisibility Visibility = dxbc::ShaderVisibility::All;
+  dxbc::StaticSamplerFlags Flags = dxbc::StaticSamplerFlags::None;
 };
 
 /// Models RootElement : RootFlags | RootConstants | RootParam

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -74,6 +74,7 @@ struct StaticSampler {
   uint32_t ShaderRegister;
   uint32_t RegisterSpace;
   dxbc::ShaderVisibility ShaderVisibility;
+  uint32_t Flags;
 };
 
 struct RootParametersContainer {

llvm/include/llvm/Object/DXContainer.h

@@ -228,11 +228,11 @@ class RootSignature {
   uint32_t Flags;
   ViewArray<dxbc::RTS0::v1::RootParameterHeader> ParametersHeaders;
   StringRef PartData;
-  ViewArray<dxbc::RTS0::v1::StaticSampler> StaticSamplers;
+  ViewArray<dxbc::RTS0::v3::StaticSampler> StaticSamplers;
 
   using param_header_iterator =
       ViewArray<dxbc::RTS0::v1::RootParameterHeader>::iterator;
-  using samplers_iterator = ViewArray<dxbc::RTS0::v1::StaticSampler>::iterator;
+  using samplers_iterator = ViewArray<dxbc::RTS0::v3::StaticSampler>::iterator;
 
 public:
   RootSignature(StringRef PD) : PartData(PD) {}

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -178,6 +178,11 @@ struct StaticSamplerYamlDesc {
   uint32_t ShaderRegister;
   uint32_t RegisterSpace;
   dxbc::ShaderVisibility ShaderVisibility;
+
+  LLVM_ABI uint32_t getEncodedFlags() const;
+
+#define STATIC_SAMPLER_FLAG(Num, Enum, Flag) bool Enum = false;
+#include "llvm/BinaryFormat/DXContainerConstants.def"
 };
 
 struct RootSignatureYamlDesc {

llvm/lib/BinaryFormat/DXContainer.cpp

@@ -89,6 +89,15 @@ ArrayRef<EnumEntry<DescriptorRangeFlags>> dxbc::getDescriptorRangeFlags() {
   return ArrayRef(DescriptorRangeFlagNames);
 }
 
+static const EnumEntry<StaticSamplerFlags> StaticSamplerFlagNames[] = {
+#define STATIC_SAMPLER_FLAG(Val, Enum, Flag) {#Enum, StaticSamplerFlags::Enum},
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<StaticSamplerFlags>> dxbc::getStaticSamplerFlags() {
+  return ArrayRef(StaticSamplerFlagNames);
+}
+
 #define SHADER_VISIBILITY(Val, Enum) {#Enum, ShaderVisibility::Enum},
 
 static const EnumEntry<ShaderVisibility> ShaderVisibilityValues[] = {

llvm/lib/Frontend/HLSL/HLSLRootSignature.cpp

@@ -172,7 +172,7 @@ raw_ostream &operator<<(raw_ostream &OS, const StaticSampler &Sampler) {
      << ", borderColor = " << Sampler.BorderColor
      << ", minLOD = " << Sampler.MinLOD << ", maxLOD = " << Sampler.MaxLOD
      << ", space = " << Sampler.Space << ", visibility = " << Sampler.Visibility
-     << ")";
+     << ", flags = " << Sampler.Flags << ")";
   return OS;
 }
 

llvm/lib/Frontend/HLSL/RootSignatureMetadata.cpp

@@ -212,6 +212,7 @@ MDNode *MetadataBuilder::BuildStaticSampler(const StaticSampler &Sampler) {
       ConstantAsMetadata::get(Builder.getInt32(Sampler.Space)),
       ConstantAsMetadata::get(
           Builder.getInt32(to_underlying(Sampler.Visibility))),
+      ConstantAsMetadata::get(Builder.getInt32(to_underlying(Sampler.Flags))),
   };
   return MDNode::get(Ctx, Operands);
 }

llvm/lib/Frontend/HLSL/RootSignatureValidations.cpp

@@ -20,7 +20,9 @@ namespace rootsig {
 
 bool verifyRootFlag(uint32_t Flags) { return (Flags & ~0xfff) == 0; }
 
-bool verifyVersion(uint32_t Version) { return (Version == 1 || Version == 2); }
+bool verifyVersion(uint32_t Version) {
+  return (Version == 1 || Version == 2 || Version == 3);
+}
 
 bool verifyRegisterValue(uint32_t RegisterValue) {
   return RegisterValue != ~0U;

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -33,7 +33,9 @@ static uint32_t rewriteOffsetToCurrentByte(raw_svector_ostream &Stream,
 size_t RootSignatureDesc::getSize() const {
   uint32_t StaticSamplersOffset = computeStaticSamplersOffset();
   size_t StaticSamplersSize =
-      StaticSamplers.size() * sizeof(dxbc::RTS0::v1::StaticSampler);
+      (Version > 2 ? sizeof(dxbc::RTS0::v3::StaticSampler)
+                   : sizeof(dxbc::RTS0::v1::StaticSampler)) *
+      StaticSamplers.size();
 
   return size_t(StaticSamplersOffset) + StaticSamplersSize;
 }
@@ -171,6 +173,9 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
     support::endian::write(BOS, S.ShaderRegister, llvm::endianness::little);
     support::endian::write(BOS, S.RegisterSpace, llvm::endianness::little);
     support::endian::write(BOS, S.ShaderVisibility, llvm::endianness::little);
+
+    if (Version > 2)
+      support::endian::write(BOS, S.Flags, llvm::endianness::little);
   }
   assert(Storage.size() == getSize());
   OS.write(Storage.data(), Storage.size());

llvm/lib/Object/DXContainer.cpp

@@ -276,10 +276,13 @@ Error DirectX::RootSignature::parse() {
       RootParametersOffset,
       NumParameters * sizeof(dxbc::RTS0::v1::RootParameterHeader));
 
-  StaticSamplers.Stride = sizeof(dxbc::RTS0::v1::StaticSampler);
-  StaticSamplers.Data = PartData.substr(
-      StaticSamplersOffset,
-      NumStaticSamplers * sizeof(dxbc::RTS0::v1::StaticSampler));
+  StaticSamplers.Stride = (Version <= 2)
+                              ? sizeof(dxbc::RTS0::v1::StaticSampler)
+                              : sizeof(dxbc::RTS0::v3::StaticSampler);
+
+  StaticSamplers.Data = PartData.substr(StaticSamplersOffset,
+                                        static_cast<size_t>(NumStaticSamplers) *
+                                            StaticSamplers.Stride);
 
   return Error::success();
 }

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -343,6 +343,9 @@ Error DXContainerWriter::writeParts(raw_ostream &OS) {
         NewSampler.RegisterSpace = Param.RegisterSpace;
         NewSampler.ShaderVisibility = Param.ShaderVisibility;
 
+        if (RS.Version > 2)
+          NewSampler.Flags = Param.getEncodedFlags();
+
         RS.StaticSamplers.push_back(NewSampler);
       }