[DX][ObjYAML] Support for parsing DXIL part

[beanz]

Summary

This patch adds support for parsing the DXIL part data into the
ObjectYAML tooling.

The DXIL part has additional headers describing the shader and bitcode
data and stores serialized bitcode after the headers.

Depends on D124945

[kuhar]

llvm/include/llvm/BinaryFormat/DXContainer.h

100

In bytes or dwords?

101

I'm confused by this comment. Did you mean [BitcodeHeader.Size]? Also what is the content of this blob? I assume that it's not the bitcode, because it's starts after .Offset bytes, right?

110

nit: add an empty line between structs?

115

Maybe: // Number of uint32_t words including this header ?

125

It seems to me that this relies on the structs being tightly packed. But does it always have to be the case? I thought that the compiler is allowed to insert padding arbitrarily as long as all fields remain aligned.

In the serialization/deserialization code I worked on in the past, we always relied on the size of individual fields for this reason, and always read/wrote all structs field wise.

llvm/include/llvm/Object/DXContainer.h

118

This performs a copy, doesn't it? I guess that this may be desired because ProgramHeader is relatively small, right?

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

108

nit: can you add an empty line in between functions?

123–160

Consider an early exit instead: https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code

155–156

Can we use OS.write_zeros(PadBytes) instead?

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

37–38

Is the whole blob necessary? Could we replace it with something shorter that is not a complete DXIL program but starts with the right magic only etc?

[beanz]

llvm/include/llvm/BinaryFormat/DXContainer.h

100

Bytes here, I'll be explicit in the comment.

101

Doh... Yes. The Offset is the offset from the start of the header.

The DXContainer format overly defines every offset and size.

125

In C struct members are "aligned in an implementation-defined manner appropriate to its type". The compiler can't arbitrarily add padding, but it can to ensure appropriate alignment. I explicitly added an Unused member to represent the padding that is expected, and that should make the remaining structure layout consistent with every architecture I'm aware of.

The assert is just intended to verify my bit math (and that none of our other architectures do anything weird) because code depends on it.

For binary file formats like object files we pretty regularly rely on structure definitions being C-standard layout. This is consistent throughout the BinaryFormat library.

llvm/include/llvm/Object/DXContainer.h

118

Yea, I don't love this. The ProgramHeader is small so the copy shouldn't hurt too much.

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

123–160

I didn't want to do an early exit here because this will grow to cover more parts as I add support for them.

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

37–38

Yea, for the purposes of this test I can do a simplified case.

[kuhar]

llvm/include/llvm/BinaryFormat/DXContainer.h

125

I think that's fair. It's very unlikely for an implementation to over-align fields and a lot of other code would break too.

llvm/lib/Object/DXContainer.cpp

13

nit: cstddef? I think we generally prefer C++ versions of C headers

[beanz]

Updates based on review feedback.

Thank you @kuhar!

[kuhar]

LGTM

llvm/include/llvm/BinaryFormat/DXContainer.h

struct PartHeader {
 
   void swapBytes() { sys::swapByteOrder(Size); }
   StringRef getName() const {
     return StringRef(reinterpret_cast<const char *>(&Name[0]), 4);
   }
   // Structure is followed directly by part data: uint8_t PartData[PartSize].
 };
 
+struct BitcodeHeader {
+  uint8_t Magic[4];     // ACSII "DXIL".
+  uint8_t MajorVersion; // DXIL version.
+  uint8_t MinorVersion; // DXIL version.
+  uint16_t Unused;
+  uint32_t Offset; // Offset to LLVM bitcode (from start of header).
+  uint32_t Size;   // Size of LLVM bitcode (in bytes).
+  // Followed by uint8_t[BitcodeHeader.Size] at &BitcodeHeader + Header.Offset
+
+  void swapBytes() {
+    sys::swapByteOrder(MinorVersion);
+    sys::swapByteOrder(MajorVersion);
+    sys::swapByteOrder(Offset);
+    sys::swapByteOrder(Size);
+  }
+};
+
+struct ProgramHeader {
+  uint8_t MinorVersion : 4;
+  uint8_t MajorVersion : 4;
+  uint8_t Unused;
+  uint16_t ShaderKind;
+  uint32_t Size; // Size in uint32_t words including this header.
+  BitcodeHeader Bitcode;
+
+  void swapBytes() {
+    sys::swapByteOrder(ShaderKind);
+    sys::swapByteOrder(Size);
+    Bitcode.swapBytes();
+  }
+};
+
+static_assert(sizeof(ProgramHeader) == 24, "ProgramHeader Size incorrect!");
+
 } // namespace dxbc
 } // namespace llvm
 
 #endif // LLVM_BINARYFORMAT_DXCONTAINER_H

