[llvm] [DX] Add support for program signatures (PR #67346)

Chris B via llvm-commits llvm-commits at lists.llvm.org
Tue Sep 26 16:32:16 PDT 2023


https://github.com/llvm-beanz updated https://github.com/llvm/llvm-project/pull/67346

>From 53ba4daa0c1493edbb0d579f8d5fe0303b43e70c Mon Sep 17 00:00:00 2001
From: Chris Bieneman <chris.bieneman at me.com>
Date: Mon, 25 Sep 2023 11:23:59 -0500
Subject: [PATCH 1/2] [DX] Add support for program signatures

For DirectX, program signatures are encoded into three different binary
sections depending on if the signature is for inputs, outputs, or
patches. All three signature types use the same data structure encoding
so they can share a lot of logic.

This patch adds support for reading and writing program signature data
as both yaml and binary data.
---
 llvm/include/llvm/BinaryFormat/DXContainer.h  |  68 +++++
 .../BinaryFormat/DXContainerConstants.def     |  51 ++++
 llvm/include/llvm/MC/DXContainerPSVInfo.h     |  28 ++
 llvm/include/llvm/Object/DXContainer.h        | 173 +++++++-----
 .../include/llvm/ObjectYAML/DXContainerYAML.h |  29 ++
 llvm/lib/BinaryFormat/DXContainer.cpp         |  30 +++
 llvm/lib/MC/DXContainerPSVInfo.cpp            |  48 ++++
 llvm/lib/Object/DXContainer.cpp               |  24 ++
 llvm/lib/ObjectYAML/DXContainerEmitter.cpp    |  14 +
 llvm/lib/ObjectYAML/DXContainerYAML.cpp       |  37 +++
 .../DXContainer/SignatureParts.yaml           | 254 ++++++++++++++++++
 llvm/tools/obj2yaml/dxcontainer2yaml.cpp      |  19 ++
 12 files changed, 710 insertions(+), 65 deletions(-)
 create mode 100644 llvm/test/ObjectYAML/DXContainer/SignatureParts.yaml

diff --git a/llvm/include/llvm/BinaryFormat/DXContainer.h b/llvm/include/llvm/BinaryFormat/DXContainer.h
index f92071e32222e1c..5d38a939202c10e 100644
--- a/llvm/include/llvm/BinaryFormat/DXContainer.h
+++ b/llvm/include/llvm/BinaryFormat/DXContainer.h
@@ -426,6 +426,74 @@ struct ResourceBindInfo : public v0::ResourceBindInfo {
 } // namespace v2
 } // namespace PSV
 
+#define COMPONENT_PRECISION(Val, Enum) Enum = Val,
+enum class SigMinPrecision : uint32_t {
+#include "DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<SigMinPrecision>> getSigMinPrecisions();
+
+#define D3D_SYSTEM_VALUE(Val, Enum) Enum = Val,
+enum class D3DSystemValue : uint32_t {
+#include "DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<D3DSystemValue>> getD3DSystemValues();
+
+#define COMPONENT_TYPE(Val, Enum) Enum = Val,
+enum class SigComponentType : uint32_t {
+#include "DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<SigComponentType>> getSigComponentTypes();
+
+struct ProgramSignatureHeader {
+  uint32_t ParamCount;
+  uint32_t FirstParamOffset;
+
+  void swapBytes() {
+    sys::swapByteOrder(ParamCount);
+    sys::swapByteOrder(ParamCount);
+  }
+};
+
+struct ProgramSignatureElement {
+  uint32_t Stream;     // Stream index (parameters must appear in non-decreasing
+                       // stream order)
+  uint32_t NameOffset; // Offset to LPCSTR from start of ProgramSignatureHeader.
+  uint32_t Index;      // Semantic Index
+  D3DSystemValue SystemValue; // Semantic type. Similar to PSV::SemanticKind.
+  SigComponentType CompType;  // Type of bits.
+  uint32_t Register;          // Register Index (row index)
+  uint8_t Mask;               // Mask (column allocation)
+
+  // The ExclusiveMask has a different meaning for input and output signatures.
+  // For an output signature, masked components of the output register are never
+  // written to.
+  // For an input signature, masked components of the input register are always
+  // read.
+  uint8_t ExclusiveMask;
+
+  uint16_t Unused;
+  SigMinPrecision MinPrecision; // Minimum precision of input/output data
+
+  void swapBytes() {
+    sys::swapByteOrder(Stream);
+    sys::swapByteOrder(NameOffset);
+    sys::swapByteOrder(Index);
+    sys::swapByteOrder(SystemValue);
+    sys::swapByteOrder(CompType);
+    sys::swapByteOrder(Register);
+    sys::swapByteOrder(Mask);
+    sys::swapByteOrder(ExclusiveMask);
+    sys::swapByteOrder(MinPrecision);
+  }
+};
+
+// Easy to get this wrong. Earlier assertions can help determine
+static_assert(sizeof(ProgramSignatureElement) == 32,
+              "ProgramSignatureElement is misaligned");
+
 } // namespace dxbc
 } // namespace llvm
 
diff --git a/llvm/include/llvm/BinaryFormat/DXContainerConstants.def b/llvm/include/llvm/BinaryFormat/DXContainerConstants.def
index e2cbad293e16252..87dd0a5cb6ba709 100644
--- a/llvm/include/llvm/BinaryFormat/DXContainerConstants.def
+++ b/llvm/include/llvm/BinaryFormat/DXContainerConstants.def
@@ -4,6 +4,9 @@ CONTAINER_PART(DXIL)
 CONTAINER_PART(SFI0)
 CONTAINER_PART(HASH)
 CONTAINER_PART(PSV0)
+CONTAINER_PART(ISG1)
+CONTAINER_PART(OSG1)
+CONTAINER_PART(PSG1)
 
 #undef CONTAINER_PART
 #endif 
@@ -101,6 +104,20 @@ COMPONENT_TYPE(9, Float64)
 #undef COMPONENT_TYPE
 #endif
 
+#ifdef COMPONENT_PRECISION
+
+COMPONENT_PRECISION(0, Default)
+COMPONENT_PRECISION(1, Float16)
+COMPONENT_PRECISION(2, Float2_8)
+COMPONENT_PRECISION(3, Reserved)
+COMPONENT_PRECISION(4, SInt16)
+COMPONENT_PRECISION(5, UInt16)
+COMPONENT_PRECISION(0xf0, Any16)
+COMPONENT_PRECISION(0xf1, Any10)
+
+#undef COMPONENT_PRECISION
+#endif
+
 #ifdef INTERPOLATION_MODE
 
 INTERPOLATION_MODE(0, Undefined)
@@ -115,3 +132,37 @@ INTERPOLATION_MODE(8, Invalid)
 
 #undef INTERPOLATION_MODE
 #endif
+
+#ifdef D3D_SYSTEM_VALUE
+
+D3D_SYSTEM_VALUE(0, Undefined)
+D3D_SYSTEM_VALUE(1, Position)
+D3D_SYSTEM_VALUE(2, ClipDistance)
+D3D_SYSTEM_VALUE(3, CullDistance)
+D3D_SYSTEM_VALUE(4, RenderTargetArrayIndex)
+D3D_SYSTEM_VALUE(5, ViewPortArrayIndex)
+D3D_SYSTEM_VALUE(6, VertexID)
+D3D_SYSTEM_VALUE(7, PrimitiveID)
+D3D_SYSTEM_VALUE(8, InstanceID)
+D3D_SYSTEM_VALUE(9, IsFrontFace)
+D3D_SYSTEM_VALUE(10, SampleIndex)
+D3D_SYSTEM_VALUE(11, FinalQuadEdgeTessfactor)
+D3D_SYSTEM_VALUE(12, FinalQuadInsideTessfactor)
+D3D_SYSTEM_VALUE(13, FinalTriEdgeTessfactor)
+D3D_SYSTEM_VALUE(14, FinalTriInsideTessfactor)
+D3D_SYSTEM_VALUE(15, FinalLineDetailTessfactor)
+D3D_SYSTEM_VALUE(16, FinalLineDensityTessfactor)
+D3D_SYSTEM_VALUE(23, Barycentrics)
+D3D_SYSTEM_VALUE(24, ShadingRate)
+D3D_SYSTEM_VALUE(25, CullPrimitive)
+D3D_SYSTEM_VALUE(64, Target)
+D3D_SYSTEM_VALUE(65, Depth)
+D3D_SYSTEM_VALUE(66, Coverage)
+D3D_SYSTEM_VALUE(67, DepthGE)
+D3D_SYSTEM_VALUE(68, DepthLE)
+D3D_SYSTEM_VALUE(69, StencilRef)
+D3D_SYSTEM_VALUE(70, InnerCoverage)
+
+#undef D3D_SYSTEM_VALUE
+
+#endif
diff --git a/llvm/include/llvm/MC/DXContainerPSVInfo.h b/llvm/include/llvm/MC/DXContainerPSVInfo.h
index 9b1892088142b63..7d21c18d252f1cc 100644
--- a/llvm/include/llvm/MC/DXContainerPSVInfo.h
+++ b/llvm/include/llvm/MC/DXContainerPSVInfo.h
@@ -86,6 +86,34 @@ struct PSVRuntimeInfo {
   }
 };
 
