[DirectX] Adding support for Root Descriptors in obj2yaml/yaml2obj

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -158,6 +158,11 @@ enum class RootElementFlag : uint32_t {
 #include "DXContainerConstants.def"
 };
 
+#define ROOT_DESCRIPTOR_FLAG(Num, Val) Val = 1ull << Num,
+enum class RootDescriptorFlag : uint32_t {
+#include "DXContainerConstants.def"
+};
+
 #define ROOT_PARAMETER(Val, Enum) Enum = Val,
 enum class RootParameterType : uint32_t {
 #include "DXContainerConstants.def"
@@ -580,7 +585,33 @@ struct ProgramSignatureElement {
 
 static_assert(sizeof(ProgramSignatureElement) == 32,
               "ProgramSignatureElement is misaligned");
+namespace RTS0 {
+namespace v1 {
+struct RootDescriptor {
+  uint32_t ShaderRegister;
+  uint32_t RegisterSpace;
+  void swapBytes() {
+    sys::swapByteOrder(ShaderRegister);
+    sys::swapByteOrder(RegisterSpace);
+  }
+};
+} // namespace v1
+
+namespace v2 {
+struct RootDescriptor : public v1::RootDescriptor {
+  uint32_t Flags;
+
+  RootDescriptor() = default;
+  RootDescriptor(v1::RootDescriptor &Base)
+      : v1::RootDescriptor(Base), Flags(0u) {}
 
+  void swapBytes() {
+    v1::RootDescriptor::swapBytes();
+    sys::swapByteOrder(Flags);
+  }
+};
+} // namespace v2
+} // namespace RTS0
 // following dx12 naming
 // https://learn.microsoft.com/en-us/windows/win32/api/d3d12/ns-d3d12-d3d12_root_constants
 struct RootConstants {

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

@@ -72,9 +72,24 @@ ROOT_ELEMENT_FLAG(11, SamplerHeapDirectlyIndexed)
 #undef ROOT_ELEMENT_FLAG
 #endif // ROOT_ELEMENT_FLAG
 
+
+// ROOT_DESCRIPTOR_FLAG(bit offset for the flag, name).
+#ifdef ROOT_DESCRIPTOR_FLAG
+
+ROOT_DESCRIPTOR_FLAG(0, NONE)
+ROOT_DESCRIPTOR_FLAG(1, DATA_VOLATILE)
+ROOT_DESCRIPTOR_FLAG(2, DATA_STATIC_WHILE_SET_AT_EXECUTE)
+ROOT_DESCRIPTOR_FLAG(3, DATA_STATIC)
+#undef ROOT_DESCRIPTOR_FLAG
+#endif // ROOT_DESCRIPTOR_FLAG
+
+
 #ifdef ROOT_PARAMETER
 
 ROOT_PARAMETER(1, Constants32Bit)
+ROOT_PARAMETER(2, CBV)
+ROOT_PARAMETER(3, SRV)
+ROOT_PARAMETER(4, UAV)
 #undef ROOT_PARAMETER
 #endif // ROOT_PARAMETER
 

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -19,6 +19,7 @@ struct RootParameter {
   dxbc::RootParameterHeader Header;
   union {
     dxbc::RootConstants Constants;
+    dxbc::RTS0::v2::RootDescriptor Descriptor;
   };
 };
 struct RootSignatureDesc {

llvm/include/llvm/Object/DXContainer.h

@@ -17,13 +17,15 @@
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Twine.h"
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Object/Error.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/TargetParser/Triple.h"
 #include <array>
 #include <cstddef>
+#include <cstdint>
 #include <variant>
 
 namespace llvm {
@@ -121,6 +123,7 @@ namespace DirectX {
 struct RootParameterView {
   const dxbc::RootParameterHeader &Header;
   StringRef ParamData;
+
   RootParameterView(const dxbc::RootParameterHeader &H, StringRef P)
       : Header(H), ParamData(P) {}
 
@@ -149,6 +152,31 @@ struct RootConstantView : RootParameterView {
   }
 };
 
+struct RootDescriptorView : RootParameterView {
+  static bool classof(const RootParameterView *V) {
+    return (V->Header.ParameterType ==
+                llvm::to_underlying(dxbc::RootParameterType::CBV) ||
+            V->Header.ParameterType ==
+                llvm::to_underlying(dxbc::RootParameterType::SRV) ||
+            V->Header.ParameterType ==
+                llvm::to_underlying(dxbc::RootParameterType::UAV));
+  }
+
+  llvm::Expected<dxbc::RTS0::v2::RootDescriptor> read(uint32_t Version) {
+    if (Version == 1) {
+      auto Descriptor = readParameter<dxbc::RTS0::v1::RootDescriptor>();
+      if (Error E = Descriptor.takeError())
+        return E;
+      return dxbc::RTS0::v2::RootDescriptor(*Descriptor);
+    }
+    if (Version != 2)
+      return make_error<GenericBinaryError>("Invalid Root Signature version: " +
+                                                Twine(Version),
+                                            object_error::parse_failed);
+    return readParameter<dxbc::RTS0::v2::RootDescriptor>();
+  }
+};
+
 static Error parseFailed(const Twine &Msg) {
   return make_error<GenericBinaryError>(Msg.str(), object_error::parse_failed);
 }
@@ -192,6 +220,14 @@ class RootSignature {
     case dxbc::RootParameterType::Constants32Bit:
       DataSize = sizeof(dxbc::RootConstants);
       break;
+    case dxbc::RootParameterType::CBV:
+    case dxbc::RootParameterType::SRV:
+    case dxbc::RootParameterType::UAV:
+      if (Version == 1)
+        DataSize = sizeof(dxbc::RTS0::v1::RootDescriptor);
+      else
+        DataSize = sizeof(dxbc::RTS0::v2::RootDescriptor);
+      break;
     }
     size_t EndOfSectionByte = getNumStaticSamplers() == 0
                                   ? PartData.size()

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -21,7 +21,6 @@
 #include "llvm/ObjectYAML/YAML.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <array>
-#include <cstdint>
 #include <optional>
 #include <string>
 #include <vector>
@@ -73,21 +72,46 @@ struct ShaderHash {
   std::vector<llvm::yaml::Hex8> Digest;
 };
 
-#define ROOT_ELEMENT_FLAG(Num, Val) bool Val = false;
-
 struct RootConstantsYaml {
   uint32_t ShaderRegister;
   uint32_t RegisterSpace;
   uint32_t Num32BitValues;
 };
 
+struct RootDescriptorYaml {
+  RootDescriptorYaml() = default;
+
+  uint32_t ShaderRegister;
+  uint32_t RegisterSpace;
+
+  uint32_t getEncodedFlags() const;
+
+#define ROOT_DESCRIPTOR_FLAG(Num, Val) bool Val = false;
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
 struct RootParameterYamlDesc {
   uint32_t Type;
   uint32_t Visibility;
   uint32_t Offset;
+  RootParameterYamlDesc() {};
+  RootParameterYamlDesc(uint32_t T) : Type(T) {
+    switch (T) {
+
+    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
+      Constants = RootConstantsYaml();
+      break;
+    case llvm::to_underlying(dxbc::RootParameterType::CBV):
+    case llvm::to_underlying(dxbc::RootParameterType::SRV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV):
+      Descriptor = RootDescriptorYaml();
+      break;
+    }
+  }
 
   union {
     RootConstantsYaml Constants;
+    RootDescriptorYaml Descriptor;
   };
 };
 
@@ -111,6 +135,7 @@ struct RootSignatureYamlDesc {
   static llvm::Expected<DXContainerYAML::RootSignatureYamlDesc>
   create(const object::DirectX::RootSignature &Data);
 
+#define ROOT_ELEMENT_FLAG(Num, Val) bool Val = false;
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 };
 
@@ -298,6 +323,10 @@ template <> struct MappingTraits<llvm::DXContainerYAML::RootConstantsYaml> {
   static void mapping(IO &IO, llvm::DXContainerYAML::RootConstantsYaml &C);
 };
 
+template <> struct MappingTraits<llvm::DXContainerYAML::RootDescriptorYaml> {
+  static void mapping(IO &IO, llvm::DXContainerYAML::RootDescriptorYaml &D);
+};
+
 } // namespace yaml
 
 } // namespace llvm

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -37,6 +37,15 @@ size_t RootSignatureDesc::getSize() const {
     case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
       Size += sizeof(dxbc::RootConstants);
       break;
+    case llvm::to_underlying(dxbc::RootParameterType::CBV):
+    case llvm::to_underlying(dxbc::RootParameterType::SRV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV):
+      if (Version == 1)
+        Size += sizeof(dxbc::RTS0::v1::RootDescriptor);
+      else
+        Size += sizeof(dxbc::RTS0::v2::RootDescriptor);
+
+      break;
     }
   }
   return Size;
@@ -80,6 +89,16 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
       support::endian::write(BOS, P.Constants.Num32BitValues,
                              llvm::endianness::little);
       break;
+    case llvm::to_underlying(dxbc::RootParameterType::CBV):
+    case llvm::to_underlying(dxbc::RootParameterType::SRV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV):
+      support::endian::write(BOS, P.Descriptor.ShaderRegister,
+                             llvm::endianness::little);
+      support::endian::write(BOS, P.Descriptor.RegisterSpace,
+                             llvm::endianness::little);
+      if (Version > 1)
+        support::endian::write(BOS, P.Descriptor.Flags,
+                               llvm::endianness::little);
     }
   }
   assert(Storage.size() == getSize());

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -284,6 +284,14 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
           NewParam.Constants.RegisterSpace = Param.Constants.RegisterSpace;
           NewParam.Constants.ShaderRegister = Param.Constants.ShaderRegister;
           break;
+        case llvm::to_underlying(dxbc::RootParameterType::SRV):
+        case llvm::to_underlying(dxbc::RootParameterType::UAV):
+        case llvm::to_underlying(dxbc::RootParameterType::CBV):
+          NewParam.Descriptor.RegisterSpace = Param.Descriptor.RegisterSpace;
+          NewParam.Descriptor.ShaderRegister = Param.Descriptor.ShaderRegister;
+          if (P.RootSignature->Version > 1)
+            NewParam.Descriptor.Flags = Param.Descriptor.getEncodedFlags();
+          break;
         }
 
         RS.Parameters.push_back(NewParam);

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -38,8 +38,9 @@ DXContainerYAML::RootSignatureYamlDesc::create(
     const object::DirectX::RootSignature &Data) {
 
   RootSignatureYamlDesc RootSigDesc;
+  uint32_t Version = Data.getVersion();
 
-  RootSigDesc.Version = Data.getVersion();
+  RootSigDesc.Version = Version;
   RootSigDesc.NumStaticSamplers = Data.getNumStaticSamplers();
   RootSigDesc.StaticSamplersOffset = Data.getStaticSamplersOffset();
   RootSigDesc.NumRootParameters = Data.getNumRootParameters();
@@ -48,13 +49,12 @@ DXContainerYAML::RootSignatureYamlDesc::create(
   uint32_t Flags = Data.getFlags();
   for (const dxbc::RootParameterHeader &PH : Data.param_headers()) {
 
-    RootParameterYamlDesc NewP;
-    NewP.Offset = PH.ParameterOffset;
-
     if (!dxbc::isValidParameterType(PH.ParameterType))
       return createStringError(std::errc::invalid_argument,
                                "Invalid value for parameter type");
 
+    RootParameterYamlDesc NewP(PH.ParameterType);
+    NewP.Offset = PH.ParameterOffset;
     NewP.Type = PH.ParameterType;
 
     if (!dxbc::isValidShaderVisibility(PH.ShaderVisibility))
@@ -79,7 +79,24 @@ DXContainerYAML::RootSignatureYamlDesc::create(
       NewP.Constants.Num32BitValues = Constants.Num32BitValues;
       NewP.Constants.ShaderRegister = Constants.ShaderRegister;
       NewP.Constants.RegisterSpace = Constants.RegisterSpace;
+    } else if (auto *RDV =
+                   dyn_cast<object::DirectX::RootDescriptorView>(&ParamView)) {
+      llvm::Expected<dxbc::RTS0::v2::RootDescriptor> DescriptorOrErr =
+          RDV->read(Version);
+      if (Error E = DescriptorOrErr.takeError())
+        return std::move(E);
+      auto Descriptor = *DescriptorOrErr;
+      NewP.Descriptor.ShaderRegister = Descriptor.ShaderRegister;
+      NewP.Descriptor.RegisterSpace = Descriptor.RegisterSpace;
+      if (Version > 1) {
+#define ROOT_DESCRIPTOR_FLAG(Num, Val)                                         \
+  NewP.Descriptor.Val =                                                        \
+      (Descriptor.Flags &                                                      \
+       llvm::to_underlying(dxbc::RootDescriptorFlag::Val)) > 0;
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+      }
     }
+
     RootSigDesc.Parameters.push_back(NewP);
   }
 #define ROOT_ELEMENT_FLAG(Num, Val)                                            \
@@ -89,6 +106,15 @@ DXContainerYAML::RootSignatureYamlDesc::create(
   return RootSigDesc;
 }
 
+uint32_t DXContainerYAML::RootDescriptorYaml::getEncodedFlags() const {
+  uint64_t Flag = 0;
+#define ROOT_DESCRIPTOR_FLAG(Num, Val)                                         \
+  if (Val)                                                                     \
+    Flag |= (uint32_t)dxbc::RootDescriptorFlag::Val;
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+  return Flag;
+}
+
 uint32_t DXContainerYAML::RootSignatureYamlDesc::getEncodedFlags() {
   uint64_t Flag = 0;
 #define ROOT_ELEMENT_FLAG(Num, Val)                                            \
@@ -276,6 +302,14 @@ void MappingTraits<llvm::DXContainerYAML::RootConstantsYaml>::mapping(
   IO.mapRequired("ShaderRegister", C.ShaderRegister);
 }
 
+void MappingTraits<llvm::DXContainerYAML::RootDescriptorYaml>::mapping(
+    IO &IO, llvm::DXContainerYAML::RootDescriptorYaml &D) {
+  IO.mapRequired("RegisterSpace", D.RegisterSpace);
+  IO.mapRequired("ShaderRegister", D.ShaderRegister);
+#define ROOT_DESCRIPTOR_FLAG(Num, Val) IO.mapOptional(#Val, D.Val, false);
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+}
+
 void MappingTraits<llvm::DXContainerYAML::RootParameterYamlDesc>::mapping(
     IO &IO, llvm::DXContainerYAML::RootParameterYamlDesc &P) {
   IO.mapRequired("ParameterType", P.Type);
@@ -285,6 +319,11 @@ void MappingTraits<llvm::DXContainerYAML::RootParameterYamlDesc>::mapping(
   case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
     IO.mapRequired("Constants", P.Constants);
     break;
+  case llvm::to_underlying(dxbc::RootParameterType::CBV):
+  case llvm::to_underlying(dxbc::RootParameterType::SRV):
+  case llvm::to_underlying(dxbc::RootParameterType::UAV):
+    IO.mapRequired("Descriptor", P.Descriptor);
+    break;
   }
 }