llvm/include/llvm/Object/DXContainer.h

 // interface for DXContainer files.
 //
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_OBJECT_DXCONTAINER_H
 #define LLVM_OBJECT_DXCONTAINER_H
 
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBufferRef.h"
 
 namespace llvm {
 namespace object {
 class DXContainer {
+public:
+  using DXILData = std::pair<dxbc::ProgramHeader, const char *>;
+
 private:
   DXContainer(MemoryBufferRef O);
 
   MemoryBufferRef Data;
   dxbc::Header Header;
   SmallVector<uint32_t, 4> PartOffsets;
+  Optional<DXILData> DXIL;
 
   Error parseHeader();
   Error parsePartOffsets();
+  Error parseDXILHeader(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:
   }
 
   PartIterator end() const { return PartIterator(*this, PartOffsets.end()); }
 
   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; }
 };
 
 } // namespace object
 } // namespace llvm
 
 #endif // LLVM_OBJECT_DXCONTAINERFILE_H

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

 struct FileHeader {
   std::vector<llvm::yaml::Hex8> Hash;
   VersionTuple Version;
   Optional<uint32_t> FileSize;
   uint32_t PartCount;
   Optional<std::vector<uint32_t>> PartOffsets;
 };
 
+struct DXILProgram {
+  uint8_t MajorVersion;
+  uint8_t MinorVersion;
+  uint16_t ShaderKind;
+  Optional<uint32_t> Size;
+  uint16_t DXILMajorVersion;
+  uint16_t DXILMinorVersion;
+  Optional<uint32_t> DXILOffset;
+  Optional<uint32_t> DXILSize;
+  Optional<std::vector<llvm::yaml::Hex8>> DXIL;
+};
+
 struct Part {
   std::string Name;
   uint32_t Size;
+  Optional<DXILProgram> Program;
 };
 
 struct Object {
   FileHeader Header;
   std::vector<Part> Parts;
 };
 
 } // namespace DXContainerYAML

 template <> struct MappingTraits<DXContainerYAML::VersionTuple> {
   static void mapping(IO &IO, DXContainerYAML::VersionTuple &Version);
 };
 
 template <> struct MappingTraits<DXContainerYAML::FileHeader> {
   static void mapping(IO &IO, DXContainerYAML::FileHeader &Header);
 };
 
+template <> struct MappingTraits<DXContainerYAML::DXILProgram> {
+  static void mapping(IO &IO, DXContainerYAML::DXILProgram &Program);
+};
+
 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/Object/DXContainer.cpp

 //===- DXContainer.cpp - DXContainer object file implementation -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Object/DXContainer.h"
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Object/Error.h"
 
 using namespace llvm;
 using namespace llvm::object;
 
 static Error parseFailed(const Twine &Msg) {
   return make_error<GenericBinaryError>(Msg.str(), object_error::parse_failed);
 }
 
 template <typename T>
 static Error readStruct(StringRef Buffer, const char *Src, T &Struct) {

 }
 
 DXContainer::DXContainer(MemoryBufferRef O) : Data(O) {}
 
 Error DXContainer::parseHeader() {
   return readStruct(Data.getBuffer(), Data.getBuffer().data(), Header);
 }
 
