[DirectX] Updating DXContainer Yaml to represent Root Signature 1.2 (#159659)

This PR updates the YAML representation of DXContainer to support Root
Signature 1.2, this also requires updating the write logic to support
testing.

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,22 @@ 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/RootSignatureValidations.h

@@ -33,6 +33,7 @@ LLVM_ABI bool verifyRangeType(uint32_t Type);
 LLVM_ABI bool verifyDescriptorRangeFlag(uint32_t Version,
                                         dxil::ResourceClass Type,
                                         dxbc::DescriptorRangeFlags FlagsVal);
+LLVM_ABI bool verifyStaticSamplerFlags(uint32_t Version, uint32_t FlagsNumber);
 LLVM_ABI bool verifyNumDescriptors(uint32_t NumDescriptors);
 LLVM_ABI bool verifyMipLODBias(float MipLODBias);
 LLVM_ABI bool verifyMaxAnisotropy(uint32_t MaxAnisotropy);

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -74,6 +74,8 @@ struct StaticSampler {
   uint32_t ShaderRegister;
   uint32_t RegisterSpace;
   dxbc::ShaderVisibility ShaderVisibility;
+  // Version 3 onwards:
+  uint32_t Flags = 0;
 };
 
 struct RootParametersContainer {

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/RootSignatureMetadata.cpp

@@ -719,6 +719,12 @@ Error MetadataParser::validateRootSignature(
           joinErrors(std::move(DeferredErrs),
                      make_error<RootSignatureValidationError<uint32_t>>(
                          "RegisterSpace", Sampler.RegisterSpace));
+
+    if (!hlsl::rootsig::verifyStaticSamplerFlags(RSD.Version, Sampler.Flags))
+      DeferredErrs =
+          joinErrors(std::move(DeferredErrs),
+                     make_error<RootSignatureValidationError<uint32_t>>(
+                         "Static Sampler Flag", Sampler.Flags));
   }
 
   return DeferredErrs;

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;
@@ -111,6 +113,25 @@ bool verifyDescriptorRangeFlag(uint32_t Version, dxil::ResourceClass Type,
   return (Flags & ~Mask) == FlagT::None;
 }
 
+bool verifyStaticSamplerFlags(uint32_t Version, uint32_t FlagsNumber) {
+  uint32_t LargestValue = llvm::to_underlying(
+      dxbc::StaticSamplerFlags::LLVM_BITMASK_LARGEST_ENUMERATOR);
+  if (FlagsNumber >= NextPowerOf2(LargestValue))
+    return false;
+
+  dxbc::StaticSamplerFlags Flags = dxbc::StaticSamplerFlags(FlagsNumber);
+  if (Version <= 2)
+    return Flags == dxbc::StaticSamplerFlags::None;
+
+  assert(Version == 3 && "Provided invalid root signature version");
+
+  dxbc::StaticSamplerFlags Mask =
+      dxbc::StaticSamplerFlags::NonNormalizedCoordinates |
+      dxbc::StaticSamplerFlags::UintBorderColor |
+      dxbc::StaticSamplerFlags::None;
+  return (Flags | Mask) == Mask;
+}
+
 bool verifyNumDescriptors(uint32_t NumDescriptors) {
   return NumDescriptors > 0;
 }

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -32,10 +32,12 @@ 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);
+  size_t StaticSamplersSize = sizeof(dxbc::RTS0::v1::StaticSampler);
+  if (Version > 2)
+    StaticSamplersSize = sizeof(dxbc::RTS0::v3::StaticSampler);
 
-  return size_t(StaticSamplersOffset) + StaticSamplersSize;
+  return size_t(StaticSamplersOffset) +
+         (StaticSamplersSize * StaticSamplers.size());
 }
 
 uint32_t RootSignatureDesc::computeRootParametersOffset() const {
@@ -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/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);
       }
 

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -245,6 +245,15 @@ uint32_t DXContainerYAML::DescriptorRangeYaml::getEncodedFlags() const {
   return Flags;
 }
 
+uint32_t DXContainerYAML::StaticSamplerYamlDesc::getEncodedFlags() const {
+  uint64_t Flags = 0;
+#define STATIC_SAMPLER_FLAG(Num, Enum, Flag)                                   \
+  if (Enum)                                                                    \
+    Flags |= (uint32_t)dxbc::StaticSamplerFlags::Enum;
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+  return Flags;
+}
+
 uint64_t DXContainerYAML::ShaderFeatureFlags::getEncodedFlags() {
   uint64_t Flag = 0;
 #define SHADER_FEATURE_FLAG(Num, DxilModuleNum, Val, Str)                      \
@@ -512,6 +521,9 @@ void MappingTraits<llvm::DXContainerYAML::StaticSamplerYamlDesc>::mapping(
   IO.mapRequired("ShaderRegister", S.ShaderRegister);
   IO.mapRequired("RegisterSpace", S.RegisterSpace);
   IO.mapRequired("ShaderVisibility", S.ShaderVisibility);
+#define STATIC_SAMPLER_FLAG(Num, Enum, Flag)                                   \
+  IO.mapOptional(#Flag, S.Enum, false);
+#include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
 void MappingTraits<DXContainerYAML::Part>::mapping(IO &IO,

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Invalid-Version.ll

@@ -0,0 +1,20 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+
+; CHECK: error: Invalid value for Version: 4
+; CHECK-NOT: Root Signature Definitions
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2, !3, !4, !5} ; list of function/root signature pairs
+!2 = !{ ptr @main, !6, i32 1 } ; function, root signature
+!3 = !{ ptr @main, !6, i32 4 } ; function, root signature
+!4 = !{ ptr @main, !6, i32 2 } ; function, root signature
+!5 = !{ ptr @main, !6, i32 3 } ; function, root signature
+!6 = !{ } ; list of root signature elements

llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp

@@ -526,3 +526,54 @@ Parts:
   EXPECT_EQ(Storage.size(), 144u);
   EXPECT_TRUE(memcmp(Buffer, Storage.data(), 144u) == 0);
 }
+
+TEST(RootSignature, ParseStaticSamplersV13) {
+  SmallString<128> Storage;
+
+  // First read a fully explicit yaml with all sizes and offsets provided
+  ASSERT_TRUE(convert(Storage, R"(--- !dxcontainer
+Header:
+  Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 
+                     0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:           1
+    Minor:           0
+  PartCount:       1
+  PartOffsets:     [ 60 ]
+Parts:
+  - Name:            RTS0
+    Size:            76
+    RootSignature:
+      Version: 3
+      NumRootParameters: 0
+      RootParametersOffset: 24
+      NumStaticSamplers: 1
+      StaticSamplersOffset: 24
+      Parameters: []
+      Samplers: 
+        - ShaderRegister: 31 
+          RegisterSpace: 32
+          ShaderVisibility:  All
+          SAMPLER_FLAG_UINT_BORDER_COLOR: true
+      AllowInputAssemblerInputLayout: true
+      DenyGeometryShaderRootAccess: true
+    )"));
+
+  uint8_t Buffer[] = {
+      0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x90, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x3c, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+      0x52, 0x54, 0x53, 0x30, 0x4c, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x18, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0x55, 0x00, 0x00, 0x00,
+      0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00,
+      0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0x7f, 0x7f,
+      0x1f, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+      0x01, 0x00, 0x00, 0x00};
+
+  EXPECT_EQ(Storage.size(), 148U);
+  EXPECT_TRUE(memcmp(Buffer, Storage.data(), 148U) == 0);
+}