[DX] Support pipeline state masks (#66425)

The DXContainer pipeline state information encodes a bunch of mask
vectors that are used to track things about the inputs and outputs from
each shader.

This adds support for reading and writing them throught he YAML test
interfaces. The writing logic in MC is extremely primitive and we'll
want to revisit the API for that, but since I'm not sure how we'll want
to generate the mask bits from DXIL during code generation I didn't want
to spend too much time on the API.

Fixes #59479

Author: llvm-beanz

llvm/include/llvm/MC/DXContainerPSVInfo.h

@@ -14,6 +14,7 @@
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/TargetParser/Triple.h"
 
+#include <array>
 #include <numeric>
 #include <stdint.h>
 
@@ -51,6 +52,19 @@ struct PSVRuntimeInfo {
   SmallVector<PSVSignatureElement> OutputElements;
   SmallVector<PSVSignatureElement> PatchOrPrimElements;
 
+  // TODO: Make this interface user-friendly.
+  // The interface here is bad, and we'll want to change this in the future. We
+  // probably will want to build out these mask vectors as vectors of bools and
+  // have this utility object convert them to the bit masks. I don't want to
+  // over-engineer this API now since we don't know what the data coming in to
+  // feed it will look like, so I kept it extremely simple for the immediate use
+  // case.
+  std::array<SmallVector<uint32_t>, 4> OutputVectorMasks;
+  SmallVector<uint32_t> PatchOrPrimMasks;
+  std::array<SmallVector<uint32_t>, 4> InputOutputMap;
+  SmallVector<uint32_t> InputPatchMap;
+  SmallVector<uint32_t> PatchOutputMap;
+
   // Serialize PSVInfo into the provided raw_ostream. The version field
   // specifies the data version to encode, the default value specifies encoding
   // the highest supported version.

llvm/include/llvm/Object/DXContainer.h

@@ -21,12 +21,25 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/TargetParser/Triple.h"
+#include <array>
 #include <variant>
 
 namespace llvm {
 namespace object {
 
 namespace DirectX {
+
+namespace detail {
+template <typename T>
+std::enable_if_t<std::is_arithmetic<T>::value, void> swapBytes(T &value) {
+  sys::swapByteOrder(value);
+}
+
+template <typename T>
+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
@@ -35,7 +48,7 @@ class PSVRuntimeInfo {
   // swaps it as appropriate.
   template <typename T> struct ViewArray {
     StringRef Data;
-    uint32_t Stride; // size of each element in the list.
+    uint32_t Stride = sizeof(T); // size of each element in the list.
 
     ViewArray() = default;
     ViewArray(StringRef D, size_t S) : Data(D), Stride(S) {}
@@ -65,7 +78,7 @@ class PSVRuntimeInfo {
         memcpy(static_cast<void *>(&Val), Current,
                std::min(Stride, MaxStride()));
         if (sys::IsBigEndianHost)
-          Val.swapBytes();
+          detail::swapBytes(Val);
         return Val;
       }
 
@@ -120,6 +133,12 @@ class PSVRuntimeInfo {
   SigElementArray SigOutputElements;
   SigElementArray SigPatchOrPrimElements;
 
+  std::array<ViewArray<uint32_t>, 4> OutputVectorMasks;
+  ViewArray<uint32_t> PatchOrPrimMasks;
+  std::array<ViewArray<uint32_t>, 4> InputOutputMap;
+  ViewArray<uint32_t> InputPatchMap;
+  ViewArray<uint32_t> PatchOutputMap;
+
 public:
   PSVRuntimeInfo(StringRef D) : Data(D), Size(0) {}
 
@@ -140,6 +159,22 @@ class PSVRuntimeInfo {
 
   const InfoStruct &getInfo() const { return BasicInfo; }
 
+  template <typename T> const T *getInfoAs() const {
+    if (const auto *P = std::get_if<dxbc::PSV::v2::RuntimeInfo>(&BasicInfo))
+      return static_cast<const T *>(P);
+    if (std::is_same<T, dxbc::PSV::v2::RuntimeInfo>::value)
+      return nullptr;
+
+    if (const auto *P = std::get_if<dxbc::PSV::v1::RuntimeInfo>(&BasicInfo))
+      return static_cast<const T *>(P);
+    if (std::is_same<T, dxbc::PSV::v1::RuntimeInfo>::value)
+      return nullptr;
+
+    if (const auto *P = std::get_if<dxbc::PSV::v0::RuntimeInfo>(&BasicInfo))
+      return static_cast<const T *>(P);
+    return nullptr;
+  }
+
   StringRef getStringTable() const { return StringTable; }
   ArrayRef<uint32_t> getSemanticIndexTable() const {
     return SemanticIndexTable;
@@ -155,7 +190,46 @@ class PSVRuntimeInfo {
     return SigPatchOrPrimElements;
   }
 
+  ViewArray<uint32_t> getOutputVectorMasks(size_t Idx) const {
+    assert(Idx < 4);
+    return OutputVectorMasks[Idx];
+  }
+
+  ViewArray<uint32_t> getPatchOrPrimMasks() const { return PatchOrPrimMasks; }
+
+  ViewArray<uint32_t> getInputOutputMap(size_t Idx) const {
+    assert(Idx < 4);
+    return InputOutputMap[Idx];
+  }
+
+  ViewArray<uint32_t> getInputPatchMap() const { return InputPatchMap; }
+  ViewArray<uint32_t> getPatchOutputMap() const { return PatchOutputMap; }
+
   uint32_t getSigElementStride() const { return SigInputElements.Stride; }
+
+  bool usesViewID() const {
+    if (const auto *P = getInfoAs<dxbc::PSV::v1::RuntimeInfo>())
+      return P->UsesViewID != 0;
+    return false;
+  }
+
+  uint8_t getInputVectorCount() const {
+    if (const auto *P = getInfoAs<dxbc::PSV::v1::RuntimeInfo>())
+      return P->SigInputVectors;
+    return 0;
+  }
+
+  ArrayRef<uint8_t> getOutputVectorCounts() const {
+    if (const auto *P = getInfoAs<dxbc::PSV::v1::RuntimeInfo>())
+      return ArrayRef<uint8_t>(P->SigOutputVectors);
+    return ArrayRef<uint8_t>();
+  }
+
+  uint8_t getPatchConstOrPrimVectorCount() const {
+    if (const auto *P = getInfoAs<dxbc::PSV::v1::RuntimeInfo>())
+      return P->GeomData.SigPatchConstOrPrimVectors;
+    return 0;
+  }
 };
 
 } // namespace DirectX

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -19,6 +19,7 @@
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/ObjectYAML/YAML.h"
 #include "llvm/Support/YAMLTraits.h"
+#include <array>
 #include <cstdint>
 #include <optional>
 #include <string>
@@ -113,6 +114,13 @@ struct PSVInfo {
   SmallVector<SignatureElement> SigOutputElements;
   SmallVector<SignatureElement> SigPatchOrPrimElements;
 
+  using MaskVector = SmallVector<llvm::yaml::Hex32>;
+  std::array<MaskVector, 4> OutputVectorMasks;
+  MaskVector PatchOrPrimMasks;
+  std::array<MaskVector, 4> InputOutputMap;
+  MaskVector InputPatchMap;
+  MaskVector PatchOutputMap;
+
   void mapInfoForVersion(yaml::IO &IO);
 
   PSVInfo();
@@ -143,6 +151,7 @@ struct Object {
 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_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/include/llvm/Support/EndianStream.h

@@ -25,6 +25,15 @@ namespace support {
 
 namespace endian {
 
+template <typename value_type>
+inline void write_array(raw_ostream &os, ArrayRef<value_type> values,
+                        endianness endian) {
+  for (const auto orig : values) {
+    value_type value = byte_swap<value_type>(orig, endian);
+    os.write((const char *)&value, sizeof(value_type));
+  }
+}
+
 template <typename value_type>
 inline void write(raw_ostream &os, value_type value, endianness endian) {
   value = byte_swap<value_type>(value, endian);

llvm/lib/MC/DXContainerPSVInfo.cpp

@@ -147,4 +147,17 @@ void PSVRuntimeInfo::write(raw_ostream &OS, uint32_t Version) const {
     OS.write(reinterpret_cast<const char *>(&SignatureElements[0]),
              SignatureElements.size() * sizeof(v0::SignatureElement));
   }
+
+  for (const auto &MaskVector : OutputVectorMasks)
+    support::endian::write_array(OS, ArrayRef<uint32_t>(MaskVector),
+                                 support::little);
+  support::endian::write_array(OS, ArrayRef<uint32_t>(PatchOrPrimMasks),
+                               support::little);
+  for (const auto &MaskVector : InputOutputMap)
+    support::endian::write_array(OS, ArrayRef<uint32_t>(MaskVector),
+                                 support::little);
+  support::endian::write_array(OS, ArrayRef<uint32_t>(InputPatchMap),
+                               support::little);
+  support::endian::write_array(OS, ArrayRef<uint32_t>(PatchOutputMap),
+                               support::little);
 }

llvm/lib/Object/DXContainer.cpp

@@ -321,6 +321,68 @@ Error DirectX::PSVRuntimeInfo::parse(uint16_t ShaderKind) {
     Current += PSize;
   }
 
+  ArrayRef<uint8_t> OutputVectorCounts = getOutputVectorCounts();
+  uint8_t PatchConstOrPrimVectorCount = getPatchConstOrPrimVectorCount();
+  uint8_t InputVectorCount = getInputVectorCount();
+
+  auto maskDwordSize = [](uint8_t Vector) {
+    return (static_cast<uint32_t>(Vector) + 7) >> 3;
+  };
+
+  auto mapTableSize = [maskDwordSize](uint8_t X, uint8_t Y) {
+    return maskDwordSize(Y) * X * 4;
+  };
+
+  if (usesViewID()) {
+    for (uint32_t I = 0; I < OutputVectorCounts.size(); ++I) {
+      // The vector mask is one bit per component and 4 components per vector.
+      // We can compute the number of dwords required by rounding up to the next
+      // multiple of 8.
+      uint32_t NumDwords =
+          maskDwordSize(static_cast<uint32_t>(OutputVectorCounts[I]));
+      size_t NumBytes = NumDwords * sizeof(uint32_t);
+      OutputVectorMasks[I].Data = Data.substr(Current - Data.begin(), NumBytes);
+      Current += NumBytes;
+    }
+
+    if (ShaderStage == Triple::Hull && PatchConstOrPrimVectorCount > 0) {
+      uint32_t NumDwords = maskDwordSize(PatchConstOrPrimVectorCount);
+      size_t NumBytes = NumDwords * sizeof(uint32_t);
+      PatchOrPrimMasks.Data = Data.substr(Current - Data.begin(), NumBytes);
+      Current += NumBytes;
+    }
+  }
+
+  // Input/Output mapping table
+  for (uint32_t I = 0; I < OutputVectorCounts.size(); ++I) {
+    if (InputVectorCount == 0 || OutputVectorCounts[I] == 0)
+      continue;
+    uint32_t NumDwords = mapTableSize(InputVectorCount, OutputVectorCounts[I]);
+    size_t NumBytes = NumDwords * sizeof(uint32_t);
+    InputOutputMap[I].Data = Data.substr(Current - Data.begin(), NumBytes);
+    Current += NumBytes;
+  }
+
+  // Hull shader: Input/Patch mapping table
+  if (ShaderStage == Triple::Hull && PatchConstOrPrimVectorCount > 0 &&
+      InputVectorCount > 0) {
+    uint32_t NumDwords =
+        mapTableSize(InputVectorCount, PatchConstOrPrimVectorCount);
+    size_t NumBytes = NumDwords * sizeof(uint32_t);
+    InputPatchMap.Data = Data.substr(Current - Data.begin(), NumBytes);
+    Current += NumBytes;
+  }
+
+  // Domain Shader: Patch/Output mapping table
+  if (ShaderStage == Triple::Domain && PatchConstOrPrimVectorCount > 0 &&
+      OutputVectorCounts[0] > 0) {
+    uint32_t NumDwords =
+        mapTableSize(PatchConstOrPrimVectorCount, OutputVectorCounts[0]);
+    size_t NumBytes = NumDwords * sizeof(uint32_t);
+    PatchOutputMap.Data = Data.substr(Current - Data.begin(), NumBytes);
+    Current += NumBytes;
+  }
+
   return Error::success();
 }
 

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -219,6 +219,26 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
             El.Allocated, El.Kind, El.Type, El.Mode, El.DynamicMask,
             El.Stream});
 
+      static_assert(PSV.OutputVectorMasks.size() == PSV.InputOutputMap.size());
+      for (unsigned I = 0; I < PSV.OutputVectorMasks.size(); ++I) {
+        PSV.OutputVectorMasks[I].insert(PSV.OutputVectorMasks[I].begin(),
+                                        P.Info->OutputVectorMasks[I].begin(),
+                                        P.Info->OutputVectorMasks[I].end());
+        PSV.InputOutputMap[I].insert(PSV.InputOutputMap[I].begin(),
+                                     P.Info->InputOutputMap[I].begin(),
+                                     P.Info->InputOutputMap[I].end());
+      }
+
+      PSV.PatchOrPrimMasks.insert(PSV.PatchOrPrimMasks.begin(),
+                                  P.Info->PatchOrPrimMasks.begin(),
+                                  P.Info->PatchOrPrimMasks.end());
+      PSV.InputPatchMap.insert(PSV.InputPatchMap.begin(),
+                               P.Info->InputPatchMap.begin(),
+                               P.Info->InputPatchMap.end());
+      PSV.PatchOutputMap.insert(PSV.PatchOutputMap.begin(),
+                                P.Info->PatchOutputMap.begin(),
+                                P.Info->PatchOutputMap.end());
+
       PSV.finalize(static_cast<Triple::EnvironmentType>(
           Triple::Pixel + P.Info->Info.ShaderStage));
       PSV.write(OS, P.Info->Version);

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -139,6 +139,24 @@ void MappingTraits<DXContainerYAML::PSVInfo>::mapping(
   IO.mapRequired("SigInputElements", PSV.SigInputElements);
   IO.mapRequired("SigOutputElements", PSV.SigOutputElements);
   IO.mapRequired("SigPatchOrPrimElements", PSV.SigPatchOrPrimElements);
+
+  Triple::EnvironmentType Stage = dxbc::getShaderStage(PSV.Info.ShaderStage);
+  if (PSV.Info.UsesViewID) {
+    MutableArrayRef<SmallVector<llvm::yaml::Hex32>> MutableOutMasks(
+        PSV.OutputVectorMasks);
+    IO.mapRequired("OutputVectorMasks", MutableOutMasks);
+    if (Stage == Triple::EnvironmentType::Hull)
+      IO.mapRequired("PatchOrPrimMasks", PSV.PatchOrPrimMasks);
+  }
+  MutableArrayRef<SmallVector<llvm::yaml::Hex32>> MutableIOMap(
+      PSV.InputOutputMap);
+  IO.mapRequired("InputOutputMap", MutableIOMap);
+
+  if (Stage == Triple::EnvironmentType::Hull)
+    IO.mapRequired("InputPatchMap", PSV.InputPatchMap);
+
+  if (Stage == Triple::EnvironmentType::Domain)
+    IO.mapRequired("PatchOutputMap", PSV.PatchOutputMap);
 }
 
 void MappingTraits<DXContainerYAML::Part>::mapping(IO &IO,