[NFC] Refactoring DXContainerYaml Root Parameter representation

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -90,29 +90,53 @@ struct RootDescriptorYaml {
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 };
 
-struct RootParameterYamlDesc {
+struct RootParameterHeaderYaml {
   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;
+
+  RootParameterHeaderYaml(){};
+  RootParameterHeaderYaml(uint32_t T) : Type(T) {}
+};
+
+struct RootParameterLocationYaml {
+  RootParameterHeaderYaml Header;
+  std::optional<size_t> IndexInSignature;
+
+  RootParameterLocationYaml(){};
+  explicit RootParameterLocationYaml(RootParameterHeaderYaml Header)
+      : Header(Header) {}
+};
+
+struct RootParameterYamlDesc {
+  SmallVector<RootParameterLocationYaml> Locations;
+
+  SmallVector<RootConstantsYaml> Constants;
+  SmallVector<RootDescriptorYaml> Descriptors;
+
+  template <typename T>
+  T &getOrInsertImpl(RootParameterLocationYaml &ParamDesc,
+                     SmallVectorImpl<T> &Container) {
+    if (!ParamDesc.IndexInSignature) {
+      ParamDesc.IndexInSignature = Container.size();
+      Container.emplace_back();
     }
+    return Container[*ParamDesc.IndexInSignature];
+  }
+
+  RootConstantsYaml &
+  getOrInsertConstants(RootParameterLocationYaml &ParamDesc) {
+    return getOrInsertImpl(ParamDesc, Constants);
+  }
+
+  RootDescriptorYaml &
+  getOrInsertDescriptor(RootParameterLocationYaml &ParamDesc) {
+    return getOrInsertImpl(ParamDesc, Descriptors);
   }
 
-  union {
-    RootConstantsYaml Constants;
-    RootDescriptorYaml Descriptor;
-  };
+  void insertLocation(RootParameterLocationYaml &Location) {
+    Locations.push_back(Location);
+  }
 };
 
 struct RootSignatureYamlDesc {
@@ -124,14 +148,10 @@ struct RootSignatureYamlDesc {
   uint32_t NumStaticSamplers;
   uint32_t StaticSamplersOffset;
 
-  SmallVector<RootParameterYamlDesc> Parameters;
+  RootParameterYamlDesc Parameters;
 
   uint32_t getEncodedFlags();
 
-  iterator_range<RootParameterYamlDesc *> params() {
-    return make_range(Parameters.begin(), Parameters.end());
-  }
-
   static llvm::Expected<DXContainerYAML::RootSignatureYamlDesc>
   create(const object::DirectX::RootSignature &Data);
 
@@ -242,7 +262,7 @@ LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::ResourceBindInfo)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::SignatureElement)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::PSVInfo::MaskVector)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::SignatureParameter)
-LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::RootParameterYamlDesc)
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::RootParameterLocationYaml)
 LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::PSV::SemanticKind)
 LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::PSV::ComponentType)
 LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::PSV::InterpolationMode)
@@ -315,8 +335,12 @@ template <> struct MappingTraits<DXContainerYAML::RootSignatureYamlDesc> {
                       DXContainerYAML::RootSignatureYamlDesc &RootSignature);
 };
 
