[ObjectYAML] Add support for DXContainer HASH

[beanz]

Summary

DXContainer files contain a part that has an MD5 of the generated
shader. This adds support to the ObjectYAML tooling to expand the hash
part data and hash iteself in preparation for adding hashing support to
DirectX code generation.

[python3kgae]

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

70

Why use vector here when the size is always 16?

[beanz]

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

70

We can't bind yaml traits to statically sized arrays because YAML doesn't have statically sized arrays.

[python3kgae]

llvm/lib/BinaryFormat/DXContainer.cpp

28

add const to Zeros?

29

Why do we want to set Digest to Zeros when Flags is 0?

[beanz]

llvm/lib/BinaryFormat/DXContainer.cpp

29

This doesn't set, it compares. If flags and the digest are all zero'd it wasn't populated, so we can omit it from the textual YAML output.

llvm/include/llvm/BinaryFormat/DXContainer.h

enum class HashFlags : uint32_t {
   IncludesSource = 1, // This flag indicates that the shader hash was computed
                       // taking into account source information (-Zss)
 };
 
 struct ShaderHash {
   uint32_t Flags; // dxbc::HashFlags
   uint8_t Digest[16];
 
+  bool isPopulated();
+
   void swapBytes() { sys::swapByteOrder(Flags); }
 };
 
 struct ContainerVersion {
   uint16_t Major;
   uint16_t Minor;
 
   void swapBytes() {

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

 
 #ifdef CONTAINER_PART
 CONTAINER_PART(DXIL)
 CONTAINER_PART(SFI0)
+CONTAINER_PART(HASH)
 
 #undef CONTAINER_PART
 #endif 
 
 #ifdef SHADER_FLAG
 
 SHADER_FLAG(0, Doubles, "Double-precision floating point")
 SHADER_FLAG(1, ComputeShadersPlusRawAndStructuredBuffers, "Raw and Structured buffers")

llvm/include/llvm/Object/DXContainer.h

 private:
   DXContainer(MemoryBufferRef O);
 
   MemoryBufferRef Data;
   dxbc::Header Header;
   SmallVector<uint32_t, 4> PartOffsets;
   Optional<DXILData> DXIL;
   Optional<uint64_t> ShaderFlags;
+  Optional<dxbc::ShaderHash> Hash;
 
   Error parseHeader();
   Error parsePartOffsets();
   Error parseDXILHeader(uint32_t Offset);
   Error parseShaderFlags(uint32_t Offset);
+  Error parseHash(uint32_t Offset);
   friend class PartIterator;
 
 public:
   // The PartIterator is a wrapper around the iterator for the PartOffsets
   // member of the DXContainer. It contains a refernce to the container, and the
   // current iterator value, as well as storage for a parsed part header.
   class PartIterator {
     const DXContainer &Container;
public:
   StringRef getData() const { return Data.getBuffer(); }
   static Expected<DXContainer> create(MemoryBufferRef Object);
 
   const dxbc::Header &getHeader() const { return Header; }
 
   Optional<DXILData> getDXIL() const { return DXIL; }
 
   Optional<uint64_t> getShaderFlags() const { return ShaderFlags; }
+
+  Optional<dxbc::ShaderHash> getShaderHash() const { return Hash; }
 };
 
 } // namespace object
 } // namespace llvm
 
 #endif // LLVM_OBJECT_DXCONTAINER_H

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

 /// of DXContainer.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_OBJECTYAML_DXCONTAINERYAML_H
 #define LLVM_OBJECTYAML_DXCONTAINERYAML_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/ObjectYAML/YAML.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <cstdint>
 #include <string>
 #include <vector>
 
 namespace llvm {
 namespace DXContainerYAML {

 #define SHADER_FLAG(Num, Val, Str) bool Val = false;
 struct ShaderFlags {
   ShaderFlags() = default;
   ShaderFlags(uint64_t FlagData);
   uint64_t getEncodedFlags();
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 };
 
+struct ShaderHash {
+  ShaderHash() = default;
+  ShaderHash(const dxbc::ShaderHash &Data);
+
+  bool IncludesSource;
+  std::vector<llvm::yaml::Hex8> Digest;
+};
+
 struct Part {
   Part() = default;
   Part(std::string N, uint32_t S) : Name(N), Size(S) {}
   std::string Name;
   uint32_t Size;
   Optional<DXILProgram> Program;
   Optional<ShaderFlags> Flags;
+  Optional<ShaderHash> Hash;
 };
 
 struct Object {
   FileHeader Header;
   std::vector<Part> Parts;
 };
 
 } // namespace DXContainerYAML

 template <> struct MappingTraits<DXContainerYAML::DXILProgram> {
   static void mapping(IO &IO, DXContainerYAML::DXILProgram &Program);
 };
 
 template <> struct MappingTraits<DXContainerYAML::ShaderFlags> {
   static void mapping(IO &IO, DXContainerYAML::ShaderFlags &Flags);
 };
 
+template <> struct MappingTraits<DXContainerYAML::ShaderHash> {
+  static void mapping(IO &IO, DXContainerYAML::ShaderHash &Hash);
+};
+
 template <> struct MappingTraits<DXContainerYAML::Part> {
   static void mapping(IO &IO, DXContainerYAML::Part &Version);
 };
 
 template <> struct MappingTraits<DXContainerYAML::Object> {
   static void mapping(IO &IO, DXContainerYAML::Object &Obj);
 };
 
 } // namespace yaml
 
 } // namespace llvm
 
 #endif // LLVM_OBJECTYAML_DXCONTAINERYAML_H

