[NFC] Updating RTS0 namespace to contain all elements related to it's representation

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -585,6 +585,7 @@ struct ProgramSignatureElement {
 
 static_assert(sizeof(ProgramSignatureElement) == 32,
               "ProgramSignatureElement is misaligned");
+
 namespace RTS0 {
 namespace v1 {
 struct RootDescriptor {
@@ -595,23 +596,7 @@ struct RootDescriptor {
     sys::swapByteOrder(RegisterSpace);
   }
 };
-} // namespace v1
-
-namespace v2 {
-struct RootDescriptor : public v1::RootDescriptor {
-  uint32_t Flags;
-
-  RootDescriptor() = default;
-  explicit 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 {
@@ -655,6 +640,24 @@ struct RootSignatureHeader {
     sys::swapByteOrder(Flags);
   }
 };
+} // namespace v1
+
+namespace v2 {
+struct RootDescriptor : public v1::RootDescriptor {
+  uint32_t Flags;
+
+  RootDescriptor() = default;
+  explicit RootDescriptor(v1::RootDescriptor &Base)
+      : v1::RootDescriptor(Base), Flags(0u) {}
+
+  void swapBytes() {
+    v1::RootDescriptor::swapBytes();
+    sys::swapByteOrder(Flags);
+  }
+};
+} // namespace v2
+} // namespace RTS0
+
 } // namespace dxbc
 } // namespace llvm
 

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -16,36 +16,36 @@ class raw_ostream;
 namespace mcdxbc {
 
 struct RootParameterInfo {
-  dxbc::RootParameterHeader Header;
+  dxbc::RTS0::v1::RootParameterHeader Header;
   size_t Location;
 
   RootParameterInfo() = default;
 
-  RootParameterInfo(dxbc::RootParameterHeader Header, size_t Location)
+  RootParameterInfo(dxbc::RTS0::v1::RootParameterHeader Header, size_t Location)
       : Header(Header), Location(Location) {}
 };
 
 struct RootParametersContainer {
   SmallVector<RootParameterInfo> ParametersInfo;
 
-  SmallVector<dxbc::RootConstants> Constants;
+  SmallVector<dxbc::RTS0::v1::RootConstants> Constants;
   SmallVector<dxbc::RTS0::v2::RootDescriptor> Descriptors;
 
-  void addInfo(dxbc::RootParameterHeader Header, size_t Location) {
+  void addInfo(dxbc::RTS0::v1::RootParameterHeader Header, size_t Location) {
     ParametersInfo.push_back(RootParameterInfo(Header, Location));
   }
 
-  void addParameter(dxbc::RootParameterHeader Header,
-                    dxbc::RootConstants Constant) {
+  void addParameter(dxbc::RTS0::v1::RootParameterHeader Header,
+                    dxbc::RTS0::v1::RootConstants Constant) {
     addInfo(Header, Constants.size());
     Constants.push_back(Constant);
   }
 
-  void addInvalidParameter(dxbc::RootParameterHeader Header) {
+  void addInvalidParameter(dxbc::RTS0::v1::RootParameterHeader Header) {
     addInfo(Header, -1);
   }
 
-  void addParameter(dxbc::RootParameterHeader Header,
+  void addParameter(dxbc::RTS0::v1::RootParameterHeader Header,
                     dxbc::RTS0::v2::RootDescriptor Descriptor) {
     addInfo(Header, Descriptors.size());
     Descriptors.push_back(Descriptor);
@@ -57,12 +57,12 @@ struct RootParametersContainer {
     return {Info.Header.ParameterType, Info.Location};
   }
 
-  const dxbc::RootParameterHeader &getHeader(size_t Location) const {
+  const dxbc::RTS0::v1::RootParameterHeader &getHeader(size_t Location) const {
     const RootParameterInfo &Info = ParametersInfo[Location];
     return Info.Header;
   }
 
-  const dxbc::RootConstants &getConstant(size_t Index) const {
+  const dxbc::RTS0::v1::RootConstants &getConstant(size_t Index) const {
     return Constants[Index];
   }
 

llvm/include/llvm/Object/DXContainer.h

@@ -121,10 +121,10 @@ template <typename T> struct ViewArray {
 
 namespace DirectX {
 struct RootParameterView {
-  const dxbc::RootParameterHeader &Header;
+  const dxbc::RTS0::v1::RootParameterHeader &Header;
   StringRef ParamData;
 
-  RootParameterView(const dxbc::RootParameterHeader &H, StringRef P)
+  RootParameterView(const dxbc::RTS0::v1::RootParameterHeader &H, StringRef P)
       : Header(H), ParamData(P) {}
 
   template <typename T> Expected<T> readParameter() {
@@ -147,8 +147,8 @@ struct RootConstantView : RootParameterView {
            (uint32_t)dxbc::RootParameterType::Constants32Bit;
   }
 
-  llvm::Expected<dxbc::RootConstants> read() {
-    return readParameter<dxbc::RootConstants>();
+  llvm::Expected<dxbc::RTS0::v1::RootConstants> read() {
+    return readParameter<dxbc::RTS0::v1::RootConstants>();
   }
 };
 
@@ -189,10 +189,11 @@ class RootSignature {
   uint32_t NumStaticSamplers;
   uint32_t StaticSamplersOffset;
   uint32_t Flags;
-  ViewArray<dxbc::RootParameterHeader> ParametersHeaders;
+  ViewArray<dxbc::RTS0::v1::RootParameterHeader> ParametersHeaders;
   StringRef PartData;
 
-  using param_header_iterator = ViewArray<dxbc::RootParameterHeader>::iterator;
+  using param_header_iterator =
+      ViewArray<dxbc::RTS0::v1::RootParameterHeader>::iterator;
 
 public:
   RootSignature(StringRef PD) : PartData(PD) {}
@@ -210,15 +211,15 @@ class RootSignature {
   uint32_t getFlags() const { return Flags; }
 
   llvm::Expected<RootParameterView>
-  getParameter(const dxbc::RootParameterHeader &Header) const {
+  getParameter(const dxbc::RTS0::v1::RootParameterHeader &Header) const {
     size_t DataSize;
 
     if (!dxbc::isValidParameterType(Header.ParameterType))
       return parseFailed("invalid parameter type");
 
     switch (static_cast<dxbc::RootParameterType>(Header.ParameterType)) {
     case dxbc::RootParameterType::Constants32Bit:
-      DataSize = sizeof(dxbc::RootConstants);
+      DataSize = sizeof(dxbc::RTS0::v1::RootConstants);
       break;
     case dxbc::RootParameterType::CBV:
     case dxbc::RootParameterType::SRV:

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -90,29 +90,16 @@ struct RootDescriptorYaml {
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 };
 
+using ParameterData = std::variant<RootConstantsYaml, RootDescriptorYaml>;
+
 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;
-    }
-  }
+  ParameterData Data;
 
-  union {
-    RootConstantsYaml Constants;
-    RootDescriptorYaml Descriptor;
-  };
+  RootParameterYamlDesc() {};
+  RootParameterYamlDesc(uint32_t T) : Type(T) {}
 };
 
 struct RootSignatureYamlDesc {

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -29,13 +29,14 @@ static void rewriteOffsetToCurrentByte(raw_svector_ostream &Stream,
 }
 
 size_t RootSignatureDesc::getSize() const {
-  size_t Size = sizeof(dxbc::RootSignatureHeader) +
-                ParametersContainer.size() * sizeof(dxbc::RootParameterHeader);
+  size_t Size =
+      sizeof(dxbc::RTS0::v1::RootSignatureHeader) +
+      ParametersContainer.size() * sizeof(dxbc::RTS0::v1::RootParameterHeader);
 
   for (const RootParameterInfo &I : ParametersContainer) {
     switch (I.Header.ParameterType) {
     case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
-      Size += sizeof(dxbc::RootConstants);
+      Size += sizeof(dxbc::RTS0::v1::RootConstants);
       break;
     case llvm::to_underlying(dxbc::RootParameterType::CBV):
     case llvm::to_underlying(dxbc::RootParameterType::SRV):
@@ -81,7 +82,7 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
     const auto &[Type, Loc] = ParametersContainer.getTypeAndLocForParameter(I);
     switch (Type) {
     case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
-      const dxbc::RootConstants &Constants =
+      const dxbc::RTS0::v1::RootConstants &Constants =
           ParametersContainer.getConstant(Loc);
       support::endian::write(BOS, Constants.ShaderRegister,
                              llvm::endianness::little);

llvm/lib/Object/DXContainer.cpp

@@ -273,7 +273,8 @@ Error DirectX::RootSignature::parse() {
   Current += sizeof(uint32_t);
 
   ParametersHeaders.Data = PartData.substr(
-      RootParametersOffset, NumParameters * sizeof(dxbc::RootParameterHeader));
+      RootParametersOffset,
+      NumParameters * sizeof(dxbc::RTS0::v1::RootParameterHeader));
 
   return Error::success();
 }

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -274,28 +274,26 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       RS.StaticSamplersOffset = P.RootSignature->StaticSamplersOffset;
 
       for (const auto &Param : P.RootSignature->Parameters) {
-        dxbc::RootParameterHeader Header{Param.Type, Param.Visibility,
-                                         Param.Offset};
-
-        switch (Param.Type) {
-        case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
-          dxbc::RootConstants Constants;
-          Constants.Num32BitValues = Param.Constants.Num32BitValues;
-          Constants.RegisterSpace = Param.Constants.RegisterSpace;
-          Constants.ShaderRegister = Param.Constants.ShaderRegister;
+        auto Header = dxbc::RTS0::v1::RootParameterHeader{
+            Param.Type, Param.Visibility, Param.Offset};
+
+        if (auto *ConstantYaml =
+                std::get_if<DXContainerYAML::RootConstantsYaml>(&Param.Data)) {
+          dxbc::RTS0::v1::RootConstants Constants;
+          Constants.Num32BitValues = ConstantYaml->Num32BitValues;
+          Constants.RegisterSpace = ConstantYaml->RegisterSpace;
+          Constants.ShaderRegister = ConstantYaml->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):
+        } else if (auto *DescriptorYaml =
+                       std::get_if<DXContainerYAML::RootDescriptorYaml>(
+                           &Param.Data)) {
           dxbc::RTS0::v2::RootDescriptor Descriptor;
-          Descriptor.RegisterSpace = Param.Descriptor.RegisterSpace;
-          Descriptor.ShaderRegister = Param.Descriptor.ShaderRegister;
+          Descriptor.RegisterSpace = DescriptorYaml->RegisterSpace;
+          Descriptor.ShaderRegister = DescriptorYaml->ShaderRegister;
           if (RS.Version > 1)
-            Descriptor.Flags = Param.Descriptor.getEncodedFlags();
+            Descriptor.Flags = DescriptorYaml->getEncodedFlags();
           RS.ParametersContainer.addParameter(Header, Descriptor);
-          break;
-        default:
+        } else {
           // 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.

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -47,7 +47,7 @@ DXContainerYAML::RootSignatureYamlDesc::create(
   RootSigDesc.RootParametersOffset = Data.getRootParametersOffset();
 
   uint32_t Flags = Data.getFlags();
-  for (const dxbc::RootParameterHeader &PH : Data.param_headers()) {
+  for (const dxbc::RTS0::v1::RootParameterHeader &PH : Data.param_headers()) {
 
     if (!dxbc::isValidParameterType(PH.ParameterType))
       return createStringError(std::errc::invalid_argument,
@@ -70,31 +70,35 @@ DXContainerYAML::RootSignatureYamlDesc::create(
     object::DirectX::RootParameterView ParamView = ParamViewOrErr.get();
 
     if (auto *RCV = dyn_cast<object::DirectX::RootConstantView>(&ParamView)) {
-      llvm::Expected<dxbc::RootConstants> ConstantsOrErr = RCV->read();
+      llvm::Expected<dxbc::RTS0::v1::RootConstants> ConstantsOrErr =
+          RCV->read();
       if (Error E = ConstantsOrErr.takeError())
         return std::move(E);
 
       auto Constants = *ConstantsOrErr;
-
-      NewP.Constants.Num32BitValues = Constants.Num32BitValues;
-      NewP.Constants.ShaderRegister = Constants.ShaderRegister;
-      NewP.Constants.RegisterSpace = Constants.RegisterSpace;
+      RootConstantsYaml ConstantYaml;
+      ConstantYaml.Num32BitValues = Constants.Num32BitValues;
+      ConstantYaml.ShaderRegister = Constants.ShaderRegister;
+      ConstantYaml.RegisterSpace = Constants.RegisterSpace;
+      NewP.Data = ConstantYaml;
     } 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;
+      RootDescriptorYaml YamlDescriptor;
+      YamlDescriptor.ShaderRegister = Descriptor.ShaderRegister;
+      YamlDescriptor.RegisterSpace = Descriptor.RegisterSpace;
       if (Version > 1) {
 #define ROOT_DESCRIPTOR_FLAG(Num, Val)                                         \
-  NewP.Descriptor.Val =                                                        \
+  YamlDescriptor.Val =                                                         \
       (Descriptor.Flags &                                                      \
        llvm::to_underlying(dxbc::RootDescriptorFlag::Val)) > 0;
 #include "llvm/BinaryFormat/DXContainerConstants.def"
       }
+      NewP.Data = YamlDescriptor;
     }
 
     RootSigDesc.Parameters.push_back(NewP);
@@ -316,14 +320,24 @@ void MappingTraits<llvm::DXContainerYAML::RootParameterYamlDesc>::mapping(
   IO.mapRequired("ShaderVisibility", P.Visibility);
 
   switch (P.Type) {
-  case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
-    IO.mapRequired("Constants", P.Constants);
-    break;
+  case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
+    DXContainerYAML::RootConstantsYaml Constants;
+    if (IO.outputting())
+      Constants = std::get<DXContainerYAML::RootConstantsYaml>(P.Data);
+    IO.mapRequired("Constants", Constants);
+    if (!IO.outputting())
+      P.Data = 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;
+  case llvm::to_underlying(dxbc::RootParameterType::UAV): {
+    DXContainerYAML::RootDescriptorYaml Descriptor;
+    if (IO.outputting())
+      Descriptor = std::get<DXContainerYAML::RootDescriptorYaml>(P.Data);
+    IO.mapRequired("Descriptor", Descriptor);
+    if (!IO.outputting())
+      P.Data = Descriptor;
+  } break;
   }
 }