[Object] Refactor code for extracting offload binaries

[jhuber6]

Summary

We currently extract offload binaries inside of the linker wrapper.
Other tools may wish to do the same extraction operation. This patch
simply factors out this handling into the OffloadBinary.h interface.

[jhuber6]

ping

[yaxunl]

LGTM. Thanks.

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

 /// Codegen flags for LTO backend.
 static codegen::RegisterCodeGenFlags CodeGenFlags;
 
 /// Global flag to indicate that the LTO pipeline threw an error.
 static std::atomic<bool> LTOError;
 
 using OffloadingImage = OffloadBinary::OffloadingImage;
 
-/// A class to contain the binary information for a single OffloadBinary.
-class OffloadFile : public OwningBinary<OffloadBinary> {
-public:
-  using TargetID = std::pair<StringRef, StringRef>;
-
-  OffloadFile(std::unique_ptr<OffloadBinary> Binary,
-              std::unique_ptr<MemoryBuffer> Buffer)
-      : OwningBinary<OffloadBinary>(std::move(Binary), std::move(Buffer)) {}
-
-  /// We use the Triple and Architecture pair to group linker inputs together.
-  /// This conversion function lets us use these files in a hash-map.
-  operator TargetID() const {
-    return std::make_pair(getBinary()->getTriple(), getBinary()->getArch());
-  }
-};
-
 namespace llvm {
 // Provide DenseMapInfo so that OffloadKind can be used in a DenseMap.
 template <> struct DenseMapInfo<OffloadKind> {
   static inline OffloadKind getEmptyKey() { return OFK_LAST; }
   static inline OffloadKind getTombstoneKey() {
     return static_cast<OffloadKind>(OFK_LAST + 1);
   }
   static unsigned getHashValue(const OffloadKind &Val) { return Val; }

 const OptTable &getOptTable() {
   static const WrapperOptTable *Table = []() {
     auto Result = std::make_unique<WrapperOptTable>();
     return Result.release();
   }();
   return *Table;
 }
 
-Error extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
-                        SmallVectorImpl<OffloadFile> &DeviceFiles);
-
 void printCommands(ArrayRef<StringRef> CmdArgs) {
   if (CmdArgs.empty())
     return;
 
   llvm::errs() << " \"" << CmdArgs.front() << "\" ";
   for (auto IC = std::next(CmdArgs.begin()), IE = CmdArgs.end(); IC != IE; ++IC)
     llvm::errs() << *IC << (std::next(IC) != IE ? " " : "\n");
 }
if (Error Err = executeCommands(LinkerPath, LinkerArgs))
     return Err;
   return Error::success();
 }
 
 void printVersion(raw_ostream &OS) {
   OS << clang::getClangToolFullVersion("clang-linker-wrapper") << '\n';
 }
 
