[Offloading] Extend OffloadBinary format to support multiple metadata entries

clang/test/Driver/linker-wrapper-image.c

@@ -25,7 +25,7 @@
 // OPENMP-REL: @.omp_offloading.device_image = internal unnamed_addr constant [[[SIZE:[0-9]+]] x i8] c"\10\FF\10\AD{{.*}}", section ".llvm.offloading.relocatable", align 8
 
 //      OPENMP: @.omp_offloading.device_image = internal unnamed_addr constant [[[SIZE:[0-9]+]] x i8] c"\10\FF\10\AD{{.*}}", section ".llvm.offloading", align 8
-// OPENMP-NEXT: @.omp_offloading.device_images = internal unnamed_addr constant [1 x %__tgt_device_image] [%__tgt_device_image { ptr getelementptr ([[[BEGIN:[0-9]+]] x i8], ptr @.omp_offloading.device_image, i64 0, i64 144), ptr getelementptr ([[[END:[0-9]+]] x i8], ptr @.omp_offloading.device_image, i64 0, i64 144), ptr @__start_llvm_offload_entries, ptr @__stop_llvm_offload_entries }]
+// OPENMP-NEXT: @.omp_offloading.device_images = internal unnamed_addr constant [1 x %__tgt_device_image] [%__tgt_device_image { ptr getelementptr ([[[IMG_OFF:[0-9]+]] x i8], ptr @.omp_offloading.device_image, i64 0, i64 [[IMG_OFF]]), ptr getelementptr ([[[IMG_OFF]] x i8], ptr @.omp_offloading.device_image, i64 0, i64 [[IMG_OFF]]), ptr @__start_llvm_offload_entries, ptr @__stop_llvm_offload_entries }]
 // OPENMP-NEXT: @.omp_offloading.descriptor = internal constant %__tgt_bin_desc { i32 1, ptr @.omp_offloading.device_images, ptr @__start_llvm_offload_entries, ptr @__stop_llvm_offload_entries }
 // OPENMP-NEXT: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 101, ptr @.omp_offloading.descriptor_reg, ptr null }]
 

llvm/include/llvm/Object/OffloadBinary.h

@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file contains the binary format used for budingling device metadata with
+// This file contains the binary format used for bundling 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
@@ -52,6 +52,13 @@ enum ImageKind : uint16_t {
   IMG_LAST,
 };
 
+/// Flags associated with the Entry.
+enum OffloadEntryFlags : uint32_t {
+  OIF_None = 0,
+  // Entry doesn't contain image. Used to keep metadata only entries.
+  OIF_NoImage = (1 << 0),
+};
+
 /// 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.
