[DirectX] adding support in obj2yaml and yaml2obj to root constants (#127840)

Adding support for Root Constant in MC, Object and obj2yaml and
yaml2obj, this PR adds:
  - new structures to dxbc definition.
  - serialize and desirialize logic from dxcontainer to yaml
  - tests validating against dxc
  - adding support to multiple parts.

Closes: #126633

---------

Co-authored-by: joaosaffran <[email protected]>

Author: joaosaffran

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -14,6 +14,7 @@
 #define LLVM_BINARYFORMAT_DXCONTAINER_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/SwapByteOrder.h"
 #include "llvm/TargetParser/Triple.h"
 
@@ -157,6 +158,40 @@ enum class RootElementFlag : uint32_t {
 #include "DXContainerConstants.def"
 };
 
+#define ROOT_PARAMETER(Val, Enum) Enum = Val,
+enum class RootParameterType : uint32_t {
+#include "DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<RootParameterType>> getRootParameterTypes();
+
+#define ROOT_PARAMETER(Val, Enum)                                              \
+  case Val:                                                                    \
+    return true;
+inline bool isValidParameterType(uint32_t V) {
+  switch (V) {
+#include "DXContainerConstants.def"
+  }
+  return false;
+}
+
+#define SHADER_VISIBILITY(Val, Enum) Enum = Val,
+enum class ShaderVisibility : uint32_t {
+#include "DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<ShaderVisibility>> getShaderVisibility();
+
+#define SHADER_VISIBILITY(Val, Enum)                                           \
+  case Val:                                                                    \
+    return true;
+inline bool isValidShaderVisibility(uint32_t V) {
+  switch (V) {
+#include "DXContainerConstants.def"
+  }
+  return false;
+}
+
 PartType parsePartType(StringRef S);
 
 struct VertexPSVInfo {
@@ -546,15 +581,49 @@ struct ProgramSignatureElement {
 static_assert(sizeof(ProgramSignatureElement) == 32,
               "ProgramSignatureElement is misaligned");
 
-struct RootSignatureValidations {
+// following dx12 naming
+// https://learn.microsoft.com/en-us/windows/win32/api/d3d12/ns-d3d12-d3d12_root_constants
+struct RootConstants {
+  uint32_t ShaderRegister;
+  uint32_t RegisterSpace;
+  uint32_t Num32BitValues;
+
+  void swapBytes() {
+    sys::swapByteOrder(ShaderRegister);
+    sys::swapByteOrder(RegisterSpace);
+    sys::swapByteOrder(Num32BitValues);
+  }
+};
 
-  static bool isValidRootFlag(uint32_t Flags) { return (Flags & ~0xfff) == 0; }
+struct RootParameterHeader {
+  uint32_t ParameterType;
+  uint32_t ShaderVisibility;
+  uint32_t ParameterOffset;
 
-  static bool isValidVersion(uint32_t Version) {
-    return (Version == 1 || Version == 2);
+  void swapBytes() {
+    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);
+  }
+};
 } // namespace dxbc
 } // namespace llvm
 

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

@@ -72,6 +72,24 @@ ROOT_ELEMENT_FLAG(11, SamplerHeapDirectlyIndexed)
 #undef ROOT_ELEMENT_FLAG
 #endif // ROOT_ELEMENT_FLAG
 
+#ifdef ROOT_PARAMETER
+
+ROOT_PARAMETER(1, Constants32Bit)
+#undef ROOT_PARAMETER
+#endif // ROOT_PARAMETER
+
+#ifdef SHADER_VISIBILITY
+
+SHADER_VISIBILITY(0, All)
+SHADER_VISIBILITY(1, Vertex)
+SHADER_VISIBILITY(2, Hull)
+SHADER_VISIBILITY(3, Domain)
+SHADER_VISIBILITY(4, Geometry)
+SHADER_VISIBILITY(5, Pixel)
+SHADER_VISIBILITY(6, Amplification)
+SHADER_VISIBILITY(7, Mesh)
+#undef SHADER_VISIBILITY
+#endif // SHADER_VISIBILITY
 
 #ifdef DXIL_MODULE_FLAG
 

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -6,23 +6,33 @@
 //
 //===----------------------------------------------------------------------===//
 
+#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 {
-  uint32_t Version = 2;
-  uint32_t NumParameters = 0;
-  uint32_t RootParametersOffset = 0;
-  uint32_t NumStaticSamplers = 0;
-  uint32_t StaticSamplersOffset = 0;
-  uint32_t Flags = 0;
+
+  uint32_t Version = 2U;
+  uint32_t Flags = 0U;
+  uint32_t RootParameterOffset = 0U;
+  uint32_t StaticSamplersOffset = 0u;
+  uint32_t NumStaticSamplers = 0u;
+  SmallVector<mcdxbc::RootParameter> Parameters;
 
   void write(raw_ostream &OS) const;
+
+  size_t getSize() const;
 };
 } // namespace mcdxbc
 } // namespace llvm

llvm/include/llvm/Object/DXContainer.h

@@ -18,10 +18,12 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/Object/Error.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/TargetParser/Triple.h"
 #include <array>
+#include <cstddef>
 #include <variant>
 
 namespace llvm {
@@ -116,6 +118,40 @@ template <typename T> struct ViewArray {
 };
 
 namespace DirectX {
+struct RootParameterView {
+  const dxbc::RootParameterHeader &Header;
+  StringRef ParamData;
+  RootParameterView(const dxbc::RootParameterHeader &H, StringRef P)
+      : Header(H), ParamData(P) {}
+
+  template <typename T> Expected<T> readParameter() {
+    T Struct;
+    if (sizeof(T) != ParamData.size())
+      return make_error<GenericBinaryError>(
+          "Reading structure out of file bounds", object_error::parse_failed);
+
+    memcpy(&Struct, ParamData.data(), sizeof(T));
+    // DXContainer is always little endian
+    if (sys::IsBigEndianHost)
+      Struct.swapBytes();
+    return Struct;
+  }
+};
+
+struct RootConstantView : RootParameterView {
+  static bool classof(const RootParameterView *V) {
+    return V->Header.ParameterType ==
+           (uint32_t)dxbc::RootParameterType::Constants32Bit;
+  }
+
+  llvm::Expected<dxbc::RootConstants> read() {
+    return readParameter<dxbc::RootConstants>();
+  }
+};
+
+static Error parseFailed(const Twine &Msg) {
+  return make_error<GenericBinaryError>(Msg.str(), object_error::parse_failed);
+}
 
 class RootSignature {
 private:
@@ -125,17 +161,49 @@ class RootSignature {
   uint32_t NumStaticSamplers;
   uint32_t StaticSamplersOffset;
   uint32_t Flags;
+  ViewArray<dxbc::RootParameterHeader> ParametersHeaders;
+  StringRef PartData;
+
+  using param_header_iterator = ViewArray<dxbc::RootParameterHeader>::iterator;
 
 public:
-  RootSignature() {}
+  RootSignature(StringRef PD) : PartData(PD) {}
 
-  Error parse(StringRef Data);
+  Error parse();
   uint32_t getVersion() const { return Version; }
   uint32_t getNumParameters() const { return NumParameters; }
   uint32_t getRootParametersOffset() const { return RootParametersOffset; }
   uint32_t getNumStaticSamplers() const { return NumStaticSamplers; }
   uint32_t getStaticSamplersOffset() const { return StaticSamplersOffset; }
+  uint32_t getNumRootParameters() const { return ParametersHeaders.size(); }
+  llvm::iterator_range<param_header_iterator> param_headers() const {
+    return llvm::make_range(ParametersHeaders.begin(), ParametersHeaders.end());
+  }
   uint32_t getFlags() const { return Flags; }
+
+  llvm::Expected<RootParameterView>
+  getParameter(const dxbc::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);
+      break;
+    }
+    size_t EndOfSectionByte = getNumStaticSamplers() == 0
+                                  ? PartData.size()
+                                  : getStaticSamplersOffset();
+
+    if (Header.ParameterOffset + DataSize > EndOfSectionByte)
+      return parseFailed("Reading structure out of file bounds");
+
+    StringRef Buff = PartData.substr(Header.ParameterOffset, DataSize);
+    RootParameterView View = RootParameterView(Header, Buff);
+    return View;
+  }
 };
 
 class PSVRuntimeInfo {

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -74,18 +74,43 @@ struct ShaderHash {
 };
 
 #define ROOT_ELEMENT_FLAG(Num, Val) bool Val = false;
+
+struct RootConstantsYaml {
+  uint32_t ShaderRegister;
+  uint32_t RegisterSpace;
+  uint32_t Num32BitValues;
+};
+
+struct RootParameterYamlDesc {
+  uint32_t Type;
+  uint32_t Visibility;
+  uint32_t Offset;
+
+  union {
+    RootConstantsYaml Constants;
+  };
+};
+
 struct RootSignatureYamlDesc {
   RootSignatureYamlDesc() = default;
-  RootSignatureYamlDesc(const object::DirectX::RootSignature &Data);
 
   uint32_t Version;
-  uint32_t NumParameters;
+  uint32_t NumRootParameters;
   uint32_t RootParametersOffset;
   uint32_t NumStaticSamplers;
   uint32_t StaticSamplersOffset;
 
+  SmallVector<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);
+
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 };
 
@@ -192,6 +217,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_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)
@@ -264,6 +290,14 @@ 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 MappingTraits<llvm::DXContainerYAML::RootConstantsYaml> {
+  static void mapping(IO &IO, llvm::DXContainerYAML::RootConstantsYaml &C);
+};
+
 } // namespace yaml
 
 } // namespace llvm

llvm/lib/BinaryFormat/DXContainer.cpp

@@ -60,6 +60,26 @@ ArrayRef<EnumEntry<SigComponentType>> dxbc::getSigComponentTypes() {
   return ArrayRef(SigComponentTypes);
 }
 
+#define SHADER_VISIBILITY(Val, Enum) {#Enum, ShaderVisibility::Enum},
+
+static const EnumEntry<ShaderVisibility> ShaderVisibilityValues[] = {
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<ShaderVisibility>> dxbc::getShaderVisibility() {
+  return ArrayRef(ShaderVisibilityValues);
+}
+
+#define ROOT_PARAMETER(Val, Enum) {#Enum, RootParameterType::Enum},
+
+static const EnumEntry<RootParameterType> RootParameterTypes[] = {
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<RootParameterType>> dxbc::getRootParameterTypes() {
+  return ArrayRef(RootParameterTypes);
+}
+
 #define SEMANTIC_KIND(Val, Enum) {#Enum, PSV::SemanticKind::Enum},
 
 static const EnumEntry<PSV::SemanticKind> SemanticKindNames[] = {