[DXContainer] Obj2yaml support for root constants

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -63,13 +63,51 @@ struct ShaderHash {
   void swapBytes() { sys::swapByteOrder(Flags); }
 };
 
-struct RootSignatureDesc {
-  uint32_t Size;
-  uint32_t Flags;
+#define ROOT_PARAMETER(Val, Enum) Enum = Val,
+enum class RootParameterType : uint8_t {
+#include "DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<RootParameterType>> getRootParameterTypes();
+
+#define SHADER_VISIBILITY(Val, Enum) Enum = Val,
+enum class ShaderVisibilityFlag : uint8_t {
+#include "DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<ShaderVisibilityFlag>> getShaderVisibilityFlags();
+
+struct RootConstants {
+  uint32_t ShaderRegister;
+  uint32_t RegisterSpace;
+  uint32_t Num32BitValues;
 
   void swapBytes() {
-    sys::swapByteOrder(Size);
-    sys::swapByteOrder(Flags);
+    sys::swapByteOrder(ShaderRegister);
+    sys::swapByteOrder(RegisterSpace);
+    sys::swapByteOrder(Num32BitValues);
+  }
+};
+
+struct RootParameter {
+  RootParameterType ParameterType;
+  union {
+    RootConstants Constants;
+  };
+  ShaderVisibilityFlag ShaderVisibility;
+
+  void swapBytes() {
+    switch (ParameterType) {
+
+    case RootParameterType::Constants32Bit:
+      Constants.swapBytes();
+      break;
+    case RootParameterType::DescriptorTable:
+    case RootParameterType::CBV:
+    case RootParameterType::SRV:
+    case RootParameterType::UAV:
+      break;
+    }
   }
 };
 

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

@@ -55,25 +55,25 @@ SHADER_FEATURE_FLAG(31, 36, NextUnusedBit, "Next reserved shader flag bit (not a
 
 #ifdef ROOT_PARAMETER
 
-ROOT_PARAMETER(DescriptorTable)
-ROOT_PARAMETER(Constants32Bit)
-ROOT_PARAMETER(CBV)
-ROOT_PARAMETER(SRV)
-ROOT_PARAMETER(UAV)
+ROOT_PARAMETER(0, DescriptorTable)
+ROOT_PARAMETER(1, Constants32Bit)
+ROOT_PARAMETER(2, CBV)
+ROOT_PARAMETER(3, SRV)
+ROOT_PARAMETER(4, UAV)
 #undef ROOT_PARAMETER
 #endif // ROOT_PARAMETER
 
 
 #ifdef SHADER_VISIBILITY
 
-SHADER_VISIBILITY(All)
-SHADER_VISIBILITY(Vertex)
-SHADER_VISIBILITY(Hull)
-SHADER_VISIBILITY(Domain)
-SHADER_VISIBILITY(Geometry)
-SHADER_VISIBILITY(Pixel)
-SHADER_VISIBILITY(Amplification)
-SHADER_VISIBILITY(Mesh)
+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
 

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -6,6 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/BinaryFormat/DXContainer.h"
 #include <cstdint>
 #include <limits>
 
@@ -16,6 +18,7 @@ class raw_ostream;
 namespace mcdxbc {
 struct RootSignatureHeader {
   uint32_t Flags;
+  SmallVector<dxbc::RootParameter> Parameters;
 
   void swapBytes();
   void write(raw_ostream &OS);

llvm/include/llvm/Object/DXContainer.h

@@ -118,10 +118,15 @@ template <typename T> struct ViewArray {
 namespace DirectX {
 
 class RootSignature {
+
+  using ParametersArray = ViewArray<dxbc::RootParameter>;
+
 private:
   StringRef Data;
   uint32_t Size;
   uint32_t Flags;
+  uint32_t NParameters;
+  ParametersArray Parameters;
 
 public:
   RootSignature(StringRef Data) : Data(Data) {}
@@ -131,6 +136,10 @@ class RootSignature {
   uint32_t getSize() const { return Size; }
 
   uint32_t getFlags() const { return Flags; }
+
+  uint32_t getNParameters() const { return NParameters; }
+
+  ParametersArray getParameters() const { return Parameters; }
 };
 
 class PSVRuntimeInfo {

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -80,6 +80,8 @@ struct RootSignatureDesc {
 
   uint32_t getEncodedFlags();
   uint32_t Size;
+  uint32_t NumParameters;
+  SmallVector<dxbc::RootParameter> Parameters;
 
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 };
@@ -187,6 +189,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_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)
@@ -195,6 +198,8 @@ LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::PSV::ResourceKind)
 LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::D3DSystemValue)
 LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::SigComponentType)
 LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::SigMinPrecision)
+LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::RootParameterType)
+LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::ShaderVisibilityFlag)
 
 namespace llvm {
 
@@ -259,6 +264,14 @@ template <> struct MappingTraits<DXContainerYAML::RootSignatureDesc> {
                       DXContainerYAML::RootSignatureDesc &RootSignature);
 };
 
+template <> struct MappingTraits<dxbc::RootParameter> {
+  static void mapping(IO &IO, dxbc::RootParameter &RootParameter);
+};
+
+template <> struct MappingTraits<dxbc::RootConstants> {
+  static void mapping(IO &IO, dxbc::RootConstants &RootConstants);
+};
+
 } // 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, ShaderVisibilityFlag::Enum},
+
+static const EnumEntry<ShaderVisibilityFlag> ShaderVisibilityFlags[] = {
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<ShaderVisibilityFlag>> dxbc::getShaderVisibilityFlags() {
+  return ArrayRef(ShaderVisibilityFlags);
+}
+
+#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[] = {

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -7,16 +7,26 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/MC/DXContainerRootSignature.h"
+#include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Support/EndianStream.h"
-#include "llvm/Support/SwapByteOrder.h"
-#include <iterator>
 
 using namespace llvm;
 using namespace llvm::mcdxbc;
 
 void RootSignatureHeader::write(raw_ostream &OS) {
 
-  uint32_t SizeInfo = sizeof(this);
+  uint32_t SizeInfo = sizeof(RootSignatureHeader);
+  uint32_t ParamsSize = Parameters.size();
   support::endian::write(OS, SizeInfo, llvm::endianness::little);
   support::endian::write(OS, Flags, llvm::endianness::little);
+  support::endian::write(OS, ParamsSize, llvm::endianness::little);
+
+  if (Parameters.size() > 0) {
+    uint32_t BindingSize = sizeof(dxbc::RootParameter);
+
+    support::endian::write(OS, BindingSize, llvm::endianness::little);
+
+    for (const auto &Param : Parameters)
+      OS.write(reinterpret_cast<const char *>(&Param), BindingSize);
+  }
 }

llvm/lib/Object/DXContainer.cpp

@@ -251,6 +251,26 @@ Error DirectX::RootSignature::parse() {
   Flags = support::endian::read<uint32_t, llvm::endianness::little>(Current);
   Current += sizeof(uint32_t);
 
+  NParameters =
+      support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  Current += sizeof(uint32_t);
+
+  if (NParameters > 0) {
+
+    Parameters.Stride =
+        support::endian::read<uint32_t, llvm::endianness::little>(Current);
+    Current += sizeof(uint32_t);
+
+    size_t BindingDataSize = Parameters.Stride * NParameters;
+    Parameters.Data = Data.substr(Current - Data.begin(), BindingDataSize);
+
+    if (Parameters.Data.size() < BindingDataSize)
+      return parseFailed(
+          "Resource binding data extends beyond the bounds of the part");
+
+    Current += BindingDataSize;
+  }
+
   return Error::success();
 }
 

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -268,6 +268,7 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
 
       mcdxbc::RootSignatureHeader Header;
       Header.Flags = P.RootSignature->getEncodedFlags();
+      Header.Parameters = P.RootSignature->Parameters;
 
       Header.write(OS);
       break;

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -32,8 +32,11 @@ DXContainerYAML::ShaderFeatureFlags::ShaderFeatureFlags(uint64_t FlagData) {
 
 DXContainerYAML::RootSignatureDesc::RootSignatureDesc(
     const object::DirectX::RootSignature &Data)
-    : Size(Data.getSize()) {
+    : Size(Data.getSize()), NumParameters(Data.getNParameters()) {
   uint32_t Flags = Data.getFlags();
+  for (const auto &Param : Data.getParameters())
+    Parameters.push_back(Param);
+
 #define ROOT_ELEMENT_FLAG(Num, Val, Str)                                       \
   Val = (Flags & (uint32_t)dxbc::RootElementFlag::Val) > 0;
 #include "llvm/BinaryFormat/DXContainerConstants.def"
@@ -207,9 +210,50 @@ void MappingTraits<DXContainerYAML::Signature>::mapping(
   IO.mapRequired("Parameters", S.Parameters);
 }
 
+void MappingTraits<dxbc::RootParameter>::mapping(IO &IO,
+                                                 dxbc::RootParameter &S) {
+
+  IO.mapRequired("Type", S.ParameterType);
+  IO.mapRequired("ShaderVisibility", S.ShaderVisibility);
+
+  switch (S.ParameterType) {
+
+  case dxbc::RootParameterType::Constants32Bit:
+    IO.mapRequired("Constants", S.Constants);
+    break;
+  case dxbc::RootParameterType::DescriptorTable:
+  case dxbc::RootParameterType::CBV:
+  case dxbc::RootParameterType::SRV:
+  case dxbc::RootParameterType::UAV:
+    break;
+  }
+}
+
+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);
+}
+
+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);
+}
+
+void MappingTraits<dxbc::RootConstants>::mapping(IO &IO,
+                                                 dxbc::RootConstants &S) {
+
+  IO.mapRequired("Num32BitValues", S.Num32BitValues);
+  IO.mapRequired("ShaderRegister", S.ShaderRegister);
+  IO.mapRequired("RegisterSpace", S.RegisterSpace);
+}
+
 void MappingTraits<DXContainerYAML::RootSignatureDesc>::mapping(
     IO &IO, DXContainerYAML::RootSignatureDesc &S) {
   IO.mapRequired("Size", S.Size);
+  IO.mapRequired("NumParameters", S.NumParameters);
+  IO.mapRequired("Parameters", S.Parameters);
 #define ROOT_ELEMENT_FLAG(Num, Val, Str) IO.mapOptional(#Val, S.Val, false);
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }

llvm/lib/Target/DirectX/CMakeLists.txt

@@ -33,7 +33,7 @@ add_llvm_target(DirectXCodeGen
   DXILResourceAccess.cpp
   DXILShaderFlags.cpp
   DXILTranslateMetadata.cpp
-
+  DXILRootSignature.cpp
   LINK_COMPONENTS
   Analysis
   AsmPrinter