+Error DXContainer::parseDXILHeader(uint32_t Offset) {
+  if (DXIL)
+    return parseFailed("More than one DXIL part is present in the file");
+  const char *Current = Data.getBuffer().data() + Offset;
+  dxbc::ProgramHeader Header;
+  if (Error Err = readStruct(Data.getBuffer(), Current, Header))
+    return Err;
+  Current += offsetof(dxbc::ProgramHeader, Bitcode) + Header.Bitcode.Offset;
+  DXIL.emplace(std::make_pair(Header, Current));
+  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
     // header. To prevent overflow, we subtract the part header size from the
     // buffer size, rather than adding to the offset. Since the file header is
     // larger than the part header we can't reach this code unless the buffer
     // is larger than the part header, so this can't underflow.
     if (PartOffset > Data.getBufferSize() - sizeof(dxbc::PartHeader))
       return parseFailed("Part offset points beyond boundary of the file");
     PartOffsets.push_back(PartOffset);
+
+    // If this isn't a dxil part stop here...
+    if (Data.getBuffer().substr(PartOffset, 4) != "DXIL")
+      continue;
+    if (Error Err = parseDXILHeader(PartOffset + sizeof(dxbc::PartHeader)))
+      return Err;
   }
   return Error::success();
 }
 
 Expected<DXContainer> DXContainer::create(MemoryBufferRef Object) {
   DXContainer Container(Object);
   if (Error Err = Container.parseHeader())
     return std::move(Err);

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

void DXContainerWriter::writeHeader(raw_ostream &OS) {
   SmallVector<uint32_t> Offsets(ObjectFile.Header.PartOffsets->begin(),
                                 ObjectFile.Header.PartOffsets->end());
   if (sys::IsBigEndianHost)
     for (auto &O : Offsets)
       sys::swapByteOrder(O);
   OS.write(reinterpret_cast<char *>(Offsets.data()),
            Offsets.size() * sizeof(uint32_t));
 }
+
 void DXContainerWriter::writeParts(raw_ostream &OS) {
   uint32_t RollingOffset =
       sizeof(dxbc::Header) + (ObjectFile.Header.PartCount * sizeof(uint32_t));
   for (auto I : llvm::zip(ObjectFile.Parts, *ObjectFile.Header.PartOffsets)) {
     if (RollingOffset < std::get<1>(I)) {
       uint32_t PadBytes = std::get<1>(I) - RollingOffset;
-      std::vector<uint8_t> FillData(PadBytes, 0);
-      OS.write(reinterpret_cast<char *>(FillData.data()), PadBytes);
+      OS.write_zeros(PadBytes);
     }
     DXContainerYAML::Part P = std::get<0>(I);
     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));
     RollingOffset = std::get<1>(I) + sizeof(dxbc::PartHeader);
 
-    // TODO: Write Part data
+    if (P.Name == "DXIL" && P.Program) {
+      dxbc::ProgramHeader Header;
+      Header.MajorVersion = P.Program->MajorVersion;
+      Header.MinorVersion = P.Program->MinorVersion;
+      Header.ShaderKind = P.Program->ShaderKind;
+      memcpy(Header.Bitcode.Magic, "DXIL", 4);
+      Header.Bitcode.MajorVersion = P.Program->DXILMajorVersion;
+      Header.Bitcode.MinorVersion = P.Program->DXILMinorVersion;
+
+      // Compute the optional fields if needed...
+      if (P.Program->DXILOffset)
+        Header.Bitcode.Offset = P.Program->DXILOffset.getValue();
+      else
+        Header.Bitcode.Offset = sizeof(dxbc::BitcodeHeader);
+
+      if (P.Program->DXILSize)
+        Header.Bitcode.Size = P.Program->DXILSize.getValue();
+      else
+        Header.Bitcode.Size = P.Program->DXIL ? P.Program->DXIL->size() : 0;
+
+      if (P.Program->Size)
+        Header.Size = P.Program->Size.getValue();
+      else
+        Header.Size = sizeof(dxbc::ProgramHeader) + Header.Bitcode.Size;
+      if (sys::IsBigEndianHost)
+        Header.swapBytes();
+      OS.write(reinterpret_cast<const char *>(&Header),
+               sizeof(dxbc::ProgramHeader));
+      if (P.Program->DXIL) {
+        if (Header.Bitcode.Offset > sizeof(dxbc::BitcodeHeader)) {
+          uint32_t PadBytes =
+              Header.Bitcode.Offset - sizeof(dxbc::BitcodeHeader);
+          OS.write_zeros(PadBytes);
+        }
+        OS.write(reinterpret_cast<char *>(P.Program->DXIL->data()),
+                 P.Program->DXIL->size());
+      }
+    }
   }
 }
 
 Error DXContainerWriter::write(raw_ostream &OS) {
   if (Error Err = computePartOffsets())
     return Err;
   writeHeader(OS);
   writeParts(OS);

llvm/lib/ObjectYAML/DXContainerYAML.cpp

void MappingTraits<DXContainerYAML::FileHeader>::mapping(
     IO &IO, DXContainerYAML::FileHeader &Header) {
   IO.mapRequired("Hash", Header.Hash);
   IO.mapRequired("Version", Header.Version);
   IO.mapOptional("FileSize", Header.FileSize);
   IO.mapRequired("PartCount", Header.PartCount);
   IO.mapOptional("PartOffsets", Header.PartOffsets);
 }
 
+void MappingTraits<DXContainerYAML::DXILProgram>::mapping(
+    IO &IO, DXContainerYAML::DXILProgram &Program) {
+  IO.mapRequired("MajorVersion", Program.MajorVersion);
+  IO.mapRequired("MinorVersion", Program.MinorVersion);
+  IO.mapRequired("ShaderKind", Program.ShaderKind);
+  IO.mapOptional("Size", Program.Size);
+  IO.mapRequired("DXIMMajorVersion", Program.DXILMajorVersion);
+  IO.mapRequired("DXILMinorVersion", Program.DXILMinorVersion);
+  IO.mapOptional("DXILSize", Program.DXILSize);
+  IO.mapOptional("DXIL", Program.DXIL);
+}
+
 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);
 }
 
 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/DXILPart.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:        3548