+class Signature {
+  struct Parameter {
+    uint32_t Stream;
+    StringRef Name;
+    uint32_t Index;
+    dxbc::D3DSystemValue SystemValue;
+    dxbc::SigComponentType CompType;
+    uint32_t Register;
+    uint8_t Mask;
+    uint8_t ExclusiveMask;
+    dxbc::SigMinPrecision MinPrecision;
+  };
+
+  SmallVector<Parameter> Params;
+
+public:
+  void addParam(uint32_t Stream, StringRef Name, uint32_t Index,
+                dxbc::D3DSystemValue SystemValue,
+                dxbc::SigComponentType CompType, uint32_t Register,
+                uint8_t Mask, uint8_t ExclusiveMask,
+                dxbc::SigMinPrecision MinPrecision) {
+    Params.push_back(Parameter{Stream, Name, Index, SystemValue, CompType,
+                               Register, Mask, ExclusiveMask, MinPrecision});
+  }
+
+  void write(raw_ostream &OS);
+};
+
 } // namespace mcdxbc
 } // namespace llvm
 
diff --git a/llvm/include/llvm/Object/DXContainer.h b/llvm/include/llvm/Object/DXContainer.h
index 8f7ab59d7f38288..2caca6a9ae73dc0 100644
--- a/llvm/include/llvm/Object/DXContainer.h
+++ b/llvm/include/llvm/Object/DXContainer.h
@@ -27,8 +27,6 @@
 namespace llvm {
 namespace object {
 
-namespace DirectX {
-
 namespace detail {
 template <typename T>
 std::enable_if_t<std::is_arithmetic<T>::value, void> swapBytes(T &value) {
@@ -40,82 +38,85 @@ std::enable_if_t<std::is_class<T>::value, void> swapBytes(T &value) {
   value.swapBytes();
 }
 } // namespace detail
-class PSVRuntimeInfo {
 
-  // This class provides a view into the underlying resource array. The Resource
-  // data is little-endian encoded and may not be properly aligned to read
-  // directly from. The dereference operator creates a copy of the data and byte
-  // swaps it as appropriate.
-  template <typename T> struct ViewArray {
-    StringRef Data;
-    uint32_t Stride = sizeof(T); // size of each element in the list.
+// This class provides a view into the underlying resource array. The Resource
+// data is little-endian encoded and may not be properly aligned to read
+// directly from. The dereference operator creates a copy of the data and byte
+// swaps it as appropriate.
+template <typename T> struct ViewArray {
+  StringRef Data;
+  uint32_t Stride = sizeof(T); // size of each element in the list.
 
-    ViewArray() = default;
-    ViewArray(StringRef D, size_t S) : Data(D), Stride(S) {}
+  ViewArray() = default;
+  ViewArray(StringRef D, size_t S) : Data(D), Stride(S) {}
 
-    using value_type = T;
-    static constexpr uint32_t MaxStride() {
-      return static_cast<uint32_t>(sizeof(value_type));
-    }
+  using value_type = T;
+  static constexpr uint32_t MaxStride() {
+    return static_cast<uint32_t>(sizeof(value_type));
+  }
 
-    struct iterator {
-      StringRef Data;
-      uint32_t Stride; // size of each element in the list.
-      const char *Current;
-
-      iterator(const ViewArray &A, const char *C)
-          : Data(A.Data), Stride(A.Stride), Current(C) {}
-      iterator(const iterator &) = default;
-
-      value_type operator*() {
-        // Explicitly zero the structure so that unused fields are zeroed. It is
-        // up to the user to know if the fields are used by verifying the PSV
-        // version.
-        value_type Val;
-        std::memset(&Val, 0, sizeof(value_type));
-        if (Current >= Data.end())
-          return Val;
-        memcpy(static_cast<void *>(&Val), Current,
-               std::min(Stride, MaxStride()));
-        if (sys::IsBigEndianHost)
-          detail::swapBytes(Val);
+  struct iterator {
+    StringRef Data;
+    uint32_t Stride; // size of each element in the list.
+    const char *Current;
+
+    iterator(const ViewArray &A, const char *C)
+        : Data(A.Data), Stride(A.Stride), Current(C) {}
+    iterator(const iterator &) = default;
+
+    value_type operator*() {
+      // Explicitly zero the structure so that unused fields are zeroed. It is
+      // up to the user to know if the fields are used by verifying the PSV
+      // version.
+      value_type Val;
+      std::memset(&Val, 0, sizeof(value_type));
+      if (Current >= Data.end())
         return Val;
-      }
+      memcpy(static_cast<void *>(&Val), Current, std::min(Stride, MaxStride()));
+      if (sys::IsBigEndianHost)
+        detail::swapBytes(Val);
+      return Val;
+    }
 
-      iterator operator++() {
-        if (Current < Data.end())
-          Current += Stride;
-        return *this;
-      }
+    iterator operator++() {
+      if (Current < Data.end())
+        Current += Stride;
+      return *this;
+    }
 
-      iterator operator++(int) {
-        iterator Tmp = *this;
-        ++*this;
-        return Tmp;
-      }
+    iterator operator++(int) {
+      iterator Tmp = *this;
+      ++*this;
+      return Tmp;
+    }
 
-      iterator operator--() {
-        if (Current > Data.begin())
-          Current -= Stride;
-        return *this;
-      }
+    iterator operator--() {
+      if (Current > Data.begin())
+        Current -= Stride;
+      return *this;
+    }
 
-      iterator operator--(int) {
-        iterator Tmp = *this;
-        --*this;
-        return Tmp;
-      }
+    iterator operator--(int) {
+      iterator Tmp = *this;
+      --*this;
+      return Tmp;
+    }
 
-      bool operator==(const iterator I) { return I.Current == Current; }
-      bool operator!=(const iterator I) { return !(*this == I); }
-    };
+    bool operator==(const iterator I) { return I.Current == Current; }
+    bool operator!=(const iterator I) { return !(*this == I); }
+  };
 
-    iterator begin() const { return iterator(*this, Data.begin()); }
+  iterator begin() const { return iterator(*this, Data.begin()); }
 
-    iterator end() const { return iterator(*this, Data.end()); }
+  iterator end() const { return iterator(*this, Data.end()); }
 
-    size_t size() const { return Data.size() / Stride; }
-  };
+  size_t size() const { return Data.size() / Stride; }
+
+  bool isEmpty() const { return Data.empty(); }
+};
+
+namespace DirectX {
+class PSVRuntimeInfo {
 
   using ResourceArray = ViewArray<dxbc::PSV::v2::ResourceBindInfo>;
   using SigElementArray = ViewArray<dxbc::PSV::v0::SignatureElement>;
@@ -232,6 +233,36 @@ class PSVRuntimeInfo {
   }
 };
 
+class Signature {
+  ViewArray<dxbc::ProgramSignatureElement> Parameters;
+  StringRef StringTable;
+
+public:
+  ViewArray<dxbc::ProgramSignatureElement>::iterator begin() const {
+    return Parameters.begin();
+  }
+
+  ViewArray<dxbc::ProgramSignatureElement>::iterator end() const {
+    return Parameters.end();
+  }
+
+  StringRef getName(uint32_t Idx) const {
+    if (Idx == 0)
+      return "";
+
+    // Name offests are from the start of the signature data, not from the start
+    // of the string table.
+    uint32_t TableOffset =
+        Idx - sizeof(dxbc::ProgramSignatureHeader) -
+        (sizeof(dxbc::ProgramSignatureElement) * Parameters.size());
+    return StringTable.slice(TableOffset, StringTable.find('\0', TableOffset));
+  }
+
+  bool isEmpty() const { return Parameters.isEmpty(); }
+
+  Error initialize(StringRef Part);
+};
+
 } // namespace DirectX
 
 class DXContainer {
@@ -248,6 +279,9 @@ class DXContainer {
   std::optional<uint64_t> ShaderFlags;
   std::optional<dxbc::ShaderHash> Hash;
   std::optional<DirectX::PSVRuntimeInfo> PSVInfo;
+  DirectX::Signature InputSignature;
+  DirectX::Signature OutputSignature;
+  DirectX::Signature PatchConstantSignature;
 
   Error parseHeader();
   Error parsePartOffsets();
@@ -255,6 +289,7 @@ class DXContainer {
   Error parseShaderFlags(StringRef Part);
   Error parseHash(StringRef Part);
   Error parsePSVInfo(StringRef Part);
+  Error parseSignature(StringRef Part, DirectX::Signature &Array);
   friend class PartIterator;
 
 public:
@@ -340,6 +375,14 @@ class DXContainer {
   const std::optional<DirectX::PSVRuntimeInfo> &getPSVInfo() const {
     return PSVInfo;
   };
+
+  const DirectX::Signature &getInputSignature() const { return InputSignature; }
+  const DirectX::Signature &getOutputSignature() const {
+    return OutputSignature;
+  }
+  const DirectX::Signature &getPatchConstantSignature() const {
+    return PatchConstantSignature;
+  }
 };
 
 } // namespace object
diff --git a/llvm/include/llvm/ObjectYAML/DXContainerYAML.h b/llvm/include/llvm/ObjectYAML/DXContainerYAML.h
index a896b1881ff34a9..cc3870bfabd76fc 100644
--- a/llvm/include/llvm/ObjectYAML/DXContainerYAML.h
+++ b/llvm/include/llvm/ObjectYAML/DXContainerYAML.h
@@ -129,6 +129,22 @@ struct PSVInfo {
   PSVInfo(const dxbc::PSV::v2::RuntimeInfo *P);
 };
 
+struct SignatureParameter {
+  uint32_t Stream;
+  std::string Name;
+  uint32_t Index;
+  dxbc::D3DSystemValue SystemValue;
+  dxbc::SigComponentType CompType;
+  uint32_t Register;
+  uint8_t Mask;
+  uint8_t ExclusiveMask;
+  dxbc::SigMinPrecision MinPrecision;
+};
+
+struct Signature {
+  llvm::SmallVector<SignatureParameter> Parameters;
+};
+
 struct Part {
   Part() = default;
   Part(std::string N, uint32_t S) : Name(N), Size(S) {}
@@ -138,6 +154,7 @@ struct Part {
   std::optional<ShaderFlags> Flags;
   std::optional<ShaderHash> Hash;
   std::optional<PSVInfo> Info;
+  std::optional<Signature> Signature;
 };
 
 struct Object {
@@ -152,9 +169,13 @@ LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::DXContainerYAML::Part)
 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_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)
+LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::D3DSystemValue)
+LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::SigComponentType)
+LLVM_YAML_DECLARE_ENUM_TRAITS(llvm::dxbc::SigMinPrecision)
 
 namespace llvm {
 
@@ -202,6 +223,14 @@ template <> struct MappingTraits<DXContainerYAML::SignatureElement> {
   static void mapping(IO &IO, llvm::DXContainerYAML::SignatureElement &El);
 };
 
+template <> struct MappingTraits<DXContainerYAML::SignatureParameter> {
+  static void mapping(IO &IO, llvm::DXContainerYAML::SignatureParameter &El);
+};
+
+template <> struct MappingTraits<DXContainerYAML::Signature> {
+  static void mapping(IO &IO, llvm::DXContainerYAML::Signature &El);
+};
+
 } // namespace yaml
 
 } // namespace llvm
diff --git a/llvm/lib/BinaryFormat/DXContainer.cpp b/llvm/lib/BinaryFormat/DXContainer.cpp
index f6613f16fafef6b..9c0e657b069697a 100644
--- a/llvm/lib/BinaryFormat/DXContainer.cpp
+++ b/llvm/lib/BinaryFormat/DXContainer.cpp
@@ -30,6 +30,36 @@ bool ShaderHash::isPopulated() {
   return Flags > 0 || 0 != memcmp(&Digest, &Zeros, 16);
 }
 
+#define COMPONENT_PRECISION(Val, Enum) {#Enum, SigMinPrecision::Enum},
+
+static const EnumEntry<SigMinPrecision> SigMinPrecisionNames[] = {
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<SigMinPrecision>> dxbc::getSigMinPrecisions() {
+  return ArrayRef(SigMinPrecisionNames);
+}
+
+#define D3D_SYSTEM_VALUE(Val, Enum) {#Enum, D3DSystemValue::Enum},
+
+static const EnumEntry<D3DSystemValue> D3DSystemValueNames[] = {
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<D3DSystemValue>> dxbc::getD3DSystemValues() {
+  return ArrayRef(D3DSystemValueNames);
+}
+
+#define COMPONENT_TYPE(Val, Enum) {#Enum, SigComponentType::Enum},
+
+static const EnumEntry<SigComponentType> SigComponentTypes[] = {
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+};
+
+ArrayRef<EnumEntry<SigComponentType>> dxbc::getSigComponentTypes() {
+  return ArrayRef(SigComponentTypes);
+}
+
 #define SEMANTIC_KIND(Val, Enum) {#Enum, PSV::SemanticKind::Enum},
 
 static const EnumEntry<PSV::SemanticKind> SemanticKindNames[] = {
diff --git a/llvm/lib/MC/DXContainerPSVInfo.cpp b/llvm/lib/MC/DXContainerPSVInfo.cpp
index 533659053c36f3d..bdc6f79a68c0c9a 100644
--- a/llvm/lib/MC/DXContainerPSVInfo.cpp
+++ b/llvm/lib/MC/DXContainerPSVInfo.cpp
@@ -161,3 +161,51 @@ void PSVRuntimeInfo::write(raw_ostream &OS, uint32_t Version) const {
   support::endian::write_array(OS, ArrayRef<uint32_t>(PatchOutputMap),
                                support::little);
 }
+
+void Signature::write(raw_ostream &OS) {
+  SmallVector<dxbc::ProgramSignatureElement> SigParams;
+  SigParams.reserve(Params.size());
+  StringTableBuilder StrTabBuilder((StringTableBuilder::DWARF));
+
+  // Name offsets are from the start of the part. Pre-calculate the offset to
+  // the start of the string table so that it can be added to the table offset.
+  uint32_t TableStart = sizeof(dxbc::ProgramSignatureHeader) +
+                        (sizeof(dxbc::ProgramSignatureElement) * Params.size());
+
+  for (const auto &P : Params) {
+    // zero out the data
+    dxbc::ProgramSignatureElement FinalElement;
+    memset(&FinalElement, 0, sizeof(dxbc::ProgramSignatureElement));
+    FinalElement.Stream = P.Stream;
+    FinalElement.NameOffset =
+        static_cast<uint32_t>(StrTabBuilder.add(P.Name)) + TableStart;
+    FinalElement.Index = P.Index;
+    FinalElement.SystemValue = P.SystemValue;
+    FinalElement.CompType = P.CompType;
+    FinalElement.Register = P.Register;
+    FinalElement.Mask = P.Mask;
+    FinalElement.ExclusiveMask = P.ExclusiveMask;
+    FinalElement.MinPrecision = P.MinPrecision;
+    SigParams.push_back(FinalElement);
+  }
+
+  StrTabBuilder.finalizeInOrder();
+  stable_sort(SigParams, [&](const dxbc::ProgramSignatureElement &L,
+                             const dxbc::ProgramSignatureElement R) {
+    return std::tie(L.Stream, L.Register, L.NameOffset) <
+           std::tie(R.Stream, R.Register, R.NameOffset);
+  });
+  if (sys::IsBigEndianHost)
+    for (auto &El : SigParams)
+      El.swapBytes();
+
+  dxbc::ProgramSignatureHeader Header = {static_cast<uint32_t>(Params.size()),
+                                         sizeof(dxbc::ProgramSignatureHeader)};
+  if (sys::IsBigEndianHost)
+    Header.swapBytes();
+  OS.write(reinterpret_cast<const char *>(&Header),
+           sizeof(dxbc::ProgramSignatureHeader));
+  OS.write(reinterpret_cast<const char *>(SigParams.data()),
+           sizeof(dxbc::ProgramSignatureElement) * SigParams.size());
+  StrTabBuilder.write(OS);
+}
diff --git a/llvm/lib/Object/DXContainer.cpp b/llvm/lib/Object/DXContainer.cpp
index a3071d2c1b72e97..511531943b4674d 100644
--- a/llvm/lib/Object/DXContainer.cpp
+++ b/llvm/lib/Object/DXContainer.cpp
@@ -101,6 +101,18 @@ Error DXContainer::parsePSVInfo(StringRef Part) {
   return Error::success();
 }
 
+Error DirectX::Signature::initialize(StringRef Part) {
+  dxbc::ProgramSignatureHeader SigHeader;
+  if (Error Err = readStruct(Part, Part.begin(), SigHeader))
+    return Err;
+  size_t Size = sizeof(dxbc::ProgramSignatureElement) * SigHeader.ParamCount;
+  if (Part.size() < Size + SigHeader.FirstParamOffset)
+    return parseFailed("Signature extends beyond the part boundary.");
+  Parameters.Data = Part.substr(SigHeader.FirstParamOffset, Size);
+  StringTable = Part.substr(SigHeader.FirstParamOffset + Size);
+  return Error::success();
+}
+
 Error DXContainer::parsePartOffsets() {
   uint32_t LastOffset =
       sizeof(dxbc::Header) + (Header.PartCount * sizeof(uint32_t));
@@ -154,6 +166,18 @@ Error DXContainer::parsePartOffsets() {
       if (Error Err = parsePSVInfo(PartData))
         return Err;
       break;
+    case dxbc::PartType::ISG1:
+      if (Error Err = InputSignature.initialize(PartData))
+        return Err;
+      break;
+    case dxbc::PartType::OSG1:
+      if (Error Err = OutputSignature.initialize(PartData))
+        return Err;
+      break;
+    case dxbc::PartType::PSG1:
+      if (Error Err = PatchConstantSignature.initialize(PartData))
+        return Err;
+      break;
     case dxbc::PartType::Unknown:
       break;
     }
diff --git a/llvm/lib/ObjectYAML/DXContainerEmitter.cpp b/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
index 6b0e1790888854e..97003fa05523756 100644
--- a/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
+++ b/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
@@ -244,6 +244,20 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       PSV.write(OS, P.Info->Version);
       break;
     }
+    case dxbc::PartType::ISG1:
+    case dxbc::PartType::OSG1:
+    case dxbc::PartType::PSG1: {
+      if (!P.Signature.has_value())
+        continue;
+      mcdxbc::Signature Sig;
+      for (const auto &Param : P.Signature->Parameters) {
+        Sig.addParam(Param.Stream, Param.Name, Param.Index, Param.SystemValue,
+                     Param.CompType, Param.Register, Param.Mask,
+                     Param.ExclusiveMask, Param.MinPrecision);
+      }
+      Sig.write(OS);
+      break;
+    }
     case dxbc::PartType::Unknown:
       break; // Skip any handling for unrecognized parts.
     }
diff --git a/llvm/lib/ObjectYAML/DXContainerYAML.cpp b/llvm/lib/ObjectYAML/DXContainerYAML.cpp
index c7cf1ec9afc1f66..1f03f2c7d399667 100644
--- a/llvm/lib/ObjectYAML/DXContainerYAML.cpp
+++ b/llvm/lib/ObjectYAML/DXContainerYAML.cpp
@@ -159,6 +159,24 @@ void MappingTraits<DXContainerYAML::PSVInfo>::mapping(
     IO.mapRequired("PatchOutputMap", PSV.PatchOutputMap);
 }
 
+void MappingTraits<DXContainerYAML::SignatureParameter>::mapping(
+    IO &IO, DXContainerYAML::SignatureParameter &S) {
+  IO.mapRequired("Stream", S.Stream);
+  IO.mapRequired("Name", S.Name);
+  IO.mapRequired("Index", S.Index);
+  IO.mapRequired("SystemValue", S.SystemValue);
+  IO.mapRequired("CompType", S.CompType);
+  IO.mapRequired("Register", S.Register);
+  IO.mapRequired("Mask", S.Mask);
+  IO.mapRequired("ExclusiveMask", S.ExclusiveMask);
+  IO.mapRequired("MinPrecision", S.MinPrecision);
+}
+
+void MappingTraits<DXContainerYAML::Signature>::mapping(
+    IO &IO, DXContainerYAML::Signature &S) {
+  IO.mapRequired("Parameters", S.Parameters);
+}
+
 void MappingTraits<DXContainerYAML::Part>::mapping(IO &IO,
                                                    DXContainerYAML::Part &P) {
   IO.mapRequired("Name", P.Name);
@@ -167,6 +185,7 @@ void MappingTraits<DXContainerYAML::Part>::mapping(IO &IO,
   IO.mapOptional("Flags", P.Flags);
   IO.mapOptional("Hash", P.Hash);
   IO.mapOptional("PSVInfo", P.Info);
+  IO.mapOptional("Signature", P.Signature);
 }
 
 void MappingTraits<DXContainerYAML::Object>::mapping(
@@ -224,6 +243,24 @@ void ScalarEnumerationTraits<dxbc::PSV::InterpolationMode>::enumeration(
     IO.enumCase(Value, E.Name.str().c_str(), E.Value);
 }
 
+void ScalarEnumerationTraits<dxbc::D3DSystemValue>::enumeration(
+    IO &IO, dxbc::D3DSystemValue &Value) {
+  for (const auto &E : dxbc::getD3DSystemValues())
+    IO.enumCase(Value, E.Name.str().c_str(), E.Value);
+}
+
+void ScalarEnumerationTraits<dxbc::SigMinPrecision>::enumeration(
+    IO &IO, dxbc::SigMinPrecision &Value) {
+  for (const auto &E : dxbc::getSigMinPrecisions())
+    IO.enumCase(Value, E.Name.str().c_str(), E.Value);
+}
+
+void ScalarEnumerationTraits<dxbc::SigComponentType>::enumeration(
+    IO &IO, dxbc::SigComponentType &Value) {
+  for (const auto &E : dxbc::getSigComponentTypes())
+    IO.enumCase(Value, E.Name.str().c_str(), E.Value);
+}
+
 } // namespace yaml
 
 void DXContainerYAML::PSVInfo::mapInfoForVersion(yaml::IO &IO) {
diff --git a/llvm/test/ObjectYAML/DXContainer/SignatureParts.yaml b/llvm/test/ObjectYAML/DXContainer/SignatureParts.yaml
new file mode 100644
index 000000000000000..81439a00c36de6d
--- /dev/null
+++ b/llvm/test/ObjectYAML/DXContainer/SignatureParts.yaml
@@ -0,0 +1,254 @@
+# RUN: yaml2obj %s | obj2yaml | FileCheck %s
+--- !dxcontainer
+Header:
+  Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 
+                     0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:           1
+    Minor:           0
+  FileSize:        600
+  PartCount:       3
+  PartOffsets:     [ 64, 124, 184 ]
+Parts:
+  - Name:            ISG1
+    Size:            52
+    Signature:
+      Parameters:
+        - Stream:          0
+          Name:            AAA_HSFoo
+          Index:           0
+          SystemValue:     Undefined
+          CompType:        Float32
+          Register:        0
+          Mask:            7
+          ExclusiveMask:   2
+          MinPrecision:    Default
+  - Name:            OSG1
+    Size:            52
+    Signature:
+      Parameters:
+        - Stream:          0
+          Name:            SV_Position
+          Index:           0
+          SystemValue:     Position
+          CompType:        Float32
+          Register:        0
+          Mask:            15
+          ExclusiveMask:   0
+          MinPrecision:    Default
+  - Name:            PSG1
+    Size:            372
+    Signature:
+      Parameters:
+        - Stream:          0
+          Name:            SV_TessFactor
+          Index:           0
+          SystemValue:     FinalQuadEdgeTessfactor
+          CompType:        Float32
+          Register:        0
+          Mask:            8
+          ExclusiveMask:   8
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            BBB
+          Index:           0
+          SystemValue:     Undefined
+          CompType:        Float32
+          Register:        0
+          Mask:            7
+          ExclusiveMask:   0
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            SV_TessFactor
+          Index:           1
+          SystemValue:     FinalQuadEdgeTessfactor
+          CompType:        Float32
+          Register:        1
+          Mask:            8
+          ExclusiveMask:   8
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            BBB
+          Index:           1
+          SystemValue:     Undefined
+          CompType:        Float32
+          Register:        1
+          Mask:            7
+          ExclusiveMask:   0
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            SV_TessFactor
+          Index:           2
+          SystemValue:     FinalQuadEdgeTessfactor
+          CompType:        Float32
+          Register:        2
+          Mask:            8
+          ExclusiveMask:   8
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            BBB
+          Index:           2
+          SystemValue:     Undefined
+          CompType:        Float32
+          Register:        2
+          Mask:            7
+          ExclusiveMask:   0
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            SV_TessFactor
+          Index:           3
+          SystemValue:     FinalQuadEdgeTessfactor
+          CompType:        Float32
+          Register:        3
+          Mask:            8
+          ExclusiveMask:   8
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            SV_InsideTessFactor
+          Index:           0
+          SystemValue:     FinalQuadInsideTessfactor
+          CompType:        Float32
+          Register:        4
+          Mask:            8
+          ExclusiveMask:   0
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            SV_InsideTessFactor
+          Index:           1
+          SystemValue:     FinalQuadInsideTessfactor
+          CompType:        Float32
+          Register:        5
+          Mask:            8
+          ExclusiveMask:   0
+          MinPrecision:    Default
+        - Stream:          0
+          Name:            AAA
+          Index:           0
+          SystemValue:     Undefined
+          CompType:        Float32
+          Register:        6
+          Mask:            15
+          ExclusiveMask:   4
+          MinPrecision:    Default
+...
+
+# CHECK: - Name:            ISG1
+# CHECK-NEXT:   Size:            52
+# CHECK-NEXT:   Signature:
+# CHECK-NEXT:     Parameters:
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            AAA_HSFoo
+# CHECK-NEXT:         Index:           0
+# CHECK-NEXT:         SystemValue:     Undefined
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        0
+# CHECK-NEXT:         Mask:            7
+# CHECK-NEXT:         ExclusiveMask:   2
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT: - Name:            OSG1
+# CHECK-NEXT:   Size:            52
+# CHECK-NEXT:   Signature:
+# CHECK-NEXT:     Parameters:
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            SV_Position
+# CHECK-NEXT:         Index:           0
+# CHECK-NEXT:         SystemValue:     Position
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        0
+# CHECK-NEXT:         Mask:            15
+# CHECK-NEXT:         ExclusiveMask:   0
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT: - Name:            PSG1
+# CHECK-NEXT:   Size:            372
+# CHECK-NEXT:   Signature:
+# CHECK-NEXT:     Parameters:
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            SV_TessFactor
+# CHECK-NEXT:         Index:           0
+# CHECK-NEXT:         SystemValue:     FinalQuadEdgeTessfactor
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        0
+# CHECK-NEXT:         Mask:            8
+# CHECK-NEXT:         ExclusiveMask:   8
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            BBB
+# CHECK-NEXT:         Index:           0
+# CHECK-NEXT:         SystemValue:     Undefined
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        0
+# CHECK-NEXT:         Mask:            7
+# CHECK-NEXT:         ExclusiveMask:   0
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            SV_TessFactor
+# CHECK-NEXT:         Index:           1
+# CHECK-NEXT:         SystemValue:     FinalQuadEdgeTessfactor
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        1
+# CHECK-NEXT:         Mask:            8
+# CHECK-NEXT:         ExclusiveMask:   8
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            BBB
+# CHECK-NEXT:         Index:           1
+# CHECK-NEXT:         SystemValue:     Undefined
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        1
+# CHECK-NEXT:         Mask:            7
+# CHECK-NEXT:         ExclusiveMask:   0
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            SV_TessFactor
+# CHECK-NEXT:         Index:           2
+# CHECK-NEXT:         SystemValue:     FinalQuadEdgeTessfactor
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        2
+# CHECK-NEXT:         Mask:            8
+# CHECK-NEXT:         ExclusiveMask:   8
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            BBB
+# CHECK-NEXT:         Index:           2
+# CHECK-NEXT:         SystemValue:     Undefined
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        2
+# CHECK-NEXT:         Mask:            7
+# CHECK-NEXT:         ExclusiveMask:   0
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            SV_TessFactor
+# CHECK-NEXT:         Index:           3
+# CHECK-NEXT:         SystemValue:     FinalQuadEdgeTessfactor
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        3
+# CHECK-NEXT:         Mask:            8
+# CHECK-NEXT:         ExclusiveMask:   8
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            SV_InsideTessFactor
+# CHECK-NEXT:         Index:           0
+# CHECK-NEXT:         SystemValue:     FinalQuadInsideTessfactor
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        4
+# CHECK-NEXT:         Mask:            8
+# CHECK-NEXT:         ExclusiveMask:   0
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            SV_InsideTessFactor
+# CHECK-NEXT:         Index:           1
+# CHECK-NEXT:         SystemValue:     FinalQuadInsideTessfactor
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        5
+# CHECK-NEXT:         Mask:            8
+# CHECK-NEXT:         ExclusiveMask:   0
+# CHECK-NEXT:         MinPrecision:    Default
+# CHECK-NEXT:       - Stream:          0
+# CHECK-NEXT:         Name:            AAA
+# CHECK-NEXT:         Index:           0
+# CHECK-NEXT:         SystemValue:     Undefined
+# CHECK-NEXT:         CompType:        Float32
+# CHECK-NEXT:         Register:        6
+# CHECK-NEXT:         Mask:            15
+# CHECK-NEXT:         ExclusiveMask:   4
+# CHECK-NEXT:         MinPrecision:    Default
diff --git a/llvm/tools/obj2yaml/dxcontainer2yaml.cpp b/llvm/tools/obj2yaml/dxcontainer2yaml.cpp
index e301ff94d82b593..b58d7cd952aff0e 100644
--- a/llvm/tools/obj2yaml/dxcontainer2yaml.cpp
+++ b/llvm/tools/obj2yaml/dxcontainer2yaml.cpp
@@ -16,6 +16,16 @@
 using namespace llvm;
 using namespace llvm::object;
 
+static DXContainerYAML::Signature dumpSignature(const DirectX::Signature &Sig) {
+  DXContainerYAML::Signature YAML;
+  for (auto Param : Sig)
+    YAML.Parameters.push_back(DXContainerYAML::SignatureParameter{
+        Param.Stream, Sig.getName(Param.NameOffset).str(), Param.Index,
+        Param.SystemValue, Param.CompType, Param.Register, Param.Mask,
+        Param.ExclusiveMask, Param.MinPrecision});
+  return YAML;
+}
+
 static Expected<DXContainerYAML::Object *>
 dumpDXContainer(MemoryBufferRef Source) {
   assert(file_magic::dxcontainer_object == identify_magic(Source.getBuffer()));
@@ -129,6 +139,15 @@ dumpDXContainer(MemoryBufferRef Source) {
 
       break;
     }
+    case dxbc::PartType::ISG1:
+      NewPart.Signature = dumpSignature(Container.getInputSignature());
+      break;
+    case dxbc::PartType::OSG1:
+      NewPart.Signature = dumpSignature(Container.getOutputSignature());
+      break;
+    case dxbc::PartType::PSG1:
+      NewPart.Signature = dumpSignature(Container.getPatchConstantSignature());
+      break;
     case dxbc::PartType::Unknown:
       break;
     }

>From e30bb9e86447660e800d335e9d008c8a6311a418 Mon Sep 17 00:00:00 2001
From: Chris Bieneman <chris.bieneman at me.com>
Date: Tue, 26 Sep 2023 18:31:53 -0500
Subject: [PATCH 2/2] Updates based on PR feedback

---
 llvm/include/llvm/BinaryFormat/DXContainer.h |   4 +-
 llvm/include/llvm/Object/DXContainer.h       |  15 +-
 llvm/lib/Object/DXContainer.cpp              |  15 +-
 llvm/lib/ObjectYAML/DXContainerEmitter.cpp   |  12 +-
 llvm/unittests/Object/DXContainerTest.cpp    | 136 +++++++++++++++++++
 5 files changed, 164 insertions(+), 18 deletions(-)

diff --git a/llvm/include/llvm/BinaryFormat/DXContainer.h b/llvm/include/llvm/BinaryFormat/DXContainer.h
index 5d38a939202c10e..7c220687f138ce9 100644
--- a/llvm/include/llvm/BinaryFormat/DXContainer.h
+++ b/llvm/include/llvm/BinaryFormat/DXContainer.h
@@ -460,7 +460,8 @@ struct ProgramSignatureHeader {
 struct ProgramSignatureElement {
   uint32_t Stream;     // Stream index (parameters must appear in non-decreasing
                        // stream order)
-  uint32_t NameOffset; // Offset to LPCSTR from start of ProgramSignatureHeader.
+  uint32_t NameOffset; // Offset from the start of the ProgramSignatureHeader to
+                       // the start of the null terminated string for the name.
   uint32_t Index;      // Semantic Index
   D3DSystemValue SystemValue; // Semantic type. Similar to PSV::SemanticKind.
   SigComponentType CompType;  // Type of bits.
@@ -490,7 +491,6 @@ struct ProgramSignatureElement {
   }
 };
 
-// Easy to get this wrong. Earlier assertions can help determine
 static_assert(sizeof(ProgramSignatureElement) == 32,
               "ProgramSignatureElement is misaligned");
 
diff --git a/llvm/include/llvm/Object/DXContainer.h b/llvm/include/llvm/Object/DXContainer.h
index 2caca6a9ae73dc0..e36037c83df626a 100644
--- a/llvm/include/llvm/Object/DXContainer.h
+++ b/llvm/include/llvm/Object/DXContainer.h
@@ -235,6 +235,7 @@ class PSVRuntimeInfo {
 
 class Signature {
   ViewArray<dxbc::ProgramSignatureElement> Parameters;
+  uint32_t StringTableOffset;
   StringRef StringTable;
 
 public:
@@ -246,15 +247,11 @@ class Signature {
     return Parameters.end();
   }
 
-  StringRef getName(uint32_t Idx) const {
-    if (Idx == 0)
-      return "";
-
-    // Name offests are from the start of the signature data, not from the start
-    // of the string table.
-    uint32_t TableOffset =
-        Idx - sizeof(dxbc::ProgramSignatureHeader) -
-        (sizeof(dxbc::ProgramSignatureElement) * Parameters.size());
+  StringRef getName(uint32_t Offset) const {
+    // Name offsets are from the start of the signature data, not from the start
+    // of the string table. The header encodes the start offset of the sting
+    // table, so we convert the offset here.
+    uint32_t TableOffset = Offset - StringTableOffset;
     return StringTable.slice(TableOffset, StringTable.find('\0', TableOffset));
   }
 
diff --git a/llvm/lib/Object/DXContainer.cpp b/llvm/lib/Object/DXContainer.cpp
index 511531943b4674d..337d65523b7d9b6 100644
--- a/llvm/lib/Object/DXContainer.cpp
+++ b/llvm/lib/Object/DXContainer.cpp
@@ -106,10 +106,23 @@ Error DirectX::Signature::initialize(StringRef Part) {
   if (Error Err = readStruct(Part, Part.begin(), SigHeader))
     return Err;
   size_t Size = sizeof(dxbc::ProgramSignatureElement) * SigHeader.ParamCount;
+
   if (Part.size() < Size + SigHeader.FirstParamOffset)
-    return parseFailed("Signature extends beyond the part boundary.");
+    return parseFailed("Signature parameters extend beyond the part boundary");
+
   Parameters.Data = Part.substr(SigHeader.FirstParamOffset, Size);
+  
+  StringTableOffset = SigHeader.FirstParamOffset + static_cast<uint32_t>(Size);
   StringTable = Part.substr(SigHeader.FirstParamOffset + Size);
+
+  for (const auto &Param : Parameters) {
+    if (Param.NameOffset < StringTableOffset)
+      return parseFailed("Invalid parameter name offset: name starts before "
+                         "the first name offset");
+    if (Param.NameOffset - StringTableOffset > StringTable.size())
+      return parseFailed("Invalid parameter name offset: name starts after the "
+                         "end of the part data");
+  }
   return Error::success();
 }
 
diff --git a/llvm/lib/ObjectYAML/DXContainerEmitter.cpp b/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
index 97003fa05523756..09a5e41c71234f3 100644
--- a/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
+++ b/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
@@ -247,13 +247,13 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
     case dxbc::PartType::ISG1:
     case dxbc::PartType::OSG1:
     case dxbc::PartType::PSG1: {
-      if (!P.Signature.has_value())
-        continue;
       mcdxbc::Signature Sig;
-      for (const auto &Param : P.Signature->Parameters) {
-        Sig.addParam(Param.Stream, Param.Name, Param.Index, Param.SystemValue,
-                     Param.CompType, Param.Register, Param.Mask,
-                     Param.ExclusiveMask, Param.MinPrecision);
+      if (P.Signature.has_value()) {
+        for (const auto &Param : P.Signature->Parameters) {
+          Sig.addParam(Param.Stream, Param.Name, Param.Index, Param.SystemValue,
+                       Param.CompType, Param.Register, Param.Mask,
+                       Param.ExclusiveMask, Param.MinPrecision);
+        }
       }
       Sig.write(OS);
       break;
diff --git a/llvm/unittests/Object/DXContainerTest.cpp b/llvm/unittests/Object/DXContainerTest.cpp
index 94d8dc40defd46f..da640225617d404 100644
--- a/llvm/unittests/Object/DXContainerTest.cpp
+++ b/llvm/unittests/Object/DXContainerTest.cpp
@@ -640,3 +640,139 @@ TEST(DXCFile, SigElementsExtendBeyondPart) {
       FailedWithMessage(
           "Signature elements extend beyond the size of the part"));
 }
+
+TEST(DXCFile, MalformedSignature) {
+  /*
+  The tests here exercise the DXContainer Signature section parser. These tests
+  are based on modifying the binary described by the following yaml:
+
+  --- !dxcontainer
+  Header:
+    Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+                      0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+    Version:
+      Major:           1
+      Minor:           0
+    FileSize:        128
+    PartCount:       1
+    PartOffsets:     [ 64 ]
+  Parts:
+    - Name:            ISG1
+      Size:            52
+      Signature:
+        Parameters:
+          - Stream:          0
+            Name:            AAA
+            Index:           0
+            SystemValue:     Undefined
+            CompType:        Float32
+            Register:        0
+            Mask:            7
+            ExclusiveMask:   2
+            MinPrecision:    Default
+  ...
+
+  The unmodified hex sequence is:
+
+  uint8_t Buffer[] = {
+    0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x80,
+  0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00,
+  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+  0x49, 0x53, 0x47, 0x31, 0x34, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08,
+  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00, 0x00, 0x00,
+  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+  0x00, 0x07, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x41, 0x41, 0x41, 0x00,
+  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
+  };
+
+  */
+
+  {
+
+    // This binary says the signature has 10 parameters, but the part data is
+    // only big enough for 1.
+    uint8_t Buffer[] = {
+        0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+        0x80, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x49, 0x53, 0x47, 0x31, 0x34, 0x00, 0x00, 0x00,
+        0x0A, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x28, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, 0x02, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x41, 0x41, 0x41, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00};
+    EXPECT_THAT_EXPECTED(
+        DXContainer::create(getMemoryBuffer<164>(Buffer)),
+        FailedWithMessage(
+            "Signature parameters extend beyond the part boundary"));
+  }
+
+  {
+
+    // This binary only has one parameter, but the start offset is beyond the
+    // size of the part.
+    uint8_t Buffer[] = {
+        0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+        0x80, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x49, 0x53, 0x47, 0x31, 0x34, 0x00, 0x00, 0x00,
+        0x01, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x28, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, 0x02, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x41, 0x41, 0x41, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00};
+    EXPECT_THAT_EXPECTED(
+        DXContainer::create(getMemoryBuffer<164>(Buffer)),
+        FailedWithMessage(
+            "Signature parameters extend beyond the part boundary"));
+  }
+
+  {
+
+    // This parameter has a name offset of 3, which is before the start of the
+    // string table.
+    uint8_t Buffer[] = {
+        0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+        0x80, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x49, 0x53, 0x47, 0x31, 0x34, 0x00, 0x00, 0x00,
+        0x01, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, 0x02, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x41, 0x41, 0x41, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00};
+    EXPECT_THAT_EXPECTED(
+        DXContainer::create(getMemoryBuffer<164>(Buffer)),
+        FailedWithMessage("Invalid parameter name offset: name starts before "
+                          "the first name offset"));
+  }
+
+  {
+    // This parameter has a name offset of 255, which is after the end of the
+    // part.
+    uint8_t Buffer[] = {
+        0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+        0x80, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x49, 0x53, 0x47, 0x31, 0x34, 0x00, 0x00, 0x00,
+        0x01, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0xFF, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, 0x02, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x41, 0x41, 0x41, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00};
+    EXPECT_THAT_EXPECTED(
+        DXContainer::create(getMemoryBuffer<164>(Buffer)),
+        FailedWithMessage("Invalid parameter name offset: name starts after "
+                          "the end of the part data"));
+  }
+}



More information about the llvm-commits mailing list