llvm/lib/BinaryFormat/DXContainer.cpp

 // This file contains utility functions for working with DXContainers.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/ADT/StringSwitch.h"
 
 using namespace llvm;
-using namespace llvm;
+using namespace llvm::dxbc;
 
 dxbc::PartType dxbc::parsePartType(StringRef S) {
 #define CONTAINER_PART(PartName) .Case(#PartName, PartType::PartName)
   return StringSwitch<dxbc::PartType>(S)
 #include "llvm/BinaryFormat/DXContainerConstants.def"
       .Default(dxbc::PartType::Unknown);
 }
+
+bool ShaderHash::isPopulated() {
+  static uint8_t Zeros[16] = {0};
+  return Flags > 0 || 0 != memcmp(&Digest, &Zeros, 16);
+}

llvm/lib/Object/DXContainer.cpp

Error DXContainer::parseShaderFlags(uint32_t Offset) {
   const char *Current = Data.getBuffer().data() + Offset;
   uint64_t FlagValue = 0;
   if (Error Err = readInteger(Data.getBuffer(), Current, FlagValue))
     return Err;
   ShaderFlags = FlagValue;
   return Error::success();
 }
 
+Error DXContainer::parseHash(uint32_t Offset) {
+  if (Hash)
+    return parseFailed("More than one HASH part is present in the file");
+  const char *Current = Data.getBuffer().data() + Offset;
+  dxbc::ShaderHash ReadHash;
+  if (Error Err = readStruct(Data.getBuffer(), Current, ReadHash))
+    return Err;
+  Hash = ReadHash;
+  return Error::success();
+}
+
 Error DXContainer::parsePartOffsets() {
   const char *Current = Data.getBuffer().data() + sizeof(dxbc::Header);
   for (uint32_t Part = 0; Part < Header.PartCount; ++Part) {
     uint32_t PartOffset;
     if (Error Err = readInteger(Data.getBuffer(), Current, PartOffset))
       return Err;
     Current += sizeof(uint32_t);
     // We need to ensure that each part offset leaves enough space for a part
for (uint32_t Part = 0; Part < Header.PartCount; ++Part) {
     case dxbc::PartType::DXIL:
       if (Error Err = parseDXILHeader(PartOffset + sizeof(dxbc::PartHeader)))
         return Err;
       break;
     case dxbc::PartType::SFI0:
       if (Error Err = parseShaderFlags(PartOffset + sizeof(dxbc::PartHeader)))
         return Err;
       break;
+    case dxbc::PartType::HASH:
+      if (Error Err = parseHash(PartOffset + sizeof(dxbc::PartHeader)))
+        return Err;
+      break;
     case dxbc::PartType::Unknown:
       break;
     }
   }
   return Error::success();
 }
 
 Expected<DXContainer> DXContainer::create(MemoryBufferRef Object) {

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

for (auto I : llvm::zip(ObjectFile.Parts, *ObjectFile.Header.PartOffsets)) {
     OS.write(P.Name.c_str(), 4);
     if (sys::IsBigEndianHost)
       sys::swapByteOrder(P.Size);
     OS.write(reinterpret_cast<const char *>(&P.Size), sizeof(uint32_t));
 
     dxbc::PartType PT = dxbc::parsePartType(P.Name);
 
     uint64_t DataStart = OS.tell();
-    if (PT == dxbc::PartType::DXIL) {
+    switch (PT) {
+    case dxbc::PartType::DXIL: {
       if (!P.Program)
         continue;
       dxbc::ProgramHeader Header;
       Header.MajorVersion = P.Program->MajorVersion;
       Header.MinorVersion = P.Program->MinorVersion;
       Header.Unused = 0;
       Header.ShaderKind = P.Program->ShaderKind;
       memcpy(Header.Bitcode.Magic, "DXIL", 4);
case dxbc::PartType::DXIL: {
       if (P.Program->DXIL) {
         if (BitcodeOffset > sizeof(dxbc::BitcodeHeader)) {
           uint32_t PadBytes = BitcodeOffset - sizeof(dxbc::BitcodeHeader);
           OS.write_zeros(PadBytes);
         }
         OS.write(reinterpret_cast<char *>(P.Program->DXIL->data()),
                  P.Program->DXIL->size());
       }
-    } else if (PT == dxbc::PartType::SFI0) {
+      break;
+    }
+    case dxbc::PartType::SFI0: {
       // If we don't have any flags we can continue here and the data will be
       // zeroed out.
       if (!P.Flags.has_value())
         continue;
       uint64_t Flags = P.Flags->getEncodedFlags();
       if (sys::IsBigEndianHost)
         sys::swapByteOrder(Flags);
       OS.write(reinterpret_cast<char *>(&Flags), sizeof(uint64_t));
+      break;
+    }
+    case dxbc::PartType::HASH: {
+      if (!P.Hash.has_value())
+        continue;
+      dxbc::ShaderHash Hash = {0, {0}};
+      if (P.Hash->IncludesSource)
+        Hash.Flags |= static_cast<uint32_t>(dxbc::HashFlags::IncludesSource);
+      memcpy(&Hash.Digest[0], &P.Hash->Digest[0], 16);
+      if (sys::IsBigEndianHost)
+        Hash.swapBytes();
+      OS.write(reinterpret_cast<char *>(&Hash), sizeof(dxbc::ShaderHash));
+      break;
+    }
+    case dxbc::PartType::Unknown:
+      break; // Skip any handling for unrecognized parts.
     }
     uint64_t BytesWritten = OS.tell() - DataStart;
     RollingOffset += BytesWritten;
     if (BytesWritten < PartSize)
       OS.write_zeros(PartSize - BytesWritten);
     RollingOffset += PartSize;
   }
 }

llvm/lib/ObjectYAML/DXContainerYAML.cpp

uint64_t DXContainerYAML::ShaderFlags::getEncodedFlags() {
   uint64_t Flag = 0;
 #define SHADER_FLAG(Num, Val, Str)                                             \
   if (Val)                                                                     \
     Flag |= (uint64_t)dxbc::FeatureFlags::Val;
 #include "llvm/BinaryFormat/DXContainerConstants.def"
   return Flag;
 }
 
+DXContainerYAML::ShaderHash::ShaderHash(const dxbc::ShaderHash &Data)
+    : IncludesSource((Data.Flags & static_cast<uint32_t>(
+                                       dxbc::HashFlags::IncludesSource)) != 0),
+      Digest(16, 0) {
+  memcpy(Digest.data(), &Data.Digest[0], 16);
+}
+
 namespace yaml {
 
 void MappingTraits<DXContainerYAML::VersionTuple>::mapping(
     IO &IO, DXContainerYAML::VersionTuple &Version) {
   IO.mapRequired("Major", Version.Major);
   IO.mapRequired("Minor", Version.Minor);
 }
 

 }
 
 void MappingTraits<DXContainerYAML::ShaderFlags>::mapping(
     IO &IO, DXContainerYAML::ShaderFlags &Flags) {
 #define SHADER_FLAG(Num, Val, Str) IO.mapRequired(#Val, Flags.Val);
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
+void MappingTraits<DXContainerYAML::ShaderHash>::mapping(
+    IO &IO, DXContainerYAML::ShaderHash &Hash) {
+  IO.mapRequired("IncludesSource", Hash.IncludesSource);
+  IO.mapRequired("Digest", Hash.Digest);
+}
+
 void MappingTraits<DXContainerYAML::Part>::mapping(IO &IO,
                                                    DXContainerYAML::Part &P) {
   IO.mapRequired("Name", P.Name);
   IO.mapRequired("Size", P.Size);
   IO.mapOptional("Program", P.Program);
   IO.mapOptional("Flags", P.Flags);
+  IO.mapOptional("Hash", P.Hash);
 }
 
 void MappingTraits<DXContainerYAML::Object>::mapping(
     IO &IO, DXContainerYAML::Object &Obj) {
   IO.mapTag("!dxcontainer", true);
   IO.mapRequired("Header", Obj.Header);
   IO.mapRequired("Parts", Obj.Parts);
 }
 
 } // namespace yaml
 } // namespace llvm

llvm/test/tools/obj2yaml/DXContainer/ExplicitSizeAndOffsets.yaml

 # 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:        172
+  FileSize:        184
   PartCount:       7
-  PartOffsets:     [ 60, 76, 92, 108, 124, 140, 156 ]
+  PartOffsets:     [ 60, 76, 92, 108, 124, 140, 168 ]
 Parts:
   - Name:            SFI0
     Size:            8
   - Name:            ISG1
     Size:            8
   - Name:            OSG1
     Size:            8
   - Name:            PSV0
     Size:            8
   - Name:            STAT
     Size:            8
   - Name:            HASH
-    Size:            8
+    Size:            20
   - Name:            CXIL
     Size:            8
 ...
 
 # CHECK: --- !dxcontainer
 # CHECK-NEXT: Header:
 # CHECK-NEXT:   Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 
 # CHECK-NEXT:                      0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
 # CHECK-NEXT:   Version:
 # CHECK-NEXT:     Major:           1
 # CHECK-NEXT:     Minor:           0
-# CHECK-NEXT:   FileSize:        172
+# CHECK-NEXT:   FileSize:        184
 # CHECK-NEXT:   PartCount:       7
-# CHECK-NEXT:   PartOffsets:     [ 60, 76, 92, 108, 124, 140, 156 ]
+# CHECK-NEXT:   PartOffsets:     [ 60, 76, 92, 108, 124, 140, 168 ]
 # CHECK-NEXT: Parts:
 # CHECK-NEXT:   - Name:            SFI0
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            ISG1
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            OSG1
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            PSV0
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            STAT
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            HASH
-# CHECK-NEXT:     Size:            8
+# CHECK-NEXT:     Size:            20
 # CHECK-NEXT:   - Name:            CXIL
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT: ...

llvm/test/tools/obj2yaml/DXContainer/OmitSizeAndOffsets.yaml

- Name:            ISG1
     Size:            8
   - Name:            OSG1
     Size:            8
   - Name:            PSV0
     Size:            8
   - Name:            STAT
     Size:            8
   - Name:            HASH
-    Size:            8
+    Size:            20
   - Name:            CXIL
     Size:            8
 ...
 
 # CHECK: --- !dxcontainer
 # CHECK-NEXT: Header:
 # CHECK-NEXT:   Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 
 # CHECK-NEXT:                      0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
 # CHECK-NEXT:   Version:
 # CHECK-NEXT:     Major:           1
 # CHECK-NEXT:     Minor:           0
-# CHECK-NEXT:   FileSize:        172
+# CHECK-NEXT:   FileSize:        184
 # CHECK-NEXT:   PartCount:       7
-# CHECK-NEXT:   PartOffsets:     [ 60, 76, 92, 108, 124, 140, 156 ]
+# CHECK-NEXT:   PartOffsets:     [ 60, 76, 92, 108, 124, 140, 168 ]
 # CHECK-NEXT: Parts:
 # CHECK-NEXT:   - Name:            SFI0
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            ISG1
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            OSG1
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            PSV0
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            STAT
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:   - Name:            HASH
-# CHECK-NEXT:     Size:            8
+# CHECK-NEXT:     Size:            20
 # CHECK-NEXT:   - Name:            CXIL
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT: ...

llvm/test/tools/obj2yaml/DXContainer/ShaderFlags.yaml

- Name:            ISG1
     Size:            8
   - Name:            OSG1
     Size:            8
   - Name:            PSV0
     Size:            8
   - Name:            STAT
     Size:            8
   - Name:            HASH
-    Size:            8
+    Size:            20
   - Name:            CXIL
     Size:            8
 ...
 
 # CHECK: Parts:
 # CHECK-NEXT:   - Name:            SFI0
 # CHECK-NEXT:     Size:            8
 # CHECK-NEXT:     Flags:

llvm/test/tools/obj2yaml/DXContainer/ShaderFlagsEmpty.yaml

- Name:            ISG1
     Size:            8
   - Name:            OSG1
     Size:            8
   - Name:            PSV0
     Size:            8
   - Name:            STAT
     Size:            8
   - Name:            HASH
-    Size:            8
+    Size:            20
   - Name:            CXIL
     Size:            8
 ...
 
 # CHECK: Parts:
 # CHECK-NEXT:   - Name:            SFI0
 # CHECK-NEXT:     Size:            8
 # CHECK-NOT:      Flags:
 # CHECK-NEXT:   - Name:            ISG1

llvm/test/tools/obj2yaml/DXContainer/ShaderHash.yaml

+# 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
+  PartCount:       1
+Parts:
+  - Name:            HASH
+    Size:            20
+    Hash:
+      IncludesSource:  true
+      Digest:          [ 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 
+                         0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80 ]
+...
+
+# CHECK: --- !dxcontainer
+# CHECK-NEXT: Header:
+# CHECK-NEXT:   Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 
+# CHECK-NEXT:                      0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+# CHECK:        FileSize:        64
+# CHECK-NEXT:   PartCount:       1
+# CHECK-NEXT:   PartOffsets:     [ 36 ]
+# CHECK:        - Name:            HASH
+# CHECK-NEXT:     Size:            20
+# CHECK-NEXT:     Hash:
+# CHECK-NEXT:       IncludesSource:  true
+# CHECK-NEXT:       Digest:          [ 0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x10,
+# CHECK-NEXT:                          0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80 ]
+# CHECK-NEXT: ...

llvm/tools/obj2yaml/dxcontainer2yaml.cpp

for (const auto P : Container) {
     }
     case dxbc::PartType::SFI0: {
       Optional<uint64_t> Flags = Container.getShaderFlags();
       // Omit the flags in the YAML if they are missing or zero.
       if (Flags && *Flags > 0)
         NewPart.Flags = DXContainerYAML::ShaderFlags(*Flags);
       break;
     }
+    case dxbc::PartType::HASH: {
+      Optional<dxbc::ShaderHash> Hash = Container.getShaderHash();
+      if (Hash && Hash->isPopulated())
+        NewPart.Hash = DXContainerYAML::ShaderHash(*Hash);
+      break;
+    }
     case dxbc::PartType::Unknown:
       break;
     }
   }
 
   return Obj.release();
 }