adding support for root constants in dxcontainer for ob2jyaml and yaml2obj tools

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 {
@@ -63,6 +64,7 @@ struct ShaderHash {
   void swapBytes() { sys::swapByteOrder(Flags); }
 };
 
+<<<<<<< HEAD
 struct RootSignatureDesc {
   uint32_t Size;
   uint32_t Flags;
@@ -73,6 +75,8 @@ struct RootSignatureDesc {
   }
 };
 
+=======
+>>>>>>> b1b967db8d32 (adding support for root constants in dxcontainer for ob2jyaml and yaml2obj tools)
 struct ContainerVersion {
   uint16_t Major;
   uint16_t Minor;
@@ -167,6 +171,50 @@ enum class RootElementFlag : uint32_t {
 #include "DXContainerConstants.def"
 };
 
+#define ROOT_PARAMETER(Val, Enum) Enum = Val,
+enum class RootParameterType : uint8_t {
+#include "DXContainerConstants.def"
+};
+
+#define SHADER_VISIBILITY(Val, Enum) Enum = Val,
+enum class ShaderVisibilityFlag : uint8_t {
+#include "DXContainerConstants.def"
+};
+
+struct RootConstants {
+  uint32_t ShaderRegister;
+  uint32_t RegisterSpace;
+  uint32_t Num32BitValues;
+
+  void swapBytes() {
+    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;
+    }
+  }
+};
+
 PartType parsePartType(StringRef S);
 
 struct VertexPSVInfo {
@@ -508,6 +556,8 @@ enum class SigComponentType : uint32_t {
 };
 
 ArrayRef<EnumEntry<SigComponentType>> getSigComponentTypes();
+ArrayRef<EnumEntry<RootParameterType>> getRootParameterTypes();
+ArrayRef<EnumEntry<ShaderVisibilityFlag>> getShaderVisibilityFlags();
 
 struct ProgramSignatureHeader {
   uint32_t ParamCount;

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

@@ -22,6 +22,7 @@
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/TargetParser/Triple.h"
 #include <array>
+#include <cstdint>
 #include <variant>
 
 namespace llvm {
@@ -118,10 +119,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 +137,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

@@ -187,6 +187,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 +196,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 +262,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,27 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/MC/DXContainerRootSignature.h"
+#include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Support/EndianStream.h"
-#include "llvm/Support/SwapByteOrder.h"
-#include <iterator>
+#include <cstdint>
 
 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

@@ -33,6 +33,8 @@ DXContainerYAML::ShaderFeatureFlags::ShaderFeatureFlags(uint64_t FlagData) {
 DXContainerYAML::RootSignatureDesc::RootSignatureDesc(
     const object::DirectX::RootSignature &Data)
     : Size(Data.getSize()) {
+    : Size(Data.getSize()), Version(Data.getVersion()),
+      NumParameters(Data.getNParameters()) {
   uint32_t Flags = Data.getFlags();
 #define ROOT_ELEMENT_FLAG(Num, Val, Str)                                       \
   Val = (Flags & (uint32_t)dxbc::RootElementFlag::Val) > 0;
@@ -210,6 +212,9 @@ void MappingTraits<DXContainerYAML::Signature>::mapping(
 void MappingTraits<DXContainerYAML::RootSignatureDesc>::mapping(
     IO &IO, DXContainerYAML::RootSignatureDesc &S) {
   IO.mapRequired("Size", S.Size);
+  IO.mapRequired("Version", S.Version);
+  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"
 }