@@ -63,11 +70,38 @@ enum ImageKind : uint16_t {
 /// offsets from the beginning of the file.
 class OffloadBinary : public Binary {
 public:
+  struct Header {
+    uint8_t Magic[4] = {0x10, 0xFF, 0x10, 0xAD}; // 0x10FF10AD magic bytes.
+    uint32_t Version = OffloadBinary::Version;   // Version identifier.
+    uint64_t Size;          // Size in bytes of this entire binary.
+    uint64_t EntriesOffset; // Offset in bytes to the start of entries block.
+    uint64_t EntriesCount;  // Number of metadata entries in the binary.
+  };
+
+  struct Entry {
+    ImageKind TheImageKind;     // The kind of the image stored.
+    OffloadKind TheOffloadKind; // The producer of this image.
+    uint32_t Flags;             // Additional flags associated with the entry.
+    uint64_t StringOffset;      // Offset in bytes to the string map.
+    uint64_t NumStrings;        // Number of entries in the string map.
+    uint64_t ImageOffset;       // Offset in bytes of the actual binary image.
+    uint64_t ImageSize;         // Size in bytes of the binary image.
+  };
+
+  struct StringEntry {
+    uint64_t KeyOffset;
+    uint64_t ValueOffset;
+    uint64_t ValueSize; // Size of the value in bytes.
+  };
+
+  using StringMap = MapVector<StringRef, StringRef>;
+  using entry_iterator = SmallVector<std::pair<const Entry *, StringMap>, 1>::const_iterator;
+  using entry_iterator_range = iterator_range<entry_iterator>;
   using string_iterator = MapVector<StringRef, StringRef>::const_iterator;
   using string_iterator_range = iterator_range<string_iterator>;
 
   /// The current version of the binary used for backwards compatibility.
-  static const uint32_t Version = 1;
+  static const uint32_t Version = 2;
 
   /// The offloading metadata that will be serialized to a memory buffer.
   struct OffloadingImage {
@@ -82,76 +116,79 @@ class OffloadBinary : public Binary {
   LLVM_ABI static Expected<std::unique_ptr<OffloadBinary>>
       create(MemoryBufferRef);
 
-  /// Serialize the contents of \p File to a binary buffer to be read later.
-  LLVM_ABI static SmallString<0> write(const OffloadingImage &);
+  /// Serialize the contents of \p OffloadingData to a binary buffer to be read
+  /// later.
+  LLVM_ABI static SmallString<0>
+  write(ArrayRef<OffloadingImage> OffloadingData);
 
   static uint64_t getAlignment() { return 8; }
 
-  ImageKind getImageKind() const { return TheEntry->TheImageKind; }
-  OffloadKind getOffloadKind() const { return TheEntry->TheOffloadKind; }
+  ImageKind getOnlyImageKind() const {
+    assert(getEntriesCount() == 1 && "Expected exactly one entry.");
+    return Entries[0].first->TheImageKind;
+  }
+
+  OffloadKind getOffloadKind() const { return Entries[0].first->TheOffloadKind; }
   uint32_t getVersion() const { return TheHeader->Version; }
-  uint32_t getFlags() const { return TheEntry->Flags; }
+  uint32_t getFlags() const { return Entries[0].first->Flags; }
   uint64_t getSize() const { return TheHeader->Size; }
+  uint64_t getEntriesCount() const { return TheHeader->EntriesCount; }
 
   StringRef getTriple() const { return getString("triple"); }
   StringRef getArch() const { return getString("arch"); }
   StringRef getImage() const {
-    return StringRef(&Buffer[TheEntry->ImageOffset], TheEntry->ImageSize);
+    return StringRef(&Buffer[Entries[0].first->ImageOffset], Entries[0].first->ImageSize);
   }
 
-  // Iterator over all the key and value pairs in the binary.
-  string_iterator_range strings() const { return StringData; }
+  // Iterator access to all entries in the binary
+  entry_iterator_range entries() const {
+    return make_range(Entries.begin(), Entries.end());
+  }
+  entry_iterator entries_begin() const { return Entries.begin(); }
+  entry_iterator entries_end() const { return Entries.end(); }
 
-  StringRef getString(StringRef Key) const { return StringData.lookup(Key); }
+  // Access specific entry by index.
+  const std::pair<const Entry *, StringMap> &getEntry(size_t Index) const {
+    assert(Index < Entries.size() && "Entry index out of bounds");
+    return Entries[Index];
+  }
 
-  static bool classof(const Binary *V) { return V->isOffloadFile(); }
+  // Iterator over all the key and value pairs in the binary.
+  string_iterator_range strings() const { return Entries[0].second; }
 
-  struct Header {
-    uint8_t Magic[4] = {0x10, 0xFF, 0x10, 0xAD}; // 0x10FF10AD magic bytes.
-    uint32_t Version = OffloadBinary::Version;   // Version identifier.
-    uint64_t Size;        // Size in bytes of this entire binary.
-    uint64_t EntryOffset; // Offset of the metadata entry in bytes.
-    uint64_t EntrySize;   // Size of the metadata entry in bytes.
-  };
+  StringRef getString(StringRef Key) const { return Entries[0].second.lookup(Key); }
 
-  struct Entry {
-    ImageKind TheImageKind;     // The kind of the image stored.
-    OffloadKind TheOffloadKind; // The producer of this image.
-    uint32_t Flags;             // Additional flags associated with the image.
-    uint64_t StringOffset;      // Offset in bytes to the string map.
-    uint64_t NumStrings;        // Number of entries in the string map.
-    uint64_t ImageOffset;       // Offset in bytes of the actual binary image.
-    uint64_t ImageSize;         // Size in bytes of the binary image.
-  };
-
-  struct StringEntry {
-    uint64_t KeyOffset;
-    uint64_t ValueOffset;
-  };
+  static bool classof(const Binary *V) { return V->isOffloadFile(); }
 
 private:
   OffloadBinary(MemoryBufferRef Source, const Header *TheHeader,
-                const Entry *TheEntry)
+                const Entry *EntriesBegin)
       : Binary(Binary::ID_Offload, Source), Buffer(Source.getBufferStart()),
-        TheHeader(TheHeader), TheEntry(TheEntry) {
-    const StringEntry *StringMapBegin =
-        reinterpret_cast<const StringEntry *>(&Buffer[TheEntry->StringOffset]);
-    for (uint64_t I = 0, E = TheEntry->NumStrings; I != E; ++I) {
-      StringRef Key = &Buffer[StringMapBegin[I].KeyOffset];
-      StringData[Key] = &Buffer[StringMapBegin[I].ValueOffset];
+        TheHeader(TheHeader) {
+    for (uint64_t EI = 0, EE = TheHeader->EntriesCount; EI != EE; ++EI) {
+      const Entry *TheEntry = &EntriesBegin[EI];
+      const StringEntry *StringMapBegin = reinterpret_cast<const StringEntry *>(
+          &Buffer[TheEntry->StringOffset]);
+      StringMap Strings;
+      for (uint64_t SI = 0, SE = TheEntry->NumStrings; SI != SE; ++SI) {
+        StringRef Key = &Buffer[StringMapBegin[SI].KeyOffset];
+        StringRef Value = StringRef(
+            &Buffer[StringMapBegin[SI].ValueOffset], StringMapBegin[SI].ValueSize);
+        Strings.insert({Key, Value});
+      }
+      Entries.push_back(std::make_pair(TheEntry, std::move(Strings)));
     }
   }
 
   OffloadBinary(const OffloadBinary &Other) = delete;
 
-  /// Map from keys to offsets in the binary.
-  MapVector<StringRef, StringRef> StringData;
+  /// Location of the metadata entries within the binary mapped to
+  /// the key-value string data.
+  SmallVector<std::pair<const Entry *, StringMap>, 1> Entries;
   /// Raw pointer to the MemoryBufferRef for convenience.
   const char *Buffer;
   /// Location of the header within the binary.
   const Header *TheHeader;
-  /// Location of the metadata entries within the binary.
-  const Entry *TheEntry;
 };
 
 /// A class to contain the binary information for a single OffloadBinary that

llvm/include/llvm/ObjectYAML/OffloadYAML.h

@@ -39,8 +39,8 @@ struct Binary {
 
   std::optional<uint32_t> Version;
   std::optional<uint64_t> Size;
-  std::optional<uint64_t> EntryOffset;
-  std::optional<uint64_t> EntrySize;
+  std::optional<uint64_t> EntriesOffset;
+  std::optional<uint64_t> EntriesCount;
   std::vector<Member> Members;
 };
 

llvm/lib/Frontend/Offloading/OffloadWrapper.cpp

@@ -161,7 +161,7 @@ GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs,
               Binary.bytes_begin());
       const auto *Entry =
           reinterpret_cast<const object::OffloadBinary::Entry *>(
-              Binary.bytes_begin() + Header->EntryOffset);
+              Binary.bytes_begin() + Header->EntriesOffset);
       BeginOffset = Entry->ImageOffset;
       EndOffset = Entry->ImageOffset + Entry->ImageSize;
     }
@@ -952,7 +952,7 @@ class SYCLWrapper {
     Constant *OffloadKindConstant = ConstantInt::get(
         Type::getInt8Ty(C), static_cast<uint8_t>(OB.getOffloadKind()));
     Constant *ImageKindConstant = ConstantInt::get(
-        Type::getInt8Ty(C), static_cast<uint8_t>(OB.getImageKind()));
+        Type::getInt8Ty(C), static_cast<uint8_t>(OB.getOnlyImageKind()));
     StringRef Triple = OB.getString("triple");
     Constant *TripleConstant =
         addStringToModule(Triple, Twine(OffloadKindTag) + "target." + ImageID);

llvm/lib/Object/OffloadBinary.cpp

@@ -189,74 +189,104 @@ OffloadBinary::create(MemoryBufferRef Buf) {
       TheHeader->Size < sizeof(Entry) || TheHeader->Size < sizeof(Header))
     return errorCodeToError(object_error::unexpected_eof);
 
-  if (TheHeader->EntryOffset > TheHeader->Size - sizeof(Entry) ||
-      TheHeader->EntrySize > TheHeader->Size - sizeof(Header))
+  uint64_t EntriesSize = sizeof(Entry) * TheHeader->EntriesCount;
+  if (TheHeader->EntriesOffset > TheHeader->Size - EntriesSize ||
+      EntriesSize > TheHeader->Size - sizeof(Header))
     return errorCodeToError(object_error::unexpected_eof);
 
-  const Entry *TheEntry =
-      reinterpret_cast<const Entry *>(&Start[TheHeader->EntryOffset]);
+  const Entry *Entries =
+      reinterpret_cast<const Entry *>(&Start[TheHeader->EntriesOffset]);
+  for (uint32_t I = 0; I < TheHeader->EntriesCount; ++I) {
+    const Entry *TheEntry = &Entries[I];
 
-  if (TheEntry->ImageOffset > Buf.getBufferSize() ||
-      TheEntry->StringOffset > Buf.getBufferSize())
-    return errorCodeToError(object_error::unexpected_eof);
+    if (TheEntry->ImageOffset > Buf.getBufferSize() ||
+        TheEntry->StringOffset > Buf.getBufferSize())
+      return errorCodeToError(object_error::unexpected_eof);
+  }
 
   return std::unique_ptr<OffloadBinary>(
-      new OffloadBinary(Buf, TheHeader, TheEntry));
+      new OffloadBinary(Buf, TheHeader, Entries));
 }
 
