1 files changed, 182 insertions, 0 deletions

diff --git a/contrib/llvm-project/llvm/lib/Object/OffloadBinary.cpp b/contrib/llvm-project/llvm/lib/Object/OffloadBinary.cpp
index 21946ec2d6fb..d8cdcdc21d39 100644
--- a/contrib/llvm-project/llvm/lib/Object/OffloadBinary.cpp
+++ b/contrib/llvm-project/llvm/lib/Object/OffloadBinary.cpp
@@ -10,14 +10,166 @@
 
 #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/COFF.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);
+    if (!isAddrAligned(Align(OffloadBinary::getAlignment()),
+                       Buffer->getBufferStart()))
+      Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(),
+                                              Buffer->getBufferIdentifier());
+    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 extractFromObject(const ObjectFile &Obj,
+                        SmallVectorImpl<OffloadFile> &Binaries) {
+  assert((Obj.isELF() || Obj.isCOFF()) && "Invalid file type");
+
+  for (SectionRef Sec : Obj.sections()) {
+    // ELF files contain a section with the LLVM_OFFLOADING type.
+    if (Obj.isELF() &&
+        static_cast<ELFSectionRef>(Sec).getType() != ELF::SHT_LLVM_OFFLOADING)
+      continue;
+
+    // COFF has no section types so we rely on the name of the section.
+    if (Obj.isCOFF()) {
+      Expected<StringRef> NameOrErr = Sec.getName();
+      if (!NameOrErr)
+        return NameOrErr.takeError();
+
+      if (!NameOrErr->equals(".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))
@@ -115,6 +267,36 @@ OffloadBinary::write(const OffloadingImage &OffloadingData) {
   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_relocatable:
+  case file_magic::elf_executable:
+  case file_magic::elf_shared_object:
+  case file_magic::coff_object: {
+    Expected<std::unique_ptr<ObjectFile>> ObjFile =
+        ObjectFile::createObjectFile(Buffer, Type);
+    if (!ObjFile)
+      return ObjFile.takeError();
+    return extractFromObject(*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)