+  PartCount:       7
+  PartOffsets:     [ 60, 76, 92, 108, 236, 1932, 1960 ]
+Parts:
+  - Name:            SFI0
+    Size:            8
+  - Name:            ISG1
+    Size:            8
+  - Name:            OSG1
+    Size:            8
+  - Name:            PSV0
+    Size:            120
+  - Name:            STAT
+    Size:            1688
+  - Name:            HASH
+    Size:            20
+  - Name:            DXIL
+    Size:            28
+    Program:
+      MajorVersion:    6
+      MinorVersion:    5
+      ShaderKind:      5
+      Size:            8
+      DXIMMajorVersion: 1
+      DXILMinorVersion: 5
+      DXILSize:        4
+      DXIL:            [ 0x42, 0x43, 0xC0, 0xDE, ]
+...
+
+
+
+
+#CHECK:        - Name:            DXIL
+#CHECK-NEXT:     Size:            28
+#CHECK-NEXT:     Program:
+#CHECK-NEXT:       MajorVersion:    6
+#CHECK-NEXT:       MinorVersion:    5
+#CHECK-NEXT:       ShaderKind:      5
+#CHECK-NEXT:       Size:            8
+#CHECK-NEXT:       DXIMMajorVersion: 1
+#CHECK-NEXT:       DXILMinorVersion: 5
+#CHECK-NEXT:       DXILSize:        4
+#CHECK-NEXT:       DXIL:            [ 0x42, 0x43, 0xC0, 0xDE

llvm/tools/obj2yaml/dxcontainer2yaml.cpp

dumpDXContainer(MemoryBufferRef Source) {
   Obj->Header.Version.Major = Container.getHeader().Version.Major;
   Obj->Header.Version.Minor = Container.getHeader().Version.Minor;
   Obj->Header.FileSize = Container.getHeader().FileSize;
   Obj->Header.PartCount = Container.getHeader().PartCount;
 
   Obj->Header.PartOffsets = std::vector<uint32_t>();
   for (const auto P : Container) {
     Obj->Header.PartOffsets->push_back(P.Offset);
+    if (P.Part.getName() == "DXIL") {
+      Optional<DXContainer::DXILData> DXIL = Container.getDXIL();
+      assert(DXIL.hasValue() && "Since we are iterating and found a DXIL part, "
+                                "this should never not have a value");
+      Obj->Parts.push_back(DXContainerYAML::Part{
+          P.Part.getName().str(), P.Part.Size,
+          DXContainerYAML::DXILProgram{
+              DXIL->first.MajorVersion, DXIL->first.MinorVersion,
+              DXIL->first.ShaderKind, DXIL->first.Size,
+              DXIL->first.Bitcode.MajorVersion,
+              DXIL->first.Bitcode.MinorVersion, DXIL->first.Bitcode.Offset,
+              DXIL->first.Bitcode.Size,
+              std::vector<llvm::yaml::Hex8>(
+                  DXIL->second, DXIL->second + DXIL->first.Bitcode.Size)}});
+    } else {
       Obj->Parts.push_back(
-        DXContainerYAML::Part{P.Part.getName().str(), P.Part.Size});
+          DXContainerYAML::Part{P.Part.getName().str(), P.Part.Size, None});
+    }
   }
 
   return Obj.release();
 }
 
 llvm::Error dxcontainer2yaml(llvm::raw_ostream &Out,
                              llvm::MemoryBufferRef Source) {
   Expected<DXContainerYAML::Object *> YAMLOrErr = dumpDXContainer(Source);