-SmallString<0> OffloadBinary::write(const OffloadingImage &OffloadingData) {
+SmallString<0> OffloadBinary::write(ArrayRef<OffloadingImage> OffloadingData) {
+  uint64_t EntriesCount = OffloadingData.size();
+  assert(EntriesCount > 0 && "At least one offloading image is required");
+
   // Create a null-terminated string table with all the used strings.
+  // Also calculate total size of images.
   StringTableBuilder StrTab(StringTableBuilder::ELF);
-  for (auto &KeyAndValue : OffloadingData.StringData) {
-    StrTab.add(KeyAndValue.first);
-    StrTab.add(KeyAndValue.second);
+  uint64_t TotalStringEntries = 0;
+  uint64_t TotalImagesSize = 0;
+  for (const OffloadingImage &Img : OffloadingData) {
+    for (auto &KeyAndValue : Img.StringData) {
+      StrTab.add(KeyAndValue.first);
+      StrTab.add(KeyAndValue.second);
+    }
+    TotalStringEntries += Img.StringData.size();
+    TotalImagesSize += Img.Image->getBufferSize();
   }
   StrTab.finalize();
 
-  uint64_t StringEntrySize =
-      sizeof(StringEntry) * OffloadingData.StringData.size();
+  uint64_t StringEntrySize = sizeof(StringEntry) * TotalStringEntries;
+  uint64_t EntriesSize = sizeof(Entry) * EntriesCount;
+  uint64_t StrTabOffset = sizeof(Header) + EntriesSize + StringEntrySize;
 
   // Make sure the image we're wrapping around is aligned as well.
-  uint64_t BinaryDataSize = alignTo(sizeof(Header) + sizeof(Entry) +
-                                        StringEntrySize + StrTab.getSize(),
-                                    getAlignment());
+  uint64_t BinaryDataSize =
+      alignTo(StrTabOffset + StrTab.getSize(), getAlignment());
 
-  // Create the header and fill in the offsets. The entry will be directly
+  // Create the header and fill in the offsets. The entries will be directly
   // placed after the header in memory. Align the size to the alignment of the
   // header so this can be placed contiguously in a single section.
   Header TheHeader;
-  TheHeader.Size = alignTo(
-      BinaryDataSize + OffloadingData.Image->getBufferSize(), getAlignment());
-  TheHeader.EntryOffset = sizeof(Header);
-  TheHeader.EntrySize = sizeof(Entry);
-
-  // Create the entry using the string table offsets. The string table will be
-  // placed directly after the entry in memory, and the image after that.
-  Entry TheEntry;
-  TheEntry.TheImageKind = OffloadingData.TheImageKind;
-  TheEntry.TheOffloadKind = OffloadingData.TheOffloadKind;
-  TheEntry.Flags = OffloadingData.Flags;
-  TheEntry.StringOffset = sizeof(Header) + sizeof(Entry);
-  TheEntry.NumStrings = OffloadingData.StringData.size();
-
-  TheEntry.ImageOffset = BinaryDataSize;
-  TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
+  TheHeader.Size = alignTo(BinaryDataSize + TotalImagesSize, getAlignment());
+  TheHeader.EntriesOffset = sizeof(Header);
+  TheHeader.EntriesCount = EntriesCount;
 
   SmallString<0> Data;
   Data.reserve(TheHeader.Size);
   raw_svector_ostream OS(Data);
   OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
-  OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
-  for (auto &KeyAndValue : OffloadingData.StringData) {
-    uint64_t Offset = sizeof(Header) + sizeof(Entry) + StringEntrySize;
-    StringEntry Map{Offset + StrTab.getOffset(KeyAndValue.first),
-                    Offset + StrTab.getOffset(KeyAndValue.second)};
-    OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
+
+  // Create the entries using the string table offsets. The string table will be
+  // placed directly after the set of entries in memory, and all the images are
+  // after that.
+  uint64_t StringEntryOffset = sizeof(Header) + EntriesSize;
+  uint64_t ImageOffset = BinaryDataSize;
+  for (const OffloadingImage &Img : OffloadingData) {
+    Entry TheEntry;
+
+    TheEntry.TheImageKind = Img.TheImageKind;
+    TheEntry.TheOffloadKind = Img.TheOffloadKind;
+    TheEntry.Flags = Img.Flags;
+
+    TheEntry.StringOffset = StringEntryOffset;
+    StringEntryOffset += sizeof(StringEntry) * Img.StringData.size();
+    TheEntry.NumStrings = Img.StringData.size();
+
+    TheEntry.ImageOffset = ImageOffset;
+    ImageOffset += Img.Image->getBufferSize();
+    TheEntry.ImageSize = Img.Image->getBufferSize();
+
+    OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
   }
+
+  // Create the string map entries.
+  for (const OffloadingImage &Img : OffloadingData) {
+    for (auto &KeyAndValue : Img.StringData) {
+      StringEntry Map{StrTabOffset + StrTab.getOffset(KeyAndValue.first),
+                      StrTabOffset + StrTab.getOffset(KeyAndValue.second),
+                      KeyAndValue.second.size()};
+      OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
+    }
+  }
+
   StrTab.write(OS);
   // Add padding to required image alignment.
-  OS.write_zeros(TheEntry.ImageOffset - OS.tell());
-  OS << OffloadingData.Image->getBuffer();
+  OS.write_zeros(BinaryDataSize - OS.tell());
+
+  for (const OffloadingImage &Img : OffloadingData)
+    OS << Img.Image->getBuffer();
 
   // Add final padding to required alignment.
   assert(TheHeader.Size >= OS.tell() && "Too much data written?");

llvm/lib/ObjectYAML/OffloadEmitter.cpp

@@ -45,10 +45,10 @@ bool yaml2offload(Binary &Doc, raw_ostream &Out, ErrorHandler EH) {
       TheHeader->Version = *Doc.Version;
     if (Doc.Size)
       TheHeader->Size = *Doc.Size;
-    if (Doc.EntryOffset)
-      TheHeader->EntryOffset = *Doc.EntryOffset;
-    if (Doc.EntrySize)
-      TheHeader->EntrySize = *Doc.EntrySize;
+    if (Doc.EntriesOffset)
+      TheHeader->EntriesOffset = *Doc.EntriesOffset;
+    if (Doc.EntriesCount)
+      TheHeader->EntriesCount = *Doc.EntriesCount;
 
     Out.write(Buffer.begin(), Buffer.size());
   }

llvm/lib/ObjectYAML/OffloadYAML.cpp

@@ -50,8 +50,8 @@ void MappingTraits<OffloadYAML::Binary>::mapping(IO &IO,
   IO.mapTag("!Offload", true);
   IO.mapOptional("Version", O.Version);
   IO.mapOptional("Size", O.Size);
-  IO.mapOptional("EntryOffset", O.EntryOffset);
-  IO.mapOptional("EntrySize", O.EntrySize);
+  IO.mapOptional("EntriesOffset", O.EntriesOffset);
+  IO.mapOptional("EntriesCount", O.EntriesCount);
   IO.mapRequired("Members", O.Members);
   IO.setContext(nullptr);
 }