-template <> struct MappingTraits<llvm::DXContainerYAML::RootParameterYamlDesc> {
-  static void mapping(IO &IO, llvm::DXContainerYAML::RootParameterYamlDesc &P);
+template <>
+struct MappingContextTraits<DXContainerYAML::RootParameterLocationYaml,
+                            DXContainerYAML::RootSignatureYamlDesc> {
+  static void mapping(IO &IO,
+                      llvm::DXContainerYAML::RootParameterLocationYaml &L,
+                      DXContainerYAML::RootSignatureYamlDesc &S);
 };
 
 template <> struct MappingTraits<llvm::DXContainerYAML::RootConstantsYaml> {

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -273,28 +273,36 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       RS.NumStaticSamplers = P.RootSignature->NumStaticSamplers;
       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):
+      for (DXContainerYAML::RootParameterLocationYaml &L :
+           P.RootSignature->Parameters.Locations) {
+        dxbc::RootParameterHeader Header{L.Header.Type, L.Header.Visibility,
+                                         L.Header.Offset};
+
+        switch (L.Header.Type) {
+        case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
+          const DXContainerYAML::RootConstantsYaml &ConstantYaml =
+              P.RootSignature->Parameters.getOrInsertConstants(L);
           dxbc::RootConstants Constants;
-          Constants.Num32BitValues = Param.Constants.Num32BitValues;
-          Constants.RegisterSpace = Param.Constants.RegisterSpace;
-          Constants.ShaderRegister = Param.Constants.ShaderRegister;
+          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):
+        case llvm::to_underlying(dxbc::RootParameterType::SRV):
+        case llvm::to_underlying(dxbc::RootParameterType::UAV): {
+          const DXContainerYAML::RootDescriptorYaml &DescriptorYaml =
+              P.RootSignature->Parameters.getOrInsertDescriptor(L);
+
           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:
           // Handling invalid parameter type edge case. We intentionally let
           // obj2yaml/yaml2obj parse and emit invalid dxcontainer data, in order

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -53,15 +53,15 @@ DXContainerYAML::RootSignatureYamlDesc::create(
       return createStringError(std::errc::invalid_argument,
                                "Invalid value for parameter type");
 
-    RootParameterYamlDesc NewP(PH.ParameterType);
-    NewP.Offset = PH.ParameterOffset;
-    NewP.Type = PH.ParameterType;
+    RootParameterHeaderYaml Header(PH.ParameterType);
+    Header.Offset = PH.ParameterOffset;
+    Header.Type = PH.ParameterType;
 
     if (!dxbc::isValidShaderVisibility(PH.ShaderVisibility))
       return createStringError(std::errc::invalid_argument,
                                "Invalid value for shader visibility");
 
-    NewP.Visibility = PH.ShaderVisibility;
+    Header.Visibility = PH.ShaderVisibility;
 
     llvm::Expected<object::DirectX::RootParameterView> ParamViewOrErr =
         Data.getParameter(PH);
@@ -75,29 +75,36 @@ DXContainerYAML::RootSignatureYamlDesc::create(
         return std::move(E);
 
       auto Constants = *ConstantsOrErr;
+      RootParameterLocationYaml Location(Header);
+      RootConstantsYaml &ConstantYaml =
+          RootSigDesc.Parameters.getOrInsertConstants(Location);
+      RootSigDesc.Parameters.insertLocation(Location);
+      ConstantYaml.Num32BitValues = Constants.Num32BitValues;
+      ConstantYaml.ShaderRegister = Constants.ShaderRegister;
+      ConstantYaml.RegisterSpace = Constants.RegisterSpace;
 
-      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;
+      RootParameterLocationYaml Location(Header);
+      RootDescriptorYaml &YamlDescriptor =
+          RootSigDesc.Parameters.getOrInsertDescriptor(Location);
+      RootSigDesc.Parameters.insertLocation(Location);
+
+      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"
       }
     }
-
-    RootSigDesc.Parameters.push_back(NewP);
   }
 #define ROOT_ELEMENT_FLAG(Num, Val)                                            \
   RootSigDesc.Val =                                                            \
@@ -290,11 +297,36 @@ void MappingTraits<DXContainerYAML::RootSignatureYamlDesc>::mapping(
   IO.mapRequired("RootParametersOffset", S.RootParametersOffset);
   IO.mapRequired("NumStaticSamplers", S.NumStaticSamplers);
   IO.mapRequired("StaticSamplersOffset", S.StaticSamplersOffset);
-  IO.mapRequired("Parameters", S.Parameters);
+  IO.mapRequired("Parameters", S.Parameters.Locations, S);
 #define ROOT_ELEMENT_FLAG(Num, Val) IO.mapOptional(#Val, S.Val, false);
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
+void MappingContextTraits<DXContainerYAML::RootParameterLocationYaml,
+                          DXContainerYAML::RootSignatureYamlDesc>::
+    mapping(IO &IO, DXContainerYAML::RootParameterLocationYaml &L,
+            DXContainerYAML::RootSignatureYamlDesc &S) {
+  IO.mapRequired("ParameterType", L.Header.Type);
+  IO.mapRequired("ShaderVisibility", L.Header.Visibility);
+
+  switch (L.Header.Type) {
+  case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
+    DXContainerYAML::RootConstantsYaml &Constants =
+        S.Parameters.getOrInsertConstants(L);
+    IO.mapRequired("Constants", Constants);
+    break;
+  }
+  case llvm::to_underlying(dxbc::RootParameterType::CBV):
+  case llvm::to_underlying(dxbc::RootParameterType::SRV):
+  case llvm::to_underlying(dxbc::RootParameterType::UAV): {
+    DXContainerYAML::RootDescriptorYaml &Descriptor =
+        S.Parameters.getOrInsertDescriptor(L);
+    IO.mapRequired("Descriptor", Descriptor);
+    break;
+  }
+  }
+}
+
 void MappingTraits<llvm::DXContainerYAML::RootConstantsYaml>::mapping(
     IO &IO, llvm::DXContainerYAML::RootConstantsYaml &C) {
   IO.mapRequired("Num32BitValues", C.Num32BitValues);
@@ -310,23 +342,6 @@ void MappingTraits<llvm::DXContainerYAML::RootDescriptorYaml>::mapping(
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
-void MappingTraits<llvm::DXContainerYAML::RootParameterYamlDesc>::mapping(
-    IO &IO, llvm::DXContainerYAML::RootParameterYamlDesc &P) {
-  IO.mapRequired("ParameterType", P.Type);
-  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::CBV):
-  case llvm::to_underlying(dxbc::RootParameterType::SRV):
-  case llvm::to_underlying(dxbc::RootParameterType::UAV):
-    IO.mapRequired("Descriptor", P.Descriptor);
-    break;
-  }
-}
-
 void MappingTraits<DXContainerYAML::Part>::mapping(IO &IO,
                                                    DXContainerYAML::Part &P) {
   IO.mapRequired("Name", P.Name);