[DXIL] Adding support to RootSignatureFlags in obj2yaml (#122396)

This PR adds:
- `RootSignatureFlags` extraction from DXContainer using `obj2yaml`


This PR is part of: #121493

---------

Co-authored-by: joaosaffran <[email protected]>

Author: joaosaffran

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -14,6 +14,8 @@
 #define LLVM_BINARYFORMAT_DXCONTAINER_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/BinaryStreamError.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/SwapByteOrder.h"
 #include "llvm/TargetParser/Triple.h"
 
@@ -152,6 +154,11 @@ enum class FeatureFlags : uint64_t {
 static_assert((uint64_t)FeatureFlags::NextUnusedBit <= 1ull << 63,
               "Shader flag bits exceed enum size.");
 
+#define ROOT_ELEMENT_FLAG(Num, Val) Val = 1ull << Num,
+enum class RootElementFlag : uint32_t {
+#include "DXContainerConstants.def"
+};
+
 PartType parsePartType(StringRef S);
 
 struct VertexPSVInfo {
@@ -541,6 +548,23 @@ struct ProgramSignatureElement {
 static_assert(sizeof(ProgramSignatureElement) == 32,
               "ProgramSignatureElement is misaligned");
 
+struct RootSignatureValidations {
+
+  static Expected<uint32_t> validateRootFlag(uint32_t Flags) {
+    if ((Flags & ~0x80000fff) != 0)
+      return llvm::make_error<BinaryStreamError>("Invalid Root Signature flag");
+    return Flags;
+  }
+
+  static Expected<uint32_t> validateVersion(uint32_t Version) {
+    if (Version == 1 || Version == 2)
+      return Version;
+
+    return llvm::make_error<BinaryStreamError>(
+        "Invalid Root Signature Version");
+  }
+};
+
 } // namespace dxbc
 } // namespace llvm
 

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

@@ -4,6 +4,7 @@ CONTAINER_PART(DXIL)
 CONTAINER_PART(SFI0)
 CONTAINER_PART(HASH)
 CONTAINER_PART(PSV0)
+CONTAINER_PART(RTS0)
 CONTAINER_PART(ISG1)
 CONTAINER_PART(OSG1)
 CONTAINER_PART(PSG1)
@@ -52,6 +53,26 @@ SHADER_FEATURE_FLAG(31, 36, NextUnusedBit, "Next reserved shader flag bit (not a
 #undef SHADER_FEATURE_FLAG
 #endif // SHADER_FEATURE_FLAG
 
+
+// ROOT_ELEMENT_FLAG(bit offset for the flag, name).
+#ifdef ROOT_ELEMENT_FLAG
+
+ROOT_ELEMENT_FLAG(0, AllowInputAssemblerInputLayout)
+ROOT_ELEMENT_FLAG(1, DenyVertexShaderRootAccess)
+ROOT_ELEMENT_FLAG(2, DenyHullShaderRootAccess)
+ROOT_ELEMENT_FLAG(3, DenyDomainShaderRootAccess)
+ROOT_ELEMENT_FLAG(4, DenyGeometryShaderRootAccess)
+ROOT_ELEMENT_FLAG(5, DenyPixelShaderRootAccess)
+ROOT_ELEMENT_FLAG(6, AllowStreamOutput)
+ROOT_ELEMENT_FLAG(7, LocalRootSignature)
+ROOT_ELEMENT_FLAG(8, DenyAmplificationShaderRootAccess)
+ROOT_ELEMENT_FLAG(9, DenyMeshShaderRootAccess)
+ROOT_ELEMENT_FLAG(10, CBVSRVUAVHeapDirectlyIndexed)
+ROOT_ELEMENT_FLAG(11, SamplerHeapDirectlyIndexed)
+#undef ROOT_ELEMENT_FLAG
+#endif // ROOT_ELEMENT_FLAG
+
+
 #ifdef DXIL_MODULE_FLAG
 
 // Only save DXIL module flags which not map to feature flags here.

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -0,0 +1,28 @@
+//===- llvm/MC/DXContainerRootSignature.h - RootSignature -*- C++ -*- ========//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <cstdint>
+#include <limits>
+
+namespace llvm {
+
+class raw_ostream;
+
+namespace mcdxbc {
+struct RootSignatureHeader {
+  uint32_t Version = 2;
+  uint32_t NumParameters = 0;
+  uint32_t RootParametersOffset = 0;
+  uint32_t NumStaticSamplers = 0;
+  uint32_t StaticSamplersOffset = 0;
+  uint32_t Flags = 0;
+
+  void write(raw_ostream &OS);
+};
+} // namespace mcdxbc
+} // namespace llvm

llvm/include/llvm/Object/DXContainer.h

@@ -116,6 +116,28 @@ template <typename T> struct ViewArray {
 };
 
 namespace DirectX {
+
+class RootSignature {
+private:
+  uint32_t Version;
+  uint32_t NumParameters;
+  uint32_t RootParametersOffset;
+  uint32_t NumStaticSamplers;
+  uint32_t StaticSamplersOffset;
+  uint32_t Flags;
+
+public:
+  RootSignature() {}
+
+  Error parse(StringRef Data);
+  uint32_t getVersion() const { return Version; }
+  uint32_t getNumParameters() const { return NumParameters; }
+  uint32_t getRootParametersOffset() const { return RootParametersOffset; }
+  uint32_t getNumStaticSamplers() const { return NumStaticSamplers; }
+  uint32_t getStaticSamplersOffset() const { return StaticSamplersOffset; }
+  uint32_t getFlags() const { return Flags; }
+};
+
 class PSVRuntimeInfo {
 
   using ResourceArray = ViewArray<dxbc::PSV::v2::ResourceBindInfo>;
@@ -287,6 +309,7 @@ class DXContainer {
   std::optional<uint64_t> ShaderFeatureFlags;
   std::optional<dxbc::ShaderHash> Hash;
   std::optional<DirectX::PSVRuntimeInfo> PSVInfo;
+  std::optional<DirectX::RootSignature> RootSignature;
   DirectX::Signature InputSignature;
   DirectX::Signature OutputSignature;
   DirectX::Signature PatchConstantSignature;
@@ -296,6 +319,7 @@ class DXContainer {
   Error parseDXILHeader(StringRef Part);
   Error parseShaderFeatureFlags(StringRef Part);
   Error parseHash(StringRef Part);
+  Error parseRootSignature(StringRef Part);
   Error parsePSVInfo(StringRef Part);
   Error parseSignature(StringRef Part, DirectX::Signature &Array);
   friend class PartIterator;
@@ -382,6 +406,10 @@ class DXContainer {
 
   std::optional<dxbc::ShaderHash> getShaderHash() const { return Hash; }
 
+  std::optional<DirectX::RootSignature> getRootSignature() const {
+    return RootSignature;
+  }
+
   const std::optional<DirectX::PSVRuntimeInfo> &getPSVInfo() const {
     return PSVInfo;
   };

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -17,6 +17,7 @@
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/Object/DXContainer.h"
 #include "llvm/ObjectYAML/YAML.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <array>
@@ -72,6 +73,22 @@ struct ShaderHash {
   std::vector<llvm::yaml::Hex8> Digest;
 };
 
+#define ROOT_ELEMENT_FLAG(Num, Val) bool Val = false;
+struct RootSignatureDesc {
+  RootSignatureDesc() = default;
+  RootSignatureDesc(const object::DirectX::RootSignature &Data);
+
+  uint32_t Version;
+  uint32_t NumParameters;
+  uint32_t RootParametersOffset;
+  uint32_t NumStaticSamplers;
+  uint32_t StaticSamplersOffset;
+
+  uint32_t getEncodedFlags();
+
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
 using ResourceFlags = dxbc::PSV::ResourceFlags;
 using ResourceBindInfo = dxbc::PSV::v2::ResourceBindInfo;
 
@@ -159,6 +176,7 @@ struct Part {
   std::optional<ShaderHash> Hash;
   std::optional<PSVInfo> Info;
   std::optional<DXContainerYAML::Signature> Signature;
+  std::optional<DXContainerYAML::RootSignatureDesc> RootSignature;
 };
 
 struct Object {
@@ -241,6 +259,11 @@ template <> struct MappingTraits<DXContainerYAML::Signature> {
   static void mapping(IO &IO, llvm::DXContainerYAML::Signature &El);
 };
 
+template <> struct MappingTraits<DXContainerYAML::RootSignatureDesc> {
+  static void mapping(IO &IO,
+                      DXContainerYAML::RootSignatureDesc &RootSignature);
+};
+
 } // namespace yaml
 
 } // namespace llvm

llvm/lib/MC/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_llvm_component_library(LLVMMC
   ConstantPools.cpp
   DXContainerPSVInfo.cpp
+  DXContainerRootSignature.cpp
   ELFObjectWriter.cpp
   GOFFObjectWriter.cpp
   MCAsmBackend.cpp

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -0,0 +1,23 @@
+//===- llvm/MC/DXContainerRootSignature.cpp - RootSignature -*- C++ -*-=======//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/DXContainerRootSignature.h"
+#include "llvm/Support/EndianStream.h"
+
+using namespace llvm;
+using namespace llvm::mcdxbc;
+
+void RootSignatureHeader::write(raw_ostream &OS) {
+
+  support::endian::write(OS, Version, llvm::endianness::little);
+  support::endian::write(OS, NumParameters, llvm::endianness::little);
+  support::endian::write(OS, RootParametersOffset, llvm::endianness::little);
+  support::endian::write(OS, NumStaticSamplers, llvm::endianness::little);
+  support::endian::write(OS, StaticSamplersOffset, llvm::endianness::little);
+  support::endian::write(OS, Flags, llvm::endianness::little);
+}

llvm/lib/Object/DXContainer.cpp

@@ -10,6 +10,7 @@
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Object/Error.h"
 #include "llvm/Support/Alignment.h"
+#include "llvm/Support/Endian.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace llvm;
@@ -92,6 +93,15 @@ Error DXContainer::parseHash(StringRef Part) {
   return Error::success();
 }
 
+Error DXContainer::parseRootSignature(StringRef Part) {
+  if (RootSignature)
+    return parseFailed("More than one RTS0 part is present in the file");
+  RootSignature = DirectX::RootSignature();
+  if (Error Err = RootSignature->parse(Part))
+    return Err;
+  return Error::success();
+}
+
 Error DXContainer::parsePSVInfo(StringRef Part) {
   if (PSVInfo)
     return parseFailed("More than one PSV0 part is present in the file");
@@ -193,6 +203,10 @@ Error DXContainer::parsePartOffsets() {
       break;
     case dxbc::PartType::Unknown:
       break;
+    case dxbc::PartType::RTS0:
+      if (Error Err = parseRootSignature(PartData))
+        return Err;
+      break;
     }
   }
 
@@ -228,6 +242,53 @@ void DXContainer::PartIterator::updateIteratorImpl(const uint32_t Offset) {
   IteratorState.Offset = Offset;
 }
 
+Error DirectX::RootSignature::parse(StringRef Data) {
+  const char *Current = Data.begin();
+
+  // Root Signature headers expects 6 integers to be present.
+  if (Data.size() < 6 * sizeof(uint32_t))
+    return parseFailed(
+        "Invalid root signature, insufficient space for header.");
+
+  uint32_t VValue =
+      support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  Current += sizeof(uint32_t);
+
+  Expected<uint32_t> MaybeVersion =
+      dxbc::RootSignatureValidations::validateVersion(VValue);
+  if (Error E = MaybeVersion.takeError())
+    return E;
+  Version = MaybeVersion.get();
+
+  NumParameters =
+      support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  Current += sizeof(uint32_t);
+
+  RootParametersOffset =
+      support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  Current += sizeof(uint32_t);
+
+  NumStaticSamplers =
+      support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  Current += sizeof(uint32_t);
+
+  StaticSamplersOffset =
+      support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  Current += sizeof(uint32_t);
+
+  uint32_t FValue =
+      support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  Current += sizeof(uint32_t);
+
+  Expected<uint32_t> MaybeFlag =
+      dxbc::RootSignatureValidations::validateRootFlag(FValue);
+  if (Error E = MaybeFlag.takeError())
+    return E;
+  Flags = MaybeFlag.get();
+
+  return Error::success();
+}
+
 Error DirectX::PSVRuntimeInfo::parse(uint16_t ShaderKind) {
   Triple::EnvironmentType ShaderStage = dxbc::getShaderStage(ShaderKind);
 

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -13,6 +13,7 @@
 
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/MC/DXContainerPSVInfo.h"
+#include "llvm/MC/DXContainerRootSignature.h"
 #include "llvm/ObjectYAML/ObjectYAML.h"
 #include "llvm/ObjectYAML/yaml2obj.h"
 #include "llvm/Support/Errc.h"
@@ -261,6 +262,20 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
     }
     case dxbc::PartType::Unknown:
       break; // Skip any handling for unrecognized parts.
+    case dxbc::PartType::RTS0:
+      if (!P.RootSignature.has_value())
+        continue;
+
+      mcdxbc::RootSignatureHeader Header;
+      Header.Flags = P.RootSignature->getEncodedFlags();
+      Header.Version = P.RootSignature->Version;
+      Header.NumParameters = P.RootSignature->NumParameters;
+      Header.RootParametersOffset = P.RootSignature->RootParametersOffset;
+      Header.NumStaticSamplers = P.RootSignature->NumStaticSamplers;
+      Header.StaticSamplersOffset = P.RootSignature->StaticSamplersOffset;
+
+      Header.write(OS);
+      break;
     }
     uint64_t BytesWritten = OS.tell() - DataStart;
     RollingOffset += BytesWritten;