Extend llvm objdump fatbin (llvm#140286)

Utilize the new extensions to the LLVM Offloading API to extend to
llvm-objdump to handle dumping fatbin offload bundles generated by HIP.
This extension to llvm-objdump adds the option --offload-fatbin.
Specifying this option will take the input object/executable and extract
all offload fatbin bundle entries into distinct code object files with
names reflecting the source file name combined with the Bundle Entry ID.
Users can also use the --arch-name option to filter offload fatbin
bundle entries by their target triple.

---------

Co-authored-by: dsalinas <[email protected]>

Author: 2 people

llvm/include/llvm/Object/OffloadBinary.h

@@ -51,31 +51,6 @@ enum ImageKind : uint16_t {
   IMG_LAST,
 };
 
-class CompressedOffloadBundle {
-private:
-  static inline const size_t MagicSize = 4;
-  static inline const size_t VersionFieldSize = sizeof(uint16_t);
-  static inline const size_t MethodFieldSize = sizeof(uint16_t);
-  static inline const size_t FileSizeFieldSize = sizeof(uint32_t);
-  static inline const size_t UncompressedSizeFieldSize = sizeof(uint32_t);
-  static inline const size_t HashFieldSize = sizeof(uint64_t);
-  static inline const size_t V1HeaderSize =
-      MagicSize + VersionFieldSize + MethodFieldSize +
-      UncompressedSizeFieldSize + HashFieldSize;
-  static inline const size_t V2HeaderSize =
-      MagicSize + VersionFieldSize + FileSizeFieldSize + MethodFieldSize +
-      UncompressedSizeFieldSize + HashFieldSize;
-  static inline const llvm::StringRef MagicNumber = "CCOB";
-  static inline const uint16_t Version = 2;
-
-public:
-  static llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
-  compress(llvm::compression::Params P, const llvm::MemoryBuffer &Input,
-           bool Verbose = false);
-  static llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
-  decompress(llvm::MemoryBufferRef &Input, bool Verbose = false);
-};
-
 /// A simple binary serialization of an offloading file. We use this format to
 /// embed the offloading image into the host executable so it can be extracted
 /// and used by the linker.
@@ -210,160 +185,11 @@ class OffloadFile : public OwningBinary<OffloadBinary> {
   }
 };
 
