refactoring to change representations in binary format

Author: joaosaffran

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -17,6 +17,7 @@
 #include "llvm/Support/SwapByteOrder.h"
 #include "llvm/TargetParser/Triple.h"
 
+#include <cstdint>
 #include <stdint.h>
 
 namespace llvm {
@@ -572,33 +573,32 @@ struct RootConstants {
   }
 };
 
-struct RootParameter {
+struct RootParameterHeader {
   dxbc::RootParameterType ParameterType;
-  union {
-    dxbc::RootConstants Constants;
-  };
   dxbc::ShaderVisibility ShaderVisibility;
+  uint32_t ParameterOffset;
 
   void swapBytes() {
-    sys::swapByteOrder(ShaderVisibility);
-    switch (ParameterType) {
-    case RootParameterType::Constants32Bit:
-      Constants.swapBytes();
-      break;
-    case RootParameterType::Empty:
-      llvm_unreachable("invalid value for ParameterType");
-      break;
-    }
     sys::swapByteOrder(ParameterType);
+    sys::swapByteOrder(ShaderVisibility);
+    sys::swapByteOrder(ParameterOffset);
   }
 };
 
 struct RootSignatureHeader {
   uint32_t Version;
+  uint32_t NumParameters;
+  uint32_t ParametersOffset;
+  uint32_t NumStaticSamplers;
+  uint32_t StaticSamplerOffset;
   uint32_t Flags;
 
   void swapBytes() {
     sys::swapByteOrder(Version);
+    sys::swapByteOrder(NumParameters);
+    sys::swapByteOrder(ParametersOffset);
+    sys::swapByteOrder(NumStaticSamplers);
+    sys::swapByteOrder(StaticSamplerOffset);
     sys::swapByteOrder(Flags);
   }
 };

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -7,19 +7,25 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/BinaryFormat/DXContainer.h"
-#include <cstdint>
-#include <limits>
 
 namespace llvm {
 
 class raw_ostream;
-
 namespace mcdxbc {
+
+struct RootParameter {
+  dxbc::RootParameterHeader Header;
+  union {
+    dxbc::RootConstants Constants;
+  };
+};
 struct RootSignatureDesc {
 
   dxbc::RootSignatureHeader Header;
-  SmallVector<dxbc::RootParameter> Parameters;
-  RootSignatureDesc() : Header(dxbc::RootSignatureHeader{2, 0}) {}
+  SmallVector<mcdxbc::RootParameter> Parameters;
+  RootSignatureDesc()
+      : Header(dxbc::RootSignatureHeader{
+            2, 0, sizeof(dxbc::RootSignatureHeader), 0, 0, 0}) {}
 
   void write(raw_ostream &OS) const;
 };

llvm/include/llvm/Object/DXContainer.h

@@ -116,7 +116,12 @@ template <typename T> struct ViewArray {
 };
 
 namespace DirectX {
-
+struct RootParameter {
+  dxbc::RootParameterHeader Header;
+  union {
+    dxbc::RootConstants Constants;
+  };
+};
 class RootSignature {
 private:
   uint32_t Version = 2;
@@ -126,7 +131,7 @@ class RootSignature {
   uint32_t StaticSamplersOffset = 0;
   uint32_t Flags = 0;
 
-  SmallVector<dxbc::RootParameter> Parameters;
+  SmallVector<DirectX::RootParameter> Parameters;
 
 public:
   RootSignature() {}
@@ -137,7 +142,7 @@ class RootSignature {
   uint32_t getRootParametersOffset() const { return RootParametersOffset; }
   uint32_t getNumStaticSamplers() const { return NumStaticSamplers; }
   uint32_t getStaticSamplersOffset() const { return StaticSamplersOffset; }
-  const SmallVector<dxbc::RootParameter> &getParameters() const {
+  const SmallVector<DirectX::RootParameter> &getParameters() const {
     return Parameters;
   }
   uint32_t getFlags() const { return Flags; }

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -19,6 +19,7 @@
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Object/DXContainer.h"
 #include "llvm/ObjectYAML/YAML.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <array>
 #include <cstdint>
@@ -74,6 +75,41 @@ struct ShaderHash {
 };
 
 #define ROOT_ELEMENT_FLAG(Num, Val) bool Val = false;
+
+struct RootConstantsYaml {
+  uint32_t ShaderRegister;
+  uint32_t RegisterSpace;
+  uint32_t Num32BitValues;
+};
+
+struct RootParameterYamlDesc {
+  dxbc::RootParameterType Type;
+  dxbc::ShaderVisibility Visibility;
+  uint32_t Offset;
+
+  RootParameterYamlDesc() = default;
+  RootParameterYamlDesc(const object::DirectX::RootParameter &Parameter) {
+    Type = Parameter.Header.ParameterType;
+    Visibility = Parameter.Header.ShaderVisibility;
+    Offset = Parameter.Header.ParameterOffset;
+    switch (Parameter.Header.ParameterType) {
+
+    case dxbc::RootParameterType::Constants32Bit: {
+      Constants.Num32BitValues = Parameter.Constants.Num32BitValues;
+      Constants.RegisterSpace = Parameter.Constants.RegisterSpace;
+      Constants.ShaderRegister = Parameter.Constants.ShaderRegister;
+    } break;
+    case dxbc::RootParameterType::Empty:
+      llvm_unreachable("Invalid Root Parameter Type. It should be verified "
+                       "before reaching here.");
+      break;
+    }
+  }
+  union {
+    RootConstantsYaml Constants;
+  };
+};
+
 struct RootSignatureYamlDesc {
   RootSignatureYamlDesc() = default;
   RootSignatureYamlDesc(const object::DirectX::RootSignature &Data);
@@ -82,7 +118,7 @@ struct RootSignatureYamlDesc {
   uint32_t NumStaticSamplers;
   uint32_t StaticSamplersOffset;
 
-  SmallVector<dxbc::RootParameter> Parameters;
+  SmallVector<RootParameterYamlDesc> Parameters;
 
   uint32_t getEncodedFlags();
 
@@ -192,7 +228,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::dxbc::RootParameter)
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::RootParameterYamlDesc)
 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)
@@ -267,12 +303,12 @@ template <> struct MappingTraits<DXContainerYAML::RootSignatureYamlDesc> {
                       DXContainerYAML::RootSignatureYamlDesc &RootSignature);
 };
 
-template <> struct MappingTraits<dxbc::RootParameter> {
-  static void mapping(IO &IO, dxbc::RootParameter &P);
+template <> struct MappingTraits<llvm::DXContainerYAML::RootParameterYamlDesc> {
+  static void mapping(IO &IO, llvm::DXContainerYAML::RootParameterYamlDesc &P);
 };
 
-template <> struct MappingTraits<dxbc::RootConstants> {
-  static void mapping(IO &IO, dxbc::RootConstants &C);
+template <> struct MappingTraits<llvm::DXContainerYAML::RootConstantsYaml> {
+  static void mapping(IO &IO, llvm::DXContainerYAML::RootConstantsYaml &C);
 };
 
 } // namespace yaml

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -44,8 +44,10 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
 
   SmallVector<uint32_t> ParamsOffsets;
   for (const auto &P : Parameters) {
-    support::endian::write(BOS, P.ParameterType, llvm::endianness::little);
-    support::endian::write(BOS, P.ShaderVisibility, llvm::endianness::little);
+    support::endian::write(BOS, P.Header.ParameterType,
+                           llvm::endianness::little);
+    support::endian::write(BOS, P.Header.ShaderVisibility,
+                           llvm::endianness::little);
 
     ParamsOffsets.push_back(writePlaceholder(BOS));
   }
@@ -55,7 +57,7 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
     rewriteOffset(BOS, ParamsOffsets[I]);
     const auto &P = Parameters[I];
 
-    switch (P.ParameterType) {
+    switch (P.Header.ParameterType) {
     case dxbc::RootParameterType::Constants32Bit: {
       support::endian::write(BOS, P.Constants.ShaderRegister,
                              llvm::endianness::little);

llvm/lib/Object/DXContainer.cpp

@@ -291,36 +291,39 @@ Error DirectX::RootSignature::parse(StringRef Data) {
 
   Current = Begin + RootParametersOffset;
   for (uint32_t It = 0; It < NumParameters; It++) {
-    dxbc::RootParameter NewParam;
+    DirectX::RootParameter NewParam;
 
-    NewParam.ParameterType =
+    NewParam.Header.ParameterType =
         support::endian::read<dxbc::RootParameterType,
                               llvm::endianness::little>(Current);
     if (!dxbc::RootSignatureValidations::isValidParameterType(
-            NewParam.ParameterType))
-      return validationFailed("unsupported parameter type value read: " +
-                              llvm::Twine((uint32_t)NewParam.ParameterType));
+            NewParam.Header.ParameterType))
+      return validationFailed(
+          "unsupported parameter type value read: " +
+          llvm::Twine((uint32_t)NewParam.Header.ParameterType));
 
     Current += sizeof(dxbc::RootParameterType);
 
-    NewParam.ShaderVisibility =
+    NewParam.Header.ShaderVisibility =
         support::endian::read<dxbc::ShaderVisibility, llvm::endianness::little>(
             Current);
     if (!dxbc::RootSignatureValidations::isValidShaderVisibility(
-            NewParam.ShaderVisibility))
-      return validationFailed("unsupported shader visility flag value read: " +
-                              llvm::Twine((uint32_t)NewParam.ShaderVisibility));
+            NewParam.Header.ShaderVisibility))
+      return validationFailed(
+          "unsupported shader visility flag value read: " +
+          llvm::Twine((uint32_t)NewParam.Header.ShaderVisibility));
 
     Current += sizeof(dxbc::ShaderVisibility);
 
-    uint32_t Offset =
+    NewParam.Header.ParameterOffset =
         support::endian::read<uint32_t, llvm::endianness::little>(Current);
     Current += sizeof(uint32_t);
 
-    switch (NewParam.ParameterType) {
+    switch (NewParam.Header.ParameterType) {
 
     case dxbc::RootParameterType::Constants32Bit:
-      if (Error Err = readStruct(Data, Begin + Offset, NewParam.Constants))
+      if (Error Err = readStruct(Data, Begin + NewParam.Header.ParameterOffset,
+                                 NewParam.Constants))
         return Err;
       break;
     case dxbc::RootParameterType::Empty:

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -18,6 +18,7 @@
 #include "llvm/ObjectYAML/yaml2obj.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
@@ -269,7 +270,25 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       mcdxbc::RootSignatureDesc RS;
       RS.Header.Flags = P.RootSignature->getEncodedFlags();
       RS.Header.Version = P.RootSignature->Version;
-      RS.Parameters = std::move(P.RootSignature->Parameters);
+      for (const auto &Param : P.RootSignature->Parameters) {
+        mcdxbc::RootParameter NewParam;
+        NewParam.Header = dxbc::RootParameterHeader{
+            Param.Type, Param.Visibility, Param.Offset};
+
+        switch (Param.Type) {
+
+        case dxbc::RootParameterType::Constants32Bit: {
+          NewParam.Constants.Num32BitValues = Param.Constants.Num32BitValues;
+          NewParam.Constants.RegisterSpace = Param.Constants.RegisterSpace;
+          NewParam.Constants.ShaderRegister = Param.Constants.ShaderRegister;
+        } break;
+        case dxbc::RootParameterType::Empty:
+          llvm_unreachable("Invalid parameter type");
+          break;
+        }
+
+        RS.Parameters.push_back(NewParam);
+      }
 
       RS.write(OS);
       break;

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -35,7 +35,9 @@ DXContainerYAML::RootSignatureYamlDesc::RootSignatureYamlDesc(
       NumStaticSamplers(Data.getNumStaticSamplers()),
       StaticSamplersOffset(Data.getStaticSamplersOffset()) {
   uint32_t Flags = Data.getFlags();
-  Parameters = Data.getParameters();
+  for (const auto &P : Data.getParameters()) {
+    Parameters.push_back(RootParameterYamlDesc(P));
+  }
 #define ROOT_ELEMENT_FLAG(Num, Val)                                            \
   Val = (Flags & (uint32_t)dxbc::RootElementFlag::Val) > 0;
 #include "llvm/BinaryFormat/DXContainerConstants.def"
@@ -219,18 +221,18 @@ void MappingTraits<DXContainerYAML::RootSignatureYamlDesc>::mapping(
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
-void MappingTraits<dxbc::RootConstants>::mapping(IO &IO,
-                                                 dxbc::RootConstants &C) {
+void MappingTraits<llvm::DXContainerYAML::RootConstantsYaml>::mapping(
+    IO &IO, llvm::DXContainerYAML::RootConstantsYaml &C) {
   IO.mapRequired("Num32BitValues", C.Num32BitValues);
   IO.mapRequired("RegisterSpace", C.RegisterSpace);
   IO.mapRequired("ShaderRegister", C.ShaderRegister);
 }
 
-void MappingTraits<dxbc::RootParameter>::mapping(IO &IO,
-                                                 dxbc::RootParameter &P) {
-  IO.mapRequired("ParameterType", P.ParameterType);
-  IO.mapRequired("ShaderVisibility", P.ShaderVisibility);
-  switch (P.ParameterType) {
+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 dxbc::RootParameterType::Constants32Bit:
     IO.mapRequired("Constants", P.Constants);
     break;

llvm/lib/Target/DirectX/DXILRootSignature.cpp

@@ -204,8 +204,8 @@ PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
     OS << indent(Space) << "Flags: " << format_hex(RS.Header.Flags, 8) << ":\n";
     OS << indent(Space) << "Version: " << RS.Header.Version << ":\n";
     OS << indent(Space) << "NumParameters: " << RS.Parameters.size() << ":\n";
-    OS << indent(Space)
-       << "RootParametersOffset: " << RS.Parameters.size_in_bytes() << ":\n";
+    OS << indent(Space) << "RootParametersOffset: " << sizeof(RS.Header)
+       << ":\n";
     OS << indent(Space) << "NumStaticSamplers: " << 0 << ":\n";
     OS << indent(Space) << "StaticSamplersOffset: "
        << sizeof(RS.Header) + RS.Parameters.size_in_bytes() << ":\n";

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll

@@ -27,14 +27,14 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 ; CHECK-NEXT:   Flags: 0x000001
 ; CHECK-NEXT:   Version: 2
 ; CHECK-NEXT:   NumParameters: 0
-; CHECK-NEXT:   RootParametersOffset: 0
+; CHECK-NEXT:   RootParametersOffset: 24
 ; CHECK-NEXT:   NumStaticSamplers: 0
-; CHECK-NEXT:   StaticSamplersOffset: 8
+; CHECK-NEXT:   StaticSamplersOffset: 24
 
 ; CHECK-LABEL: Definition for 'anotherMain':
 ; CHECK-NEXT:   Flags: 0x000002
 ; CHECK-NEXT:   Version: 2
 ; CHECK-NEXT:   NumParameters: 0
-; CHECK-NEXT:   RootParametersOffset: 0
+; CHECK-NEXT:   RootParametersOffset: 24
 ; CHECK-NEXT:   NumStaticSamplers: 0
-; CHECK-NEXT:   StaticSamplersOffset: 8
+; CHECK-NEXT:   StaticSamplersOffset: 24