[DirectX] adding support to read/write descriptor table data using obj2yaml/yaml2obj

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -163,13 +163,25 @@ enum class RootDescriptorFlag : uint32_t {
 #include "DXContainerConstants.def"
 };
 
+#define DESCRIPTOR_RANGE_FLAG(Num, Val) Val = 1ull << Num,
+enum class DescriptorRangeFlag : 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 DESCRIPTOR_RANGE(Val, Enum) Enum = Val,
+enum class DescriptorRangeType : uint32_t {
+#include "DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<DescriptorRangeType>> getDescriptorRangeTypes();
+
 #define ROOT_PARAMETER(Val, Enum)                                              \
   case Val:                                                                    \
     return true;
@@ -588,6 +600,21 @@ static_assert(sizeof(ProgramSignatureElement) == 32,
 
 namespace RTS0 {
 namespace v1 {
+struct DescriptorRange {
+  uint32_t RangeType;
+  uint32_t NumDescriptors;
+  uint32_t BaseShaderRegister;
+  uint32_t RegisterSpace;
+  uint32_t OffsetInDescriptorsFromTableStart;
+  void swapBytes() {
+    sys::swapByteOrder(RangeType);
+    sys::swapByteOrder(NumDescriptors);
+    sys::swapByteOrder(BaseShaderRegister);
+    sys::swapByteOrder(RegisterSpace);
+    sys::swapByteOrder(OffsetInDescriptorsFromTableStart);
+  }
+};
+
 struct RootDescriptor {
   uint32_t ShaderRegister;
   uint32_t RegisterSpace;
@@ -655,6 +682,14 @@ struct RootDescriptor : public v1::RootDescriptor {
     sys::swapByteOrder(Flags);
   }
 };
+
+struct DescriptorRange : public v1::DescriptorRange {
+  uint32_t Flags;
+  void swapBytes() {
+    v1::DescriptorRange::swapBytes();
+    sys::swapByteOrder(Flags);
+  }
+};
 } // namespace v2
 } // namespace RTS0
 

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

@@ -86,8 +86,21 @@ ROOT_DESCRIPTOR_FLAG(3, DATA_STATIC)
 #endif // ROOT_DESCRIPTOR_FLAG
 
 
+// DESCRIPTOR_RANGE_FLAG(bit offset for the flag, name).
+#ifdef DESCRIPTOR_RANGE_FLAG
+
+DESCRIPTOR_RANGE_FLAG(0, NONE)
+DESCRIPTOR_RANGE_FLAG(1, DESCRIPTORS_VOLATILE)
+DESCRIPTOR_RANGE_FLAG(2, DATA_VOLATILE)
+DESCRIPTOR_RANGE_FLAG(3, DATA_STATIC_WHILE_SET_AT_EXECUTE)
+DESCRIPTOR_RANGE_FLAG(4, DATA_STATIC)
+DESCRIPTOR_RANGE_FLAG(16, DESCRIPTORS_STATIC_KEEPING_BUFFER_BOUNDS_CHECKS)
+#undef DESCRIPTOR_RANGE_FLAG
+#endif // DESCRIPTOR_RANGE_FLAG
+
 #ifdef ROOT_PARAMETER
 
+ROOT_PARAMETER(0, DescriptorTable)
 ROOT_PARAMETER(1, Constants32Bit)
 ROOT_PARAMETER(2, CBV)
 ROOT_PARAMETER(3, SRV)

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -25,11 +25,22 @@ struct RootParameterInfo {
       : Header(Header), Location(Location) {}
 };
 
+struct DescriptorTable {
+  SmallVector<dxbc::RTS0::v2::DescriptorRange> Ranges;
+  SmallVector<dxbc::RTS0::v2::DescriptorRange>::const_iterator begin() const {
+    return Ranges.begin();
+  }
+  SmallVector<dxbc::RTS0::v2::DescriptorRange>::const_iterator end() const {
+    return Ranges.end();
+  }
+};
+
 struct RootParametersContainer {
   SmallVector<RootParameterInfo> ParametersInfo;
 
   SmallVector<dxbc::RTS0::v1::RootConstants> Constants;
   SmallVector<dxbc::RTS0::v2::RootDescriptor> Descriptors;
+  SmallVector<DescriptorTable> Tables;
 
   void addInfo(dxbc::RTS0::v1::RootParameterHeader Header, size_t Location) {
     ParametersInfo.push_back(RootParameterInfo(Header, Location));
@@ -51,6 +62,12 @@ struct RootParametersContainer {
     Descriptors.push_back(Descriptor);
   }
 
+  void addParameter(dxbc::RTS0::v1::RootParameterHeader Header,
+                    DescriptorTable Table) {
+    addInfo(Header, Tables.size());
+    Tables.push_back(Table);
+  }
+
   const std::pair<uint32_t, uint32_t>
   getTypeAndLocForParameter(uint32_t Location) const {
     const RootParameterInfo &Info = ParametersInfo[Location];
@@ -70,6 +87,10 @@ struct RootParametersContainer {
     return Descriptors[Index];
   }
 
+  const DescriptorTable &getDescriptorTable(size_t Index) const {
+    return Tables[Index];
+  }
+
   size_t size() const { return ParametersInfo.size(); }
 
   SmallVector<RootParameterInfo>::const_iterator begin() const {

llvm/include/llvm/Object/DXContainer.h

@@ -20,6 +20,7 @@
 #include "llvm/ADT/Twine.h"
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Object/Error.h"
+#include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/TargetParser/Triple.h"
@@ -177,6 +178,50 @@ struct RootDescriptorView : RootParameterView {
   }
 };
 
+struct DescriptorTable {
+  uint32_t NumRanges;
+  uint32_t RangesOffset;
+  ViewArray<dxbc::RTS0::v2::DescriptorRange> Ranges;
+
+  typename ViewArray<dxbc::RTS0::v2::DescriptorRange>::iterator begin() const {
+    return Ranges.begin();
+  }
+
+  typename ViewArray<dxbc::RTS0::v2::DescriptorRange>::iterator end() const {
+    return Ranges.end();
+  }
+};
+
+struct DescriptorTableView : RootParameterView {
+  static bool classof(const RootParameterView *V) {
+    return (V->Header.ParameterType ==
+            llvm::to_underlying(dxbc::RootParameterType::DescriptorTable));
+  }
+
+  // Define a type alias to access the template parameter from inside classof
+  llvm::Expected<DescriptorTable> read(uint32_t Version) {
+    const char *Current = ParamData.begin();
+    DescriptorTable Table;
+
+    Table.NumRanges =
+        support::endian::read<uint32_t, llvm::endianness::little>(Current);
+    Current += sizeof(uint32_t);
+
+    Table.RangesOffset =
+        support::endian::read<uint32_t, llvm::endianness::little>(Current);
+    Current += sizeof(uint32_t);
+
+    size_t RangeSize = sizeof(dxbc::RTS0::v1::DescriptorRange);
+    if (Version > 1)
+      RangeSize = sizeof(dxbc::RTS0::v2::DescriptorRange);
+
+    Table.Ranges.Stride = RangeSize;
+    Table.Ranges.Data =
+        ParamData.substr(2 * sizeof(uint32_t), Table.NumRanges * RangeSize);
+    return Table;
+  }
+};
+
 static Error parseFailed(const Twine &Msg) {
   return make_error<GenericBinaryError>(Msg.str(), object_error::parse_failed);
 }
@@ -213,6 +258,9 @@ class RootSignature {
   llvm::Expected<RootParameterView>
   getParameter(const dxbc::RTS0::v1::RootParameterHeader &Header) const {
     size_t DataSize;
+    size_t EndOfSectionByte = getNumStaticSamplers() == 0
+                                  ? PartData.size()
+                                  : getStaticSamplersOffset();
 
     if (!dxbc::isValidParameterType(Header.ParameterType))
       return parseFailed("invalid parameter type");
@@ -229,10 +277,24 @@ class RootSignature {
       else
         DataSize = sizeof(dxbc::RTS0::v2::RootDescriptor);
       break;
+    case dxbc::RootParameterType::DescriptorTable:
+      if (Header.ParameterOffset + sizeof(uint32_t) > EndOfSectionByte)
+        return parseFailed("Reading structure out of file bounds");
+
+      uint32_t NumRanges =
+          support::endian::read<uint32_t, llvm::endianness::little>(
+              PartData.begin() + Header.ParameterOffset);
+      if (Version == 1)
+        DataSize = sizeof(dxbc::RTS0::v1::DescriptorRange) * NumRanges;
+      else
+        DataSize = sizeof(dxbc::RTS0::v2::DescriptorRange) * NumRanges;
+
+      // 4 bytes for the number of ranges in table and
+      // 4 bytes for the ranges offset
+      DataSize += 2 * sizeof(uint32_t);
+      break;
     }
-    size_t EndOfSectionByte = getNumStaticSamplers() == 0
-                                  ? PartData.size()
-                                  : getStaticSamplersOffset();
+
 
     if (Header.ParameterOffset + DataSize > EndOfSectionByte)
       return parseFailed("Reading structure out of file bounds");

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -90,6 +90,25 @@ struct RootDescriptorYaml {
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 };
 
+struct DescriptorRangeYaml {
+  uint32_t RangeType;
+  uint32_t NumDescriptors;
+  uint32_t BaseShaderRegister;
+  uint32_t RegisterSpace;
+  uint32_t OffsetInDescriptorsFromTableStart;
+
+  uint32_t getEncodedFlags() const;
+
+#define DESCRIPTOR_RANGE_FLAG(Num, Val) bool Val = false;
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
+struct DescriptorTableYaml {
+  uint32_t NumRanges;
+  uint32_t RangesOffset;
+  SmallVector<DescriptorRangeYaml> Ranges;
+};
+
 struct RootParameterHeaderYaml {
   uint32_t Type;
   uint32_t Visibility;
@@ -113,6 +132,7 @@ struct RootParameterYamlDesc {
 
   SmallVector<RootConstantsYaml> Constants;
   SmallVector<RootDescriptorYaml> Descriptors;
+  SmallVector<DescriptorTableYaml> Tables;
 
   template <typename T>
   T &getOrInsertImpl(RootParameterLocationYaml &ParamDesc,
@@ -134,6 +154,11 @@ struct RootParameterYamlDesc {
     return getOrInsertImpl(ParamDesc, Descriptors);
   }
 
+  DescriptorTableYaml &
+  getOrInsertTable(RootParameterLocationYaml &ParamDesc) {
+    return getOrInsertImpl(ParamDesc, Tables);
+  }
+
   void insertLocation(RootParameterLocationYaml &Location) {
     Locations.push_back(Location);
   }
@@ -263,6 +288,7 @@ 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::RootParameterLocationYaml)
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::DescriptorRangeYaml)
 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)
@@ -351,6 +377,14 @@ template <> struct MappingTraits<llvm::DXContainerYAML::RootDescriptorYaml> {
   static void mapping(IO &IO, llvm::DXContainerYAML::RootDescriptorYaml &D);
 };
 
+template <> struct MappingTraits<llvm::DXContainerYAML::DescriptorTableYaml> {
+  static void mapping(IO &IO, llvm::DXContainerYAML::DescriptorTableYaml &D);
+};
+
+template <> struct MappingTraits<llvm::DXContainerYAML::DescriptorRangeYaml> {
+  static void mapping(IO &IO, llvm::DXContainerYAML::DescriptorRangeYaml &D);
+};
+
 } // namespace yaml
 
 } // namespace llvm

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -47,6 +47,18 @@ size_t RootSignatureDesc::getSize() const {
         Size += sizeof(dxbc::RTS0::v2::RootDescriptor);
 
       break;
+    case llvm::to_underlying(dxbc::RootParameterType::DescriptorTable):
+      const DescriptorTable &Table =
+          ParametersContainer.getDescriptorTable(I.Location);
+
+      // 4 bytes for the number of ranges in table and
+      // 4 bytes for the ranges offset
+      Size += 2 * sizeof(uint32_t);
+      if (Version == 1)
+        Size += sizeof(dxbc::RTS0::v1::DescriptorRange) * Table.Ranges.size();
+      else
+        Size += sizeof(dxbc::RTS0::v2::DescriptorRange) * Table.Ranges.size();
+      break;
     }
   }
   return Size;
@@ -106,6 +118,27 @@ void RootSignatureDesc::write(raw_ostream &OS) const {
         support::endian::write(BOS, Descriptor.Flags, llvm::endianness::little);
       break;
     }
+    case llvm::to_underlying(dxbc::RootParameterType::DescriptorTable): {
+      const DescriptorTable &Table =
+          ParametersContainer.getDescriptorTable(Loc);
+      support::endian::write(BOS, (uint32_t)Table.Ranges.size(),
+                             llvm::endianness::little);
+      rewriteOffsetToCurrentByte(BOS, writePlaceholder(BOS));
+      for (const auto &Range : Table) {
+        support::endian::write(BOS, Range.RangeType, llvm::endianness::little);
+        support::endian::write(BOS, Range.NumDescriptors,
+                               llvm::endianness::little);
+        support::endian::write(BOS, Range.BaseShaderRegister,
+                               llvm::endianness::little);
+        support::endian::write(BOS, Range.RegisterSpace,
+                               llvm::endianness::little);
+        support::endian::write(BOS, Range.OffsetInDescriptorsFromTableStart,
+                               llvm::endianness::little);
+        if (Version > 1)
+          support::endian::write(BOS, Range.Flags, llvm::endianness::little);
+      }
+      break;
+    }
     }
   }
   assert(Storage.size() == getSize());

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -303,6 +303,26 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
           RS.ParametersContainer.addParameter(Header, Descriptor);
           break;
         }
+        case llvm::to_underlying(dxbc::RootParameterType::DescriptorTable): {
+          const DXContainerYAML::DescriptorTableYaml &TableYaml =
+          P.RootSignature->Parameters.getOrInsertTable(L);
+          mcdxbc::DescriptorTable Table;
+          for (const auto &R : TableYaml.Ranges) {
+
+            dxbc::RTS0::v2::DescriptorRange Range;
+            Range.RangeType = R.RangeType;
+            Range.NumDescriptors = R.NumDescriptors;
+            Range.BaseShaderRegister = R.BaseShaderRegister;
+            Range.RegisterSpace = R.RegisterSpace;
+            Range.OffsetInDescriptorsFromTableStart =
+                R.OffsetInDescriptorsFromTableStart;
+            if (RS.Version > 1)
+              Range.Flags = R.getEncodedFlags();
+            Table.Ranges.push_back(Range);
+          }
+          RS.ParametersContainer.addParameter(Header, Table);
+          break;
+        }
         default:
           // Handling invalid parameter type edge case. We intentionally let
           // obj2yaml/yaml2obj parse and emit invalid dxcontainer data, in order