[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.
@@ -67,7 +74,7 @@ class OffloadBinary : public Binary {
   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 {
@@ -78,12 +85,52 @@ class OffloadBinary : public Binary {
     std::unique_ptr<MemoryBuffer> Image;
   };
 
+  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.
+  };
+
+  struct StringEntryV1 {
+    uint64_t KeyOffset;
+    uint64_t ValueOffset;
+  };
+
+  /// Attempt to extract and validate the header from the offloading binary in
+  /// \p Buf.
+  LLVM_ABI
+  static Expected<const Header *> extractHeader(MemoryBufferRef Buf);
+
   /// Attempt to parse the offloading binary stored in \p Data.
   LLVM_ABI static Expected<std::unique_ptr<OffloadBinary>>
-      create(MemoryBufferRef);
+      createV1(MemoryBufferRef);
 
-  /// Serialize the contents of \p File to a binary buffer to be read later.
-  LLVM_ABI static SmallString<0> write(const OffloadingImage &);
+  /// Attempt to parse the offloading binary stored in \p Data.
+  LLVM_ABI static Expected<std::unique_ptr<OffloadBinary>>
+      createV2(MemoryBufferRef);
+
+  /// 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; }
 
@@ -106,29 +153,6 @@ class OffloadBinary : public Binary {
 
   static bool classof(const Binary *V) { return V->isOffloadFile(); }
 
-  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.
-  };
-
-  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;
-  };
-
 private:
   OffloadBinary(MemoryBufferRef Source, const Header *TheHeader,
                 const Entry *TheEntry)
@@ -138,7 +162,7 @@ class OffloadBinary : public Binary {
         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];
+      StringData[Key] = StringRef(&Buffer[StringMapBegin[I].ValueOffset], StringMapBegin[I].ValueSize);
     }
   }
 
@@ -171,7 +195,7 @@ class OffloadFile : public OwningBinary<OffloadBinary> {
         getBinary()->getMemoryBufferRef().getBufferIdentifier());
 
     // This parsing should never fail because it has already been parsed.
-    auto NewBinaryOrErr = OffloadBinary::create(*Buffer);
+    auto NewBinaryOrErr = OffloadBinary::createV1(*Buffer);
     assert(NewBinaryOrErr && "Failed to parse a copy of the binary?");
     if (!NewBinaryOrErr)
       llvm::consumeError(NewBinaryOrErr.takeError());

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;
     }

llvm/lib/Object/Binary.cpp

@@ -94,7 +94,7 @@ Expected<std::unique_ptr<Binary>> object::createBinary(MemoryBufferRef Buffer,
     // Unrecognized object file format.
     return errorCodeToError(object_error::invalid_file_type);
   case file_magic::offload_binary:
-    return OffloadBinary::create(Buffer);
+    return OffloadBinary::createV1(Buffer);
   case file_magic::minidump:
     return MinidumpFile::create(Buffer);
   case file_magic::tapi_file: