[HLSL] Adding support for root descriptors in root signature metadata representation

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -9,18 +9,81 @@
 #include "llvm/BinaryFormat/DXContainer.h"
 #include <cstdint>
 #include <limits>
+#include <variant>
 
 namespace llvm {
 
 class raw_ostream;
 namespace mcdxbc {
 
-struct RootParameter {
+struct RootParameterInfo {
   dxbc::RootParameterHeader Header;
-  union {
-    dxbc::RootConstants Constants;
-    dxbc::RTS0::v2::RootDescriptor Descriptor;
-  };
+  size_t Location;
+
+  RootParameterInfo() = default;
+
+  RootParameterInfo(dxbc::RootParameterHeader H, size_t L)
+      : Header(H), Location(L) {}
+};
+
+using RootDescriptor = std::variant<dxbc::RTS0::v1::RootDescriptor,
+                                    dxbc::RTS0::v2::RootDescriptor>;
+using ParametersView = std::variant<const dxbc::RootConstants *,
+                                    const dxbc::RTS0::v1::RootDescriptor *,
+                                    const dxbc::RTS0::v2::RootDescriptor *>;
+struct RootParametersContainer {
+  SmallVector<RootParameterInfo> ParametersInfo;
+
+  SmallVector<dxbc::RootConstants> Constants;
+  SmallVector<RootDescriptor> Descriptors;
+
+  void addInfo(dxbc::RootParameterHeader H, size_t L) {
+    ParametersInfo.push_back(RootParameterInfo(H, L));
+  }
+
+  void addParameter(dxbc::RootParameterHeader H, dxbc::RootConstants C) {
+    addInfo(H, Constants.size());
+    Constants.push_back(C);
+  }
+
+  void addParameter(dxbc::RootParameterHeader H,
+                    dxbc::RTS0::v1::RootDescriptor D) {
+    addInfo(H, Descriptors.size());
+    Descriptors.push_back(D);
+  }
+
+  void addParameter(dxbc::RootParameterHeader H,
+                    dxbc::RTS0::v2::RootDescriptor D) {
+    addInfo(H, Descriptors.size());
+    Descriptors.push_back(D);
+  }
+
+  std::optional<ParametersView> getParameter(const RootParameterInfo *H) const {
+    switch (H->Header.ParameterType) {
+    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
+      return &Constants[H->Location];
+    case llvm::to_underlying(dxbc::RootParameterType::CBV):
+    case llvm::to_underlying(dxbc::RootParameterType::SRV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV):
+      const RootDescriptor &VersionedParam = Descriptors[H->Location];
+      if (std::holds_alternative<dxbc::RTS0::v1::RootDescriptor>(
+              VersionedParam)) {
+        return &std::get<dxbc::RTS0::v1::RootDescriptor>(VersionedParam);
+      }
+      return &std::get<dxbc::RTS0::v2::RootDescriptor>(VersionedParam);
+    }
+
+    return std::nullopt;
+  }
+
+  size_t size() const { return ParametersInfo.size(); }
+
+  SmallVector<RootParameterInfo>::const_iterator begin() const {
+    return ParametersInfo.begin();
+  }
+  SmallVector<RootParameterInfo>::const_iterator end() const {
+    return ParametersInfo.end();
+  }
 };
 struct RootSignatureDesc {
 
@@ -29,7 +92,7 @@ struct RootSignatureDesc {
   uint32_t RootParameterOffset = 0U;
   uint32_t StaticSamplersOffset = 0u;
   uint32_t NumStaticSamplers = 0u;
-  SmallVector<mcdxbc::RootParameter> Parameters;
+  mcdxbc::RootParametersContainer ParametersContainer;
 
   void write(raw_ostream &OS) const;
 

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -94,7 +94,7 @@ struct RootParameterYamlDesc {
   uint32_t Type;
   uint32_t Visibility;
   uint32_t Offset;
-  RootParameterYamlDesc() {};
+  RootParameterYamlDesc(){};
   RootParameterYamlDesc(uint32_t T) : Type(T) {
     switch (T) {
 

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -30,24 +30,20 @@ static void rewriteOffsetToCurrentByte(raw_svector_ostream &Stream,
 
 size_t RootSignatureDesc::getSize() const {
   size_t Size = sizeof(dxbc::RootSignatureHeader) +
-                Parameters.size() * sizeof(dxbc::RootParameterHeader);
+                ParametersContainer.size() * sizeof(dxbc::RootParameterHeader);
 
-  for (const mcdxbc::RootParameter &P : Parameters) {
-    switch (P.Header.ParameterType) {
-    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;
-    }
+  for (const auto &I : ParametersContainer) {
+    std::optional<ParametersView> P = ParametersContainer.getParameter(&I);
+    if (!P)
+      continue;
+    std::visit(
+        [&Size](auto &Value) -> void {
+          using T = std::decay_t<decltype(*Value)>;
+          Size += sizeof(T);
+        },
+        *P);
   }
+
   return Size;
 }
 
@@ -56,7 +52,7 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
   raw_svector_ostream BOS(Storage);
   BOS.reserveExtraSpace(getSize());
 
-  const uint32_t NumParameters = Parameters.size();
+  const uint32_t NumParameters = ParametersContainer.size();
 
   support::endian::write(BOS, Version, llvm::endianness::little);
   support::endian::write(BOS, NumParameters, llvm::endianness::little);
@@ -66,7 +62,7 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
   support::endian::write(BOS, Flags, llvm::endianness::little);
 
   SmallVector<uint32_t> ParamsOffsets;
-  for (const mcdxbc::RootParameter &P : Parameters) {
+  for (const auto &P : ParametersContainer) {
     support::endian::write(BOS, P.Header.ParameterType,
                            llvm::endianness::little);
     support::endian::write(BOS, P.Header.ShaderVisibility,
@@ -76,29 +72,38 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
   }
 
   assert(NumParameters == ParamsOffsets.size());
-  for (size_t I = 0; I < NumParameters; ++I) {
+  const RootParameterInfo *H = ParametersContainer.begin();
+  for (size_t I = 0; I < NumParameters; ++I, H++) {
     rewriteOffsetToCurrentByte(BOS, ParamsOffsets[I]);
-    const mcdxbc::RootParameter &P = Parameters[I];
-
-    switch (P.Header.ParameterType) {
-    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
-      support::endian::write(BOS, P.Constants.ShaderRegister,
+    auto P = ParametersContainer.getParameter(H);
+    if (!P)
+      continue;
+    if (std::holds_alternative<const dxbc::RootConstants *>(P.value())) {
+      auto *Constants = std::get<const dxbc::RootConstants *>(P.value());
+      support::endian::write(BOS, Constants->ShaderRegister,
                              llvm::endianness::little);
-      support::endian::write(BOS, P.Constants.RegisterSpace,
+      support::endian::write(BOS, Constants->RegisterSpace,
                              llvm::endianness::little);
-      support::endian::write(BOS, P.Constants.Num32BitValues,
+      support::endian::write(BOS, 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,
+    } else if (std::holds_alternative<const dxbc::RTS0::v1::RootDescriptor *>(
+                   *P)) {
+      auto *Descriptor =
+          std::get<const dxbc::RTS0::v1::RootDescriptor *>(P.value());
+      support::endian::write(BOS, Descriptor->ShaderRegister,
+                             llvm::endianness::little);
+      support::endian::write(BOS, Descriptor->RegisterSpace,
+                             llvm::endianness::little);
+    } else if (std::holds_alternative<const dxbc::RTS0::v2::RootDescriptor *>(
+                   *P)) {
+      auto *Descriptor =
+          std::get<const dxbc::RTS0::v2::RootDescriptor *>(P.value());
+
+      support::endian::write(BOS, Descriptor->ShaderRegister,
                              llvm::endianness::little);
-      support::endian::write(BOS, P.Descriptor.RegisterSpace,
+      support::endian::write(BOS, Descriptor->RegisterSpace,
                              llvm::endianness::little);
-      if (Version > 1)
-        support::endian::write(BOS, P.Descriptor.Flags,
-                               llvm::endianness::little);
+      support::endian::write(BOS, Descriptor->Flags, llvm::endianness::little);
     }
   }
   assert(Storage.size() == getSize());

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -274,27 +274,39 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       RS.StaticSamplersOffset = P.RootSignature->StaticSamplersOffset;
 
       for (const auto &Param : P.RootSignature->Parameters) {
-        mcdxbc::RootParameter NewParam;
-        NewParam.Header = dxbc::RootParameterHeader{
-            Param.Type, Param.Visibility, Param.Offset};
+        auto Header = dxbc::RootParameterHeader{Param.Type, Param.Visibility,
+                                                Param.Offset};
 
         switch (Param.Type) {
         case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
-          NewParam.Constants.Num32BitValues = Param.Constants.Num32BitValues;
-          NewParam.Constants.RegisterSpace = Param.Constants.RegisterSpace;
-          NewParam.Constants.ShaderRegister = Param.Constants.ShaderRegister;
+          dxbc::RootConstants Constants;
+          Constants.Num32BitValues = Param.Constants.Num32BitValues;
+          Constants.RegisterSpace = Param.Constants.RegisterSpace;
+          Constants.ShaderRegister = Param.Constants.ShaderRegister;
+          RS.ParametersContainer.addParameter(Header, Constants);
           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();
+          if (RS.Version == 1) {
+            dxbc::RTS0::v1::RootDescriptor Descriptor;
+            Descriptor.RegisterSpace = Param.Descriptor.RegisterSpace;
+            Descriptor.ShaderRegister = Param.Descriptor.ShaderRegister;
+            RS.ParametersContainer.addParameter(Header, Descriptor);
+          } else {
+            dxbc::RTS0::v2::RootDescriptor Descriptor;
+            Descriptor.RegisterSpace = Param.Descriptor.RegisterSpace;
+            Descriptor.ShaderRegister = Param.Descriptor.ShaderRegister;
+            Descriptor.Flags = Param.Descriptor.getEncodedFlags();
+            RS.ParametersContainer.addParameter(Header, Descriptor);
+          }
           break;
+        default:
+          // Handling invalid parameter type edge case. We intentionally let
+          // obj2yaml/yaml2obj parse and emit invalid dxcontainer data, in order
+          // for that to be used as a testing tool more effectively.
+          RS.ParametersContainer.addInfo(Header, -1);
         }
-
-        RS.Parameters.push_back(NewParam);
       }
 
       RS.write(OS);