-/// Bundle entry in binary clang-offload-bundler format.
-struct OffloadBundleEntry {
-  uint64_t Offset = 0u;
-  uint64_t Size = 0u;
-  uint64_t IDLength = 0u;
-  StringRef ID;
-  OffloadBundleEntry(uint64_t O, uint64_t S, uint64_t I, StringRef T)
-      : Offset(O), Size(S), IDLength(I), ID(T) {}
-  void dumpInfo(raw_ostream &OS) {
-    OS << "Offset = " << Offset << ", Size = " << Size
-       << ", ID Length = " << IDLength << ", ID = " << ID;
-  }
-  void dumpURI(raw_ostream &OS, StringRef filePath) {
-    OS << ID.data() << "\tfile:\/\/" << filePath << "#offset=" << Offset
-       << "&size=" << Size << "\n";
-  }
-};
-
-/// Fat binary embedded in object files in clang-offload-bundler format
-class OffloadBundleFatBin {
-
-private:
-  uint64_t Size = 0u;
-  StringRef FileName;
-  uint64_t NumberOfEntries;
-  SmallVector<OffloadBundleEntry> Entries;
-
-public:
-  SmallVector<OffloadBundleEntry> getEntries() { return Entries; }
-  uint64_t getSize() const { return Size; }
-  StringRef getFileName() const { return FileName; }
-  uint64_t getNumEntries() const { return NumberOfEntries; }
-
-  static Expected<std::unique_ptr<OffloadBundleFatBin>>
-  create(MemoryBufferRef, uint64_t SectionOffset, StringRef fileName);
-  Error extractBundle(const ObjectFile &Source);
-
-  Error DumpEntryToCodeObject();
-
-  Error ReadEntries(StringRef Section, uint64_t SectionOffset);
-  void DumpEntries() {
-    SmallVectorImpl<OffloadBundleEntry>::iterator it = Entries.begin();
-    for (uint64_t I = 0; I < Entries.size(); I++) {
-      it->dumpInfo(outs());
-      ++it;
-    }
-  }
-
-  void PrintEntriesAsURI() {
-    SmallVectorImpl<OffloadBundleEntry>::iterator it = Entries.begin();
-    for (uint64_t I = 0; I < NumberOfEntries; I++) {
-      it->dumpURI(outs(), FileName);
-      ++it;
-    }
-  }
-
-  OffloadBundleFatBin(MemoryBufferRef Source, StringRef file) : FileName(file) {
-    NumberOfEntries = 0;
-    Entries = SmallVector<OffloadBundleEntry>();
-  }
-
-  SmallVector<OffloadBundleEntry> EntryIDContains(StringRef str) {
-    SmallVector<OffloadBundleEntry> found = SmallVector<OffloadBundleEntry>();
-    SmallVectorImpl<OffloadBundleEntry>::iterator it = Entries.begin();
-    for (uint64_t I = 0; I < NumberOfEntries; I++) {
-      if (it->ID.contains(str)) {
-        found.push_back(*it);
-      }
-
-      ++it;
-    }
-    return found;
-  }
-};
-
-enum uri_type_t { FILE_URI, MEMORY_URI };
-
-struct OffloadBundleURI {
-  int64_t Offset = 0;
-  int64_t Size = 0;
-  uint64_t ProcessID = 0;
-  StringRef FileName;
-  uri_type_t URIType;
-
-  // Constructors
-  // TODO: add a Copy ctor ?
-  OffloadBundleURI(StringRef file, int64_t off, int64_t size)
-      : Offset(off), Size(size), ProcessID(0), FileName(file),
-        URIType(FILE_URI) {}
-
-  OffloadBundleURI(StringRef str, uri_type_t type) {
-    URIType = type;
-    switch (URIType) {
-    case FILE_URI:
-      parseFileName(str);
-      break;
-    case MEMORY_URI:
-      parseMemoryURI(str);
-      break;
-    default:
-      report_fatal_error("Unrecognized URI type.");
-    }
-  }
-
-  void parseFileName(StringRef str) {
-    ProcessID = 0;
-    URIType = FILE_URI;
-    if (str.consume_front("file://")) {
-      StringRef FilePathname =
-          str.take_until([](char c) { return (c == '#') || (c == '?'); });
-      FileName = FilePathname;
-      str = str.drop_front(FilePathname.size());
-
-      if (str.consume_front("#offset=")) {
-        StringRef OffsetStr = str.take_until([](char c) { return c == '&'; });
-        OffsetStr.getAsInteger(10, Offset);
-        str = str.drop_front(OffsetStr.size());
-
-        if (str.consume_front("&size=")) {
-          Size;
-          str.getAsInteger(10, Size);
-        } else
-          report_fatal_error("Reading 'size' in URI.");
-      } else
-        report_fatal_error("Reading 'offset' in URI.");
-    } else
-      report_fatal_error("Reading type of URI.");
-  }
-
-  void parseMemoryURI(StringRef str) {
-    // TODO: add parseMemoryURI type
-  }
-
-  StringRef getFileName() const { return FileName; }
-};
-
 /// Extracts embedded device offloading code from a memory \p Buffer to a list
 /// of \p Binaries.
 Error extractOffloadBinaries(MemoryBufferRef Buffer,
                              SmallVectorImpl<OffloadFile> &Binaries);
 
-/// Extracts fat binary in binary clang-offload-bundler format from object \p
-/// Obj and return it in \p Bundles
-Error extractOffloadBundleFatBinary(
-    const ObjectFile &Obj, SmallVectorImpl<OffloadBundleFatBin> &Bundles);
-
-/// Extract code object memory from the given \p Source object file at \p Offset
-/// and of \p Size, and copy into \p OutputFileName.
-Error extractCodeObject(const ObjectFile &Source, int64_t Offset, int64_t Size,
-                        StringRef OutputFileName);
-
-/// Extracts an Offload Bundle Entry given by URI
-Error extractOffloadBundleByURI(StringRef URIstr);
-
 /// Convert a string \p Name to an image kind.
 ImageKind getImageKind(StringRef Name);
 

