adding support for obj2yaml and initial tests

joaosaffran · joaosaffran · commit b2329676db39 · 2025-02-19T21:38:17.000Z
diff --git a/llvm/include/llvm/BinaryFormat/DXContainer.h b/llvm/include/llvm/BinaryFormat/DXContainer.h
@@ -160,14 +160,14 @@ enum class RootElementFlag : uint32_t {
 };
 
 #define ROOT_PARAMETER(Val, Enum) Enum = Val,
-enum class RootParameterType: uint32_t {
+enum class RootParameterType : uint32_t {
 #include "DXContainerConstants.def"
 };
 
 ArrayRef<EnumEntry<RootParameterType>> getRootParameterTypes();
 
 #define SHADER_VISIBILITY(Val, Enum) Enum = Val,
-enum class ShaderVisibilityFlag: uint32_t {
+enum class ShaderVisibilityFlag : uint32_t {
 #include "DXContainerConstants.def"
 };
 
@@ -566,21 +566,40 @@ struct RootConstants {
   uint32_t ShaderRegister;
   uint32_t RegisterSpace;
   uint32_t Num32BitValues;
+
+  void swapBytes() {
+    sys::swapByteOrder(ShaderRegister);
+    sys::swapByteOrder(RegisterSpace);
+    sys::swapByteOrder(Num32BitValues);
+  }
 };
 
 struct RootParameter {
-dxbc::RootParameterType ParameterType;
-union {
-  RootConstants Constants;
-};
-dxbc::ShaderVisibilityFlag ShaderVisibility;
+  dxbc::RootParameterType ParameterType;
+  union {
+    RootConstants Constants;
+  };
+  dxbc::ShaderVisibilityFlag ShaderVisibility;
+
+  void swapBytes() {
+    sys::swapByteOrder(ParameterType);
+    sys::swapByteOrder(ShaderVisibility);
+    switch (ParameterType) {
+
+    case RootParameterType::Constants32Bit:
+      Constants.swapBytes();
+      break;
+    }
+  }
 };
 
 struct RootSignatureHeader {
   uint32_t Version = 2;
   uint32_t Flags = 0;
+  void swapBytes() {
+    // nothing to swap
+  }
 };
-            
 
 struct RootSignatureValidations {
 
diff --git a/llvm/include/llvm/MC/DXContainerRootSignature.h b/llvm/include/llvm/MC/DXContainerRootSignature.h
@@ -15,7 +15,7 @@ namespace llvm {
 class raw_ostream;
 
 namespace mcdxbc {
-  
+
 struct RootSignatureDesc {
   dxbc::RootSignatureHeader Header;
   SmallVector<dxbc::RootParameter> Parameters;
diff --git a/llvm/include/llvm/Object/DXContainer.h b/llvm/include/llvm/Object/DXContainer.h
@@ -22,6 +22,7 @@
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/TargetParser/Triple.h"
 #include <array>
+#include <cstdint>
 #include <variant>
 
 namespace llvm {
@@ -126,6 +127,8 @@ class RootSignature {
   uint32_t StaticSamplersOffset;
   uint32_t Flags;
 
+  SmallVector<dxbc::RootParameter> Parameters;
+
 public:
   RootSignature() {}
 
@@ -135,6 +138,7 @@ class RootSignature {
   uint32_t getRootParametersOffset() const { return RootParametersOffset; }
   uint32_t getNumStaticSamplers() const { return NumStaticSamplers; }
   uint32_t getStaticSamplersOffset() const { return StaticSamplersOffset; }
+  SmallVector<dxbc::RootParameter> getParameters() const { return Parameters; }
   uint32_t getFlags() const { return Flags; }
 };
 
diff --git a/llvm/include/llvm/ObjectYAML/DXContainerYAML.h b/llvm/include/llvm/ObjectYAML/DXContainerYAML.h
@@ -84,7 +84,6 @@ struct RootSignatureYamlDesc {
 
   SmallVector<dxbc::RootParameter> Parameters;
 
-
   uint32_t getEncodedFlags();
 
 #include "llvm/BinaryFormat/DXContainerConstants.def"
@@ -269,13 +268,11 @@ template <> struct MappingTraits<DXContainerYAML::RootSignatureYamlDesc> {
 };
 
 template <> struct MappingTraits<dxbc::RootParameter> {
-  static void mapping(IO &IO,
-    dxbc::RootParameter &P);
+  static void mapping(IO &IO, dxbc::RootParameter &P);
 };
 
 template <> struct MappingTraits<dxbc::RootConstants> {
-  static void mapping(IO &IO,
-    dxbc::RootConstants &C);
+  static void mapping(IO &IO, dxbc::RootConstants &C);
 };
 
 } // namespace yaml
diff --git a/llvm/lib/MC/DXContainerRootSignature.cpp b/llvm/lib/MC/DXContainerRootSignature.cpp
@@ -16,56 +16,52 @@
 using namespace llvm;
 using namespace llvm::mcdxbc;
 
-template <typename T>
-static uint32_t getSizeOf () {
-    return static_cast<uint32_t>(sizeof(T));
+template <typename T> static uint32_t getSizeOf() {
+  return static_cast<uint32_t>(sizeof(T));
 }
 
-
-
 void RootSignatureDesc::write(raw_ostream &OS) const {
   uint32_t Offset = 16;
-  const uint32_t ParametersOffset = getSizeOf<dxbc::RootSignatureHeader>() + Offset;
+  const uint32_t ParametersOffset =
+      getSizeOf<dxbc::RootSignatureHeader>() + Offset;
   const uint32_t ParameterByteSize = Parameters.size_in_bytes();
 
-
   // Writing header information
   support::endian::write(OS, Header.Version, llvm::endianness::little);
-  Offset += getSizeOf<uint32_t>();    
+  Offset += getSizeOf<uint32_t>();
 
-  support::endian::write(OS, (uint32_t)Parameters.size(), llvm::endianness::little);
-  Offset += getSizeOf<uint32_t>();    
+  support::endian::write(OS, (uint32_t)Parameters.size(),
+                         llvm::endianness::little);
+  Offset += getSizeOf<uint32_t>();
 
   support::endian::write(OS, ParametersOffset, llvm::endianness::little);
-  Offset += getSizeOf<uint32_t>();    
+  Offset += getSizeOf<uint32_t>();
 
   support::endian::write(OS, ((uint32_t)0), llvm::endianness::little);
-  Offset += getSizeOf<uint32_t>();    
+  Offset += getSizeOf<uint32_t>();
 
-  support::endian::write(OS, ParameterByteSize + ParametersOffset, llvm::endianness::little);
-  Offset += getSizeOf<uint32_t>();    
+  support::endian::write(OS, ParameterByteSize + ParametersOffset,
+                         llvm::endianness::little);
+  Offset += getSizeOf<uint32_t>();
 
   support::endian::write(OS, Header.Flags, llvm::endianness::little);
 
-  for (const dxbc::RootParameter &P : Parameters){
-    support::endian::write(OS, P.ParameterType, llvm::endianness::little);    
+  for (const dxbc::RootParameter &P : Parameters) {
+    support::endian::write(OS, P.ParameterType, llvm::endianness::little);
     support::endian::write(OS, P.ShaderVisibility, llvm::endianness::little);
     support::endian::write(OS, Offset, llvm::endianness::little);
-    Offset += getSizeOf<uint32_t>();    
-
-    
-    switch(P.ParameterType){
-      case dxbc::RootParameterType::Constants32Bit:{
-        support::endian::write(OS, P.Constants.ShaderRegister, llvm::endianness::little);
-        Offset += getSizeOf<uint32_t>();
-        
-        support::endian::write(OS, P.Constants.RegisterSpace, llvm::endianness::little);
-        Offset += getSizeOf<uint32_t>();
-        
-        support::endian::write(OS, P.Constants.Num32BitValues, llvm::endianness::little);
-        Offset += getSizeOf<uint32_t>();
 
-      } break;
-      }
+    switch (P.ParameterType) {
+    case dxbc::RootParameterType::Constants32Bit: {
+      support::endian::write(OS, P.Constants.ShaderRegister,
+                             llvm::endianness::little);
+      support::endian::write(OS, P.Constants.RegisterSpace,
+                             llvm::endianness::little);
+      support::endian::write(OS, P.Constants.Num32BitValues,
+                             llvm::endianness::little);
+      Offset += getSizeOf<dxbc::RootConstants>() + 3 * getSizeOf<uint32_t>();
+
+    } break;
+    }
   }
 }
diff --git a/llvm/lib/Object/DXContainer.cpp b/llvm/lib/Object/DXContainer.cpp
@@ -12,6 +12,8 @@
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/FormatVariadic.h"
+#include <cstddef>
+#include <cstdint>
 
 using namespace llvm;
 using namespace llvm::object;
@@ -247,6 +249,7 @@ void DXContainer::PartIterator::updateIteratorImpl(const uint32_t Offset) {
 }
 
 Error DirectX::RootSignature::parse(StringRef Data) {
+  const char *Begin = Data.begin();
   const char *Current = Data.begin();
 
   // Root Signature headers expects 6 integers to be present.
@@ -288,6 +291,37 @@ Error DirectX::RootSignature::parse(StringRef Data) {
                             llvm::Twine(FValue));
   Flags = FValue;
 
+  Current = Begin + RootParametersOffset;
+  for (uint32_t It = 0; It < NumParameters; It++) {
+    dxbc::RootParameter NewParam;
+
+    NewParam.ParameterType =
+        support::endian::read<dxbc::RootParameterType,
+                              llvm::endianness::little>(Current);
+    Current += sizeof(dxbc::RootParameterType);
+
+    NewParam.ShaderVisibility =
+        support::endian::read<dxbc::ShaderVisibilityFlag,
+                              llvm::endianness::little>(Current);
+    Current += sizeof(dxbc::ShaderVisibilityFlag);
+
+    uint32_t Offset =
+        support::endian::read<uint32_t, llvm::endianness::little>(Current);
+    Current += sizeof(uint32_t);
+
+    switch (NewParam.ParameterType) {
+
+    case dxbc::RootParameterType::Constants32Bit: {
+      if (Error Err = readStruct(Data, Current, NewParam.Constants))
+        return Err;
+      if (sys::IsBigEndianHost)
+        NewParam.Constants.swapBytes();
+    } break;
+    }
+
+    Parameters.push_back(NewParam);
+  }
+
   return Error::success();
 }
 
diff --git a/llvm/lib/ObjectYAML/DXContainerYAML.cpp b/llvm/lib/ObjectYAML/DXContainerYAML.cpp
@@ -35,6 +35,7 @@ DXContainerYAML::RootSignatureYamlDesc::RootSignatureYamlDesc(
       NumStaticSamplers(Data.getNumStaticSamplers()),
       StaticSamplersOffset(Data.getStaticSamplersOffset()) {
   uint32_t Flags = Data.getFlags();
+  Parameters = Data.getParameters();
 #define ROOT_ELEMENT_FLAG(Num, Val)                                            \
   Val = (Flags & (uint32_t)dxbc::RootElementFlag::Val) > 0;
 #include "llvm/BinaryFormat/DXContainerConstants.def"
@@ -218,26 +219,25 @@ void MappingTraits<DXContainerYAML::RootSignatureYamlDesc>::mapping(
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
-void MappingTraits<dxbc::RootConstants>::mapping(
-  IO &IO, dxbc::RootConstants &C) {
-    IO.mapRequired("Num32BitValues", C.Num32BitValues);
-    IO.mapRequired("RegisterSpace", C.RegisterSpace);
-    IO.mapRequired("ShaderRegister", C.ShaderRegister);
-
-  }
+void MappingTraits<dxbc::RootConstants>::mapping(IO &IO,
+                                                 dxbc::RootConstants &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<dxbc::RootParameter>::mapping(IO &IO,
+                                                 dxbc::RootParameter &P) {
+  IO.mapRequired("ParameterType", P.ParameterType);
+  IO.mapRequired("ShaderVisibility", P.ShaderVisibility);
+  switch (P.ParameterType) {
 
-    case dxbc::RootParameterType::Constants32Bit:
-      IO.mapRequired("Constants", P.Constants);
+  case dxbc::RootParameterType::Constants32Bit:
+    IO.mapRequired("Constants", P.Constants);
 
-      break;
-    }
+    break;
   }
+}
 
 void MappingTraits<DXContainerYAML::Part>::mapping(IO &IO,
                                                    DXContainerYAML::Part &P) {
@@ -343,15 +343,15 @@ void ScalarEnumerationTraits<dxbc::SigComponentType>::enumeration(
 }
 
 void ScalarEnumerationTraits<dxbc::RootParameterType>::enumeration(
-  IO &IO, dxbc::RootParameterType &Value) {
-for (const auto &E : dxbc::getRootParameterTypes())
-  IO.enumCase(Value, E.Name.str().c_str(), E.Value);
+    IO &IO, dxbc::RootParameterType &Value) {
+  for (const auto &E : dxbc::getRootParameterTypes())
+    IO.enumCase(Value, E.Name.str().c_str(), E.Value);
 }
 
 void ScalarEnumerationTraits<dxbc::ShaderVisibilityFlag>::enumeration(
-  IO &IO, dxbc::ShaderVisibilityFlag &Value) {
-for (const auto &E : dxbc::getShaderVisibilityFlags())
-  IO.enumCase(Value, E.Name.str().c_str(), E.Value);
+    IO &IO, dxbc::ShaderVisibilityFlag &Value) {
+  for (const auto &E : dxbc::getShaderVisibilityFlags())
+    IO.enumCase(Value, E.Name.str().c_str(), E.Value);
 }
 
 } // namespace yaml
diff --git a/llvm/lib/Target/DirectX/DXILRootSignature.cpp b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
@@ -201,12 +201,11 @@ 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) << "NumStaticSamplers: " << 0
-       << ":\n";
-    OS << indent(Space) << "StaticSamplersOffset: " << sizeof(RS.Header) + RS.Parameters.size_in_bytes()
-       << ":\n";
+    OS << indent(Space)
+       << "RootParametersOffset: " << RS.Parameters.size_in_bytes() << ":\n";
+    OS << indent(Space) << "NumStaticSamplers: " << 0 << ":\n";
+    OS << indent(Space) << "StaticSamplersOffset: "
+       << sizeof(RS.Header) + RS.Parameters.size_in_bytes() << ":\n";
     Space--;
     // end root signature header
   }
diff --git a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll
@@ -22,8 +22,7 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 ; DXC-NEXT:    Size:            24
 ; DXC-NEXT:    RootSignature:
 ; DXC-NEXT:      Version:         2
-; DXC-NEXT:      NumParameters:   0
-; DXC-NEXT:      RootParametersOffset: 0
 ; DXC-NEXT:      NumStaticSamplers: 0
-; DXC-NEXT:      StaticSamplersOffset: 0
+; DXC-NEXT:      StaticSamplersOffset: 24
+; DXC-NEXT:      Parameters: []
 ; DXC-NEXT:      AllowInputAssemblerInputLayout: true
diff --git a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-MultipleEntryFunctions.ll
@@ -29,12 +29,12 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 ; CHECK-NEXT:   NumParameters: 0
 ; CHECK-NEXT:   RootParametersOffset: 0
 ; CHECK-NEXT:   NumStaticSamplers: 0
-; CHECK-NEXT:   StaticSamplersOffset: 0
+; CHECK-NEXT:   StaticSamplersOffset: 8
 
 ; CHECK-LABEL: Definition for 'anotherMain':
 ; CHECK-NEXT:   Flags: 0x000002
 ; CHECK-NEXT:   Version: 2
 ; CHECK-NEXT:   NumParameters: 0
 ; CHECK-NEXT:   RootParametersOffset: 0
 ; CHECK-NEXT:   NumStaticSamplers: 0
-; CHECK-NEXT:   StaticSamplersOffset: 0
+; CHECK-NEXT:   StaticSamplersOffset: 8
diff --git a/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml b/llvm/test/ObjectYAML/DXContainer/RootSignature-Flags.yaml
diff --git a/llvm/unittests/Object/DXContainerTest.cpp b/llvm/unittests/Object/DXContainerTest.cpp
diff --git a/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp b/llvm/unittests/ObjectYAML/DXContainerYAMLTest.cpp