-/// Attempts to extract all the embedded device images contained inside the
-/// buffer \p Contents. The buffer is expected to contain a valid offloading
-/// binary format.
-Error extractOffloadFiles(MemoryBufferRef Contents,
-                          SmallVectorImpl<OffloadFile> &DeviceFiles) {
-  uint64_t Offset = 0;
-  // There could be multiple offloading binaries stored at this section.
-  while (Offset < Contents.getBuffer().size()) {
-    std::unique_ptr<MemoryBuffer> Buffer =
-        MemoryBuffer::getMemBuffer(Contents.getBuffer().drop_front(Offset), "",
-                                   /*RequiresNullTerminator*/ false);
-    auto BinaryOrErr = OffloadBinary::create(*Buffer);
-    if (!BinaryOrErr)
-      return BinaryOrErr.takeError();
-    OffloadBinary &Binary = **BinaryOrErr;
-
-    // Create a new owned binary with a copy of the original memory.
-    std::unique_ptr<MemoryBuffer> BufferCopy = MemoryBuffer::getMemBufferCopy(
-        Binary.getData().take_front(Binary.getSize()),
-        Contents.getBufferIdentifier());
-    auto NewBinaryOrErr = OffloadBinary::create(*BufferCopy);
-    if (!NewBinaryOrErr)
-      return NewBinaryOrErr.takeError();
-    DeviceFiles.emplace_back(std::move(*NewBinaryOrErr), std::move(BufferCopy));
-
-    Offset += Binary.getSize();
-  }
-
-  return Error::success();
-}
-
-// Extract offloading binaries from an Object file \p Obj.
-Error extractFromBinary(const ObjectFile &Obj,
-                        SmallVectorImpl<OffloadFile> &DeviceFiles) {
-  for (ELFSectionRef Sec : Obj.sections()) {
-    if (Sec.getType() != ELF::SHT_LLVM_OFFLOADING)
-      continue;
-
-    Expected<StringRef> Buffer = Sec.getContents();
-    if (!Buffer)
-      return Buffer.takeError();
-
-    MemoryBufferRef Contents(*Buffer, Obj.getFileName());
-    if (Error Err = extractOffloadFiles(Contents, DeviceFiles))
-      return Err;
-  }
-
-  return Error::success();
-}
-
-Error extractFromBitcode(std::unique_ptr<MemoryBuffer> Buffer,
-                         SmallVectorImpl<OffloadFile> &DeviceFiles) {
-  LLVMContext Context;
-  SMDiagnostic Err;
-  std::unique_ptr<Module> M = getLazyIRModule(std::move(Buffer), Err, Context);
-  if (!M)
-    return createStringError(inconvertibleErrorCode(),
-                             "Failed to create module");
-
-  // Extract offloading data from globals referenced by the
-  // `llvm.embedded.object` metadata with the `.llvm.offloading` section.
-  auto *MD = M->getNamedMetadata("llvm.embedded.objects");
-  if (!MD)
-    return Error::success();
-
-  for (const MDNode *Op : MD->operands()) {
-    if (Op->getNumOperands() < 2)
-      continue;
-
-    MDString *SectionID = dyn_cast<MDString>(Op->getOperand(1));
-    if (!SectionID || SectionID->getString() != ".llvm.offloading")
-      continue;
-
-    GlobalVariable *GV =
-        mdconst::dyn_extract_or_null<GlobalVariable>(Op->getOperand(0));
-    if (!GV)
-      continue;
-
-    auto *CDS = dyn_cast<ConstantDataSequential>(GV->getInitializer());
-    if (!CDS)
-      continue;
-
-    MemoryBufferRef Contents(CDS->getAsString(), M->getName());
-    if (Error Err = extractOffloadFiles(Contents, DeviceFiles))
-      return Err;
-  }
-
-  return Error::success();
-}
-
-Error extractFromArchive(const Archive &Library,
-                         SmallVectorImpl<OffloadFile> &DeviceFiles) {
-  // Try to extract device code from each file stored in the static archive.
-  Error Err = Error::success();
-  for (auto Child : Library.children(Err)) {
-    auto ChildBufferOrErr = Child.getMemoryBufferRef();
-    if (!ChildBufferOrErr)
-      return ChildBufferOrErr.takeError();
-    std::unique_ptr<MemoryBuffer> ChildBuffer =
-        MemoryBuffer::getMemBuffer(*ChildBufferOrErr, false);
-
-    // Check if the buffer has the required alignment.
-    if (!isAddrAligned(Align(OffloadBinary::getAlignment()),
-                       ChildBuffer->getBufferStart()))
-      ChildBuffer = MemoryBuffer::getMemBufferCopy(
-          ChildBufferOrErr->getBuffer(),
-          ChildBufferOrErr->getBufferIdentifier());
-
-    if (Error Err = extractFromBuffer(std::move(ChildBuffer), DeviceFiles))
-      return Err;
-  }
-
-  if (Err)
-    return Err;
-  return Error::success();
-}
-
-/// Extracts embedded device offloading code from a memory \p Buffer to a list
-/// of \p DeviceFiles.
-Error extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
-                        SmallVectorImpl<OffloadFile> &DeviceFiles) {
-  file_magic Type = identify_magic(Buffer->getBuffer());
-  switch (Type) {
-  case file_magic::bitcode:
-    return extractFromBitcode(std::move(Buffer), DeviceFiles);
-  case file_magic::elf_relocatable: {
-    Expected<std::unique_ptr<ObjectFile>> ObjFile =
-        ObjectFile::createObjectFile(*Buffer, Type);
-    if (!ObjFile)
-      return ObjFile.takeError();
-    return extractFromBinary(*ObjFile->get(), DeviceFiles);
-  }
-  case file_magic::archive: {
-    Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
-        object::Archive::create(*Buffer);
-    if (!LibFile)
-      return LibFile.takeError();
-    return extractFromArchive(*LibFile->get(), DeviceFiles);
-  }
-  default:
-    return Error::success();
-  }
-}
-
 namespace nvptx {
 Expected<StringRef> assemble(StringRef InputFile, const ArgList &Args,
                              bool RDC = true) {
   llvm::TimeTraceScope TimeScope("NVPTX Assembler");
   // NVPTX uses the ptxas binary to create device object files.
   Expected<std::string> PtxasPath = findProgram("ptxas", {CudaBinaryPath});
   if (!PtxasPath)
     return PtxasPath.takeError();
for (const opt::Arg *Arg : Args.filtered(OPT_INPUT)) {
 
     ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
         MemoryBuffer::getFileOrSTDIN(Filename);
     if (std::error_code EC = BufferOrErr.getError())
       return createFileError(Filename, EC);
 
     bool IsLazy =
         identify_magic((*BufferOrErr)->getBuffer()) == file_magic::archive;
-    if (Error Err = extractFromBuffer(std::move(*BufferOrErr),
-                                      IsLazy ? LazyInputFiles : InputFiles))
+    if (Error Err = extractOffloadBinaries(
+            **BufferOrErr, IsLazy ? LazyInputFiles : InputFiles))
       return std::move(Err);
   }
 
   // Try to extract input from input libraries.
   for (const opt::Arg *Arg : Args.filtered(OPT_library)) {
     if (auto Library = searchLibrary(Arg->getValue(), Root, LibraryPaths)) {
       ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
           MemoryBuffer::getFileOrSTDIN(*Library);
       if (std::error_code EC = BufferOrErr.getError())
         reportError(createFileError(*Library, EC));
 
-      if (Error Err =
-              extractFromBuffer(std::move(*BufferOrErr), LazyInputFiles))
+      if (Error Err = extractOffloadBinaries(**BufferOrErr, LazyInputFiles))
         return std::move(Err);
     }
   }
 
   for (StringRef Library : Args.getAllArgValues(OPT_bitcode_library_EQ)) {
     auto FileOrErr = getInputBitcodeLibrary(Library);
     if (!FileOrErr)
       return FileOrErr.takeError();

llvm/include/llvm/Object/OffloadBinary.h

private:
   /// 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
+/// owns its memory.
+class OffloadFile : public OwningBinary<OffloadBinary> {
+public:
+  using TargetID = std::pair<StringRef, StringRef>;
+
+  OffloadFile(std::unique_ptr<OffloadBinary> Binary,
+              std::unique_ptr<MemoryBuffer> Buffer)
+      : OwningBinary<OffloadBinary>(std::move(Binary), std::move(Buffer)) {}
+
+  /// We use the Triple and Architecture pair to group linker inputs together.
+  /// This conversion function lets us use these inputs in a hash-map.
+  operator TargetID() const {
+    return std::make_pair(getBinary()->getTriple(), getBinary()->getArch());
+  }
+};
+
+/// Extracts embedded device offloading code from a memory \p Buffer to a list
+/// of \p Binaries.
+Error extractOffloadBinaries(MemoryBufferRef Buffer,
+                             SmallVectorImpl<OffloadFile> &Binaries);
+
 /// Convert a string \p Name to an image kind.
 ImageKind getImageKind(StringRef Name);
 
 /// Convert an image kind to its string representation.
 StringRef getImageKindName(ImageKind Name);
 
 /// Convert a string \p Name to an offload kind.
 OffloadKind getOffloadKind(StringRef Name);
 
 /// Convert an offload kind to its string representation.
 StringRef getOffloadKindName(OffloadKind Name);
 
 } // namespace object
 
 } // namespace llvm
 #endif

llvm/lib/Object/CMakeLists.txt

add_llvm_component_library(LLVMObject
   DEPENDS
   intrinsics_gen
   llvm_vcsrevision_h
 
   LINK_COMPONENTS
   BitReader
   Core
   MC
+  IRReader
   BinaryFormat
   MCParser
   Support
   TextAPI
   )

llvm/lib/Object/OffloadBinary.cpp

 //===- Offloading.cpp - Utilities for handling offloading code  -*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Object/OffloadBinary.h"
 
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/BinaryFormat/Magic.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IRReader/IRReader.h"
 #include "llvm/MC/StringTableBuilder.h"
+#include "llvm/Object/Archive.h"
+#include "llvm/Object/ArchiveWriter.h"
+#include "llvm/Object/Binary.h"
+#include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/Error.h"
+#include "llvm/Object/IRObjectFile.h"
+#include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/FileOutputBuffer.h"
+#include "llvm/Support/SourceMgr.h"
 
 using namespace llvm;
 using namespace llvm::object;
 
+namespace {
+
+/// Attempts to extract all the embedded device images contained inside the
+/// buffer \p Contents. The buffer is expected to contain a valid offloading
+/// binary format.
+Error extractOffloadFiles(MemoryBufferRef Contents,
+                          SmallVectorImpl<OffloadFile> &Binaries) {
+  uint64_t Offset = 0;
+  // There could be multiple offloading binaries stored at this section.
+  while (Offset < Contents.getBuffer().size()) {
+    std::unique_ptr<MemoryBuffer> Buffer =
+        MemoryBuffer::getMemBuffer(Contents.getBuffer().drop_front(Offset), "",
+                                   /*RequiresNullTerminator*/ false);
+    auto BinaryOrErr = OffloadBinary::create(*Buffer);
+    if (!BinaryOrErr)
+      return BinaryOrErr.takeError();
+    OffloadBinary &Binary = **BinaryOrErr;
+
+    // Create a new owned binary with a copy of the original memory.
+    std::unique_ptr<MemoryBuffer> BufferCopy = MemoryBuffer::getMemBufferCopy(
+        Binary.getData().take_front(Binary.getSize()),
+        Contents.getBufferIdentifier());
+    auto NewBinaryOrErr = OffloadBinary::create(*BufferCopy);
+    if (!NewBinaryOrErr)
+      return NewBinaryOrErr.takeError();
+    Binaries.emplace_back(std::move(*NewBinaryOrErr), std::move(BufferCopy));
+
+    Offset += Binary.getSize();
+  }
+
+  return Error::success();
+}
+
+// Extract offloading binaries from an Object file \p Obj.
+Error extractFromBinary(const ObjectFile &Obj,
+                        SmallVectorImpl<OffloadFile> &Binaries) {
+  for (ELFSectionRef Sec : Obj.sections()) {
+    if (Sec.getType() != ELF::SHT_LLVM_OFFLOADING)
+      continue;
+
+    Expected<StringRef> Buffer = Sec.getContents();
+    if (!Buffer)
+      return Buffer.takeError();
+
+    MemoryBufferRef Contents(*Buffer, Obj.getFileName());
+    if (Error Err = extractOffloadFiles(Contents, Binaries))
+      return Err;
+  }
+
+  return Error::success();
+}
+
+Error extractFromBitcode(MemoryBufferRef Buffer,
+                         SmallVectorImpl<OffloadFile> &Binaries) {
+  LLVMContext Context;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = getLazyIRModule(
+      MemoryBuffer::getMemBuffer(Buffer, /*RequiresNullTerminator=*/false), Err,
+      Context);
+  if (!M)
+    return createStringError(inconvertibleErrorCode(),
+                             "Failed to create module");
+
+  // Extract offloading data from globals referenced by the
+  // `llvm.embedded.object` metadata with the `.llvm.offloading` section.
+  auto *MD = M->getNamedMetadata("llvm.embedded.objects");
+  if (!MD)
+    return Error::success();
+
+  for (const MDNode *Op : MD->operands()) {
+    if (Op->getNumOperands() < 2)
+      continue;
+
+    MDString *SectionID = dyn_cast<MDString>(Op->getOperand(1));
+    if (!SectionID || SectionID->getString() != ".llvm.offloading")
+      continue;
+
+    GlobalVariable *GV =
+        mdconst::dyn_extract_or_null<GlobalVariable>(Op->getOperand(0));
+    if (!GV)
+      continue;
+
+    auto *CDS = dyn_cast<ConstantDataSequential>(GV->getInitializer());
+    if (!CDS)
+      continue;
+
+    MemoryBufferRef Contents(CDS->getAsString(), M->getName());
+    if (Error Err = extractOffloadFiles(Contents, Binaries))
+      return Err;
+  }
+
+  return Error::success();
+}
+
+Error extractFromArchive(const Archive &Library,
+                         SmallVectorImpl<OffloadFile> &Binaries) {
+  // Try to extract device code from each file stored in the static archive.
+  Error Err = Error::success();
+  for (auto Child : Library.children(Err)) {
+    auto ChildBufferOrErr = Child.getMemoryBufferRef();
+    if (!ChildBufferOrErr)
+      return ChildBufferOrErr.takeError();
+    std::unique_ptr<MemoryBuffer> ChildBuffer =
+        MemoryBuffer::getMemBuffer(*ChildBufferOrErr, false);
+
+    // Check if the buffer has the required alignment.
+    if (!isAddrAligned(Align(OffloadBinary::getAlignment()),
+                       ChildBuffer->getBufferStart()))
+      ChildBuffer = MemoryBuffer::getMemBufferCopy(
+          ChildBufferOrErr->getBuffer(),
+          ChildBufferOrErr->getBufferIdentifier());
+
+    if (Error Err = extractOffloadBinaries(*ChildBuffer, Binaries))
+      return Err;
+  }
+
+  if (Err)
+    return Err;
+  return Error::success();
+}
+
+} // namespace
+
 Expected<std::unique_ptr<OffloadBinary>>
 OffloadBinary::create(MemoryBufferRef Buf) {
   if (Buf.getBufferSize() < sizeof(Header) + sizeof(Entry))
     return errorCodeToError(object_error::parse_failed);
 
   // Check for 0x10FF1OAD magic bytes.
   if (identify_magic(Buf.getBuffer()) != file_magic::offload_binary)
     return errorCodeToError(object_error::parse_failed);
OffloadBinary::write(const OffloadingImage &OffloadingData) {
   // Add final padding to required alignment.
   assert(TheHeader.Size >= OS.tell() && "Too much data written?");
   OS.write_zeros(TheHeader.Size - OS.tell());
   assert(TheHeader.Size == OS.tell() && "Size mismatch");
 
   return MemoryBuffer::getMemBufferCopy(OS.str());
 }
 
+Error object::extractOffloadBinaries(MemoryBufferRef Buffer,
+                                     SmallVectorImpl<OffloadFile> &Binaries) {
+  file_magic Type = identify_magic(Buffer.getBuffer());
+  switch (Type) {
+  case file_magic::bitcode:
+    return extractFromBitcode(Buffer, Binaries);
+  case file_magic::elf_executable:
+  case file_magic::elf_relocatable:
+  case file_magic::elf_shared_object: {
+    Expected<std::unique_ptr<ObjectFile>> ObjFile =
+        ObjectFile::createObjectFile(Buffer, Type);
+    if (!ObjFile)
+      return ObjFile.takeError();
+    return extractFromBinary(*ObjFile->get(), Binaries);
+  }
+  case file_magic::archive: {
+    Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
+        object::Archive::create(Buffer);
+    if (!LibFile)
+      return LibFile.takeError();
+    return extractFromArchive(*LibFile->get(), Binaries);
+  }
+  case file_magic::offload_binary:
+    return extractOffloadFiles(Buffer, Binaries);
+  default:
+    return Error::success();
+  }
+}
+
 OffloadKind object::getOffloadKind(StringRef Name) {
   return llvm::StringSwitch<OffloadKind>(Name)
       .Case("openmp", OFK_OpenMP)
       .Case("cuda", OFK_Cuda)
       .Case("hip", OFK_HIP)
       .Default(OFK_None);
 }