llvm/include/llvm/Object/OffloadBundle.h

@@ -0,0 +1,199 @@
+//===- OffloadBundle.h - Utilities for offload bundles---*- C++ -*-===//
+//
+// 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
+//
+//===-------------------------------------------------------------------------===//
+//
+// This file contains the binary format used for budingling device metadata with
+// an associated device image. The data can then be stored inside a host object
+// file to create a fat binary and read by the linker. This is intended to be a
+// thin wrapper around the image itself. If this format becomes sufficiently
+// complex it should be moved to a standard binary format like msgpack or ELF.
+//
+//===-------------------------------------------------------------------------===//
+
+#ifndef LLVM_OBJECT_OFFLOADBUNDLE_H
+#define LLVM_OBJECT_OFFLOADBUNDLE_H
+
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Object/Binary.h"
+#include "llvm/Object/ObjectFile.h"
+#include "llvm/Support/Compression.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include <memory>
+
+namespace llvm {
+
+namespace object {
+
+class CompressedOffloadBundle {
+private:
+  static inline const size_t MagicSize = 4;
+  static inline const size_t VersionFieldSize = sizeof(uint16_t);
+  static inline const size_t MethodFieldSize = sizeof(uint16_t);
+  static inline const size_t FileSizeFieldSize = sizeof(uint32_t);
+  static inline const size_t UncompressedSizeFieldSize = sizeof(uint32_t);
+  static inline const size_t HashFieldSize = sizeof(uint64_t);
+  static inline const size_t V1HeaderSize =
+      MagicSize + VersionFieldSize + MethodFieldSize +
+      UncompressedSizeFieldSize + HashFieldSize;
+  static inline const size_t V2HeaderSize =
+      MagicSize + VersionFieldSize + FileSizeFieldSize + MethodFieldSize +
+      UncompressedSizeFieldSize + HashFieldSize;
+  static inline const llvm::StringRef MagicNumber = "CCOB";
+  static inline const uint16_t Version = 2;
+
+public:
+  static llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
+  compress(llvm::compression::Params P, const llvm::MemoryBuffer &Input,
+           bool Verbose = false);
+  static llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
+  decompress(llvm::MemoryBufferRef &Input, bool Verbose = false);
+};
+
+/// Bundle entry in binary clang-offload-bundler format.
+struct OffloadBundleEntry {
+  uint64_t Offset = 0u;
+  uint64_t Size = 0u;
+  uint64_t IDLength = 0u;
+  StringRef ID;
+  OffloadBundleEntry(uint64_t O, uint64_t S, uint64_t I, StringRef T)
+      : Offset(O), Size(S), IDLength(I), ID(T) {}
+  void dumpInfo(raw_ostream &OS) {
+    OS << "Offset = " << Offset << ", Size = " << Size
+       << ", ID Length = " << IDLength << ", ID = " << ID;
+  }
+  void dumpURI(raw_ostream &OS, StringRef FilePath) {
+    OS << ID.data() << "\tfile://" << FilePath << "#offset=" << Offset
+       << "&size=" << Size << "\n";
+  }
+};
+
+/// Fat binary embedded in object files in clang-offload-bundler format
+class OffloadBundleFatBin {
+
+  uint64_t Size = 0u;
+  StringRef FileName;
+  uint64_t NumberOfEntries;
+  SmallVector<OffloadBundleEntry> Entries;
+
+public:
+  SmallVector<OffloadBundleEntry> getEntries() { return Entries; }
+  uint64_t getSize() const { return Size; }
+  StringRef getFileName() const { return FileName; }
+  uint64_t getNumEntries() const { return NumberOfEntries; }
+
+  static Expected<std::unique_ptr<OffloadBundleFatBin>>
+  create(MemoryBufferRef, uint64_t SectionOffset, StringRef FileName);
+  Error extractBundle(const ObjectFile &Source);
+
+  Error dumpEntryToCodeObject();
+
+  Error readEntries(StringRef Section, uint64_t SectionOffset);
+  void dumpEntries() {
+    for (OffloadBundleEntry &Entry : Entries)
+      Entry.dumpInfo(outs());
+  }
+
+  void printEntriesAsURI() {
+    for (OffloadBundleEntry &Entry : Entries)
+      Entry.dumpURI(outs(), FileName);
+  }
+
+  OffloadBundleFatBin(MemoryBufferRef Source, StringRef File)
+      : FileName(File), NumberOfEntries(0),
+        Entries(SmallVector<OffloadBundleEntry>()) {}
+};
+
+enum UriTypeT { FILE_URI, MEMORY_URI };
+
+struct OffloadBundleURI {
+  int64_t Offset = 0;
+  int64_t Size = 0;
+  uint64_t ProcessID = 0;
+  StringRef FileName;
+  UriTypeT URIType;
+
+  // Constructors
+  // TODO: add a Copy ctor ?
+  OffloadBundleURI(StringRef File, int64_t Off, int64_t Size)
+      : Offset(Off), Size(Size), ProcessID(0), FileName(File),
+        URIType(FILE_URI) {}
+
+public:
+  static Expected<std::unique_ptr<OffloadBundleURI>>
+  createOffloadBundleURI(StringRef Str, UriTypeT Type) {
+    switch (Type) {
+    case FILE_URI:
+      return createFileURI(Str);
+      break;
+    case MEMORY_URI:
+      return createMemoryURI(Str);
+      break;
+    }
+    llvm_unreachable("Unknown UriTypeT enum");
+  }
+
+  static Expected<std::unique_ptr<OffloadBundleURI>>
+  createFileURI(StringRef Str) {
+    int64_t O = 0;
+    int64_t S = 0;
+
+    if (!Str.consume_front("file://"))
+      return createStringError(object_error::parse_failed,
+                               "Reading type of URI");
+
+    StringRef FilePathname =
+        Str.take_until([](char C) { return (C == '#') || (C == '?'); });
+    Str = Str.drop_front(FilePathname.size());
+
+    if (!Str.consume_front("#offset="))
+      return createStringError(object_error::parse_failed,
+                               "Reading 'offset' in URI");
+
+    StringRef OffsetStr = Str.take_until([](char C) { return C == '&'; });
+    OffsetStr.getAsInteger(10, O);
+    Str = Str.drop_front(OffsetStr.size());
+
+    if (Str.consume_front("&size="))
+      return createStringError(object_error::parse_failed,
+                               "Reading 'size' in URI");
+
+    Str.getAsInteger(10, S);
+    std::unique_ptr<OffloadBundleURI> OffloadingURI(
+        new OffloadBundleURI(FilePathname, O, S));
+    return std::move(OffloadingURI);
+  }
+
+  static Expected<std::unique_ptr<OffloadBundleURI>>
+  createMemoryURI(StringRef Str) {
+    // TODO: add parseMemoryURI type
+    return createStringError(object_error::parse_failed,
+                             "Memory Type URI is not currently supported.");
+  }
+
+  StringRef getFileName() const { return FileName; }
+};
+
+/// Extracts fat binary in binary clang-offload-bundler format from object \p
+/// Obj and return it in \p Bundles
+Error extractOffloadBundleFatBinary(
+    const ObjectFile &Obj, SmallVectorImpl<OffloadBundleFatBin> &Bundles);
+
+/// Extract code object memory from the given \p Source object file at \p Offset
+/// and of \p Size, and copy into \p OutputFileName.
+Error extractCodeObject(const ObjectFile &Source, int64_t Offset, int64_t Size,
+                        StringRef OutputFileName);
+
+/// Extracts an Offload Bundle Entry given by URI
+Error extractOffloadBundleByURI(StringRef URIstr);
+
+} // namespace object
+
+} // namespace llvm
+#endif