[NFC] Refactoring DXContainerYaml Root Parameter representation

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -158,13 +158,30 @@ enum class RootElementFlag : uint32_t {
 #include "DXContainerConstants.def"
 };
 
+#define ROOT_DESCRIPTOR_FLAG(Num, Val) Val = 1ull << Num,
+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;
@@ -580,7 +597,51 @@ struct ProgramSignatureElement {
 
 static_assert(sizeof(ProgramSignatureElement) == 32,
               "ProgramSignatureElement is misaligned");
+namespace RST0 {
+namespace v0 {
+struct RootDescriptor {
+  uint32_t ShaderRegister;
+  uint32_t RegisterSpace;
+  void swapBytes() {
+    sys::swapByteOrder(ShaderRegister);
+    sys::swapByteOrder(RegisterSpace);
+  }
+};
+
+struct DescriptorRange {
+  uint32_t RangeType;
+  uint32_t NumDescriptors;
+  uint32_t BaseShaderRegister;
+  uint32_t RegisterSpace;
+  int32_t OffsetInDescriptorsFromTableStart;
+  void swapBytes() {
+    sys::swapByteOrder(RangeType);
+    sys::swapByteOrder(NumDescriptors);
+    sys::swapByteOrder(BaseShaderRegister);
+    sys::swapByteOrder(RegisterSpace);
+    sys::swapByteOrder(OffsetInDescriptorsFromTableStart);
+  }
+};
+} // namespace v0
 
+namespace v1 {
+struct RootDescriptor : public v0::RootDescriptor {
+  uint32_t Flags;
+  void swapBytes() {
+    v0::RootDescriptor::swapBytes();
+    sys::swapByteOrder(Flags);
+  }
+};
+
+struct DescriptorRange : public v0::DescriptorRange {
+  uint32_t Flags;
+  void swapBytes() {
+    v0::DescriptorRange::swapBytes();
+    sys::swapByteOrder(Flags);
+  }
+};
+} // namespace v1
+} // namespace RST0
 // following dx12 naming
 // https://learn.microsoft.com/en-us/windows/win32/api/d3d12/ns-d3d12-d3d12_root_constants
 struct RootConstants {

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

@@ -72,12 +72,50 @@ ROOT_ELEMENT_FLAG(11, SamplerHeapDirectlyIndexed)
 #undef ROOT_ELEMENT_FLAG
 #endif // ROOT_ELEMENT_FLAG
 
+
+// ROOT_DESCRIPTOR_FLAG(bit offset for the flag, name).
+#ifdef ROOT_DESCRIPTOR_FLAG
+
+ROOT_DESCRIPTOR_FLAG(0, NONE)
+ROOT_DESCRIPTOR_FLAG(1, DATA_VOLATILE)
+ROOT_DESCRIPTOR_FLAG(2, DATA_STATIC_WHILE_SET_AT_EXECUTE)
+ROOT_DESCRIPTOR_FLAG(3, DATA_STATIC)
+#undef ROOT_DESCRIPTOR_FLAG
+#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)
+ROOT_PARAMETER(4, UAV)
 #undef ROOT_PARAMETER
 #endif // ROOT_PARAMETER
 
+
+#ifdef DESCRIPTOR_RANGE
+
+DESCRIPTOR_RANGE(0, SRV)
+DESCRIPTOR_RANGE(1, UAV)
+DESCRIPTOR_RANGE(2, CBV)
+DESCRIPTOR_RANGE(3, Sampler)
+#undef DESCRIPTOR_RANGE
+#endif // DESCRIPTOR_RANGE
+
 #ifdef SHADER_VISIBILITY
 
 SHADER_VISIBILITY(0, All)

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -6,20 +6,109 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/STLForwardCompat.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/BinaryFormat/DXContainer.h"
+#include <cstddef>
 #include <cstdint>
-#include <limits>
+#include <optional>
+#include <utility>
+#include <variant>
 
 namespace llvm {
 
 class raw_ostream;
 namespace mcdxbc {
 
-struct RootParameter {
+struct RootParameterInfo {
   dxbc::RootParameterHeader Header;
-  union {
-    dxbc::RootConstants Constants;
-  };
+  size_t Location;
+
+  RootParameterInfo() = default;
+
+  RootParameterInfo(dxbc::RootParameterHeader H, size_t L)
+      : Header(H), Location(L) {}
+};
+using DescriptorRanges = std::variant<dxbc::RST0::v0::DescriptorRange,
+                                      dxbc::RST0::v1::DescriptorRange>;
+struct DescriptorTable {
+  SmallVector<DescriptorRanges> Ranges;
+
+  SmallVector<DescriptorRanges>::const_iterator begin() const {
+    return Ranges.begin();
+  }
+  SmallVector<DescriptorRanges>::const_iterator end() const {
+    return Ranges.end();
+  }
+};
+
+using RootDescriptor = std::variant<dxbc::RST0::v0::RootDescriptor,
+                                    dxbc::RST0::v1::RootDescriptor>;
+
+using ParametersView = std::variant<
+    const dxbc::RootConstants *, const dxbc::RST0::v0::RootDescriptor *,
+    const dxbc::RST0::v1::RootDescriptor *, const DescriptorTable *>;
+struct RootParametersContainer {
+  SmallVector<RootParameterInfo> ParametersInfo;
+  SmallVector<dxbc::RootConstants> Constants;
+  SmallVector<RootDescriptor> Descriptors;
+  SmallVector<DescriptorTable> Tables;
+
+  void addInfo(dxbc::RootParameterHeader H, size_t L) {
+    ParametersInfo.push_back(RootParameterInfo(H, L));
+  }
+
+  void addParameter(dxbc::RootParameterHeader H, dxbc::RootConstants C) {
+    addInfo(H, Constants.size());
+    Constants.push_back(C);
+  }
+
+  void addParameter(dxbc::RootParameterHeader H,
+                    dxbc::RST0::v0::RootDescriptor D) {
+    addInfo(H, Descriptors.size());
+    Descriptors.push_back(D);
+  }
+
+  void addParameter(dxbc::RootParameterHeader H,
+                    dxbc::RST0::v1::RootDescriptor D) {
+    addInfo(H, Descriptors.size());
+    Descriptors.push_back(D);
+  }
+
+  void addParameter(dxbc::RootParameterHeader H, DescriptorTable D) {
+    addInfo(H, Tables.size());
+    Tables.push_back(D);
+  }
+
+  std::optional<ParametersView> getParameter(const RootParameterInfo *H) const {
+    switch (H->Header.ParameterType) {
+    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
+      return &Constants[H->Location];
+    case llvm::to_underlying(dxbc::RootParameterType::CBV):
+    case llvm::to_underlying(dxbc::RootParameterType::SRV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV): {
+      const RootDescriptor &VersionedParam = Descriptors[H->Location];
+      if (std::holds_alternative<dxbc::RST0::v0::RootDescriptor>(
+              VersionedParam)) {
+        return &std::get<dxbc::RST0::v0::RootDescriptor>(VersionedParam);
+      }
+      return &std::get<dxbc::RST0::v1::RootDescriptor>(VersionedParam);
+    }
+    case llvm::to_underlying(dxbc::RootParameterType::DescriptorTable):
+      return &Tables[H->Location];
+    }
+
+    return std::nullopt;
+  }
+
+  size_t size() const { return ParametersInfo.size(); }
+
+  SmallVector<RootParameterInfo>::const_iterator begin() const {
+    return ParametersInfo.begin();
+  }
+  SmallVector<RootParameterInfo>::const_iterator end() const {
+    return ParametersInfo.end();
+  }
 };
 struct RootSignatureDesc {
 
@@ -28,7 +117,7 @@ struct RootSignatureDesc {
   uint32_t RootParameterOffset = 0U;
   uint32_t StaticSamplersOffset = 0u;
   uint32_t NumStaticSamplers = 0u;
-  SmallVector<mcdxbc::RootParameter> Parameters;
+  mcdxbc::RootParametersContainer ParametersContainer;
 
   void write(raw_ostream &OS) const;