jhuber6 created this revision.
jhuber6 added reviewers: saiislam, JonChesterfield, yaxunl, ronlieb, tra, 
jdoerfert.
Herald added a subscriber: mgorny.
Herald added a project: All.
jhuber6 requested review of this revision.
Herald added subscribers: cfe-commits, sstefan1.
Herald added a project: clang.

A previous patch added support for extracting images from offloading
binaries. Users may wish to extract these files from the file types they
are most commonly emebedded in, such as an ELF or bitcode. This can be
difficult for the user to do manually, as these could be stored in
different section names potentially. This patch addsp support for
extracting these file types.

Note, this patch duplicated a lot of code from the `ClangLinkerWrapper`
that does a nearly identical extraction. Should these be merged into a
public interface via `Object/OffloadBinary.h`?


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D132607

Files:
  clang/test/Driver/offload-packager.c
  clang/tools/clang-offload-packager/CMakeLists.txt
  clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

Index: clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
===================================================================
--- clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
+++ clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
@@ -14,7 +14,14 @@
 
 #include "clang/Basic/Version.h"
 
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Object/Archive.h"
+#include "llvm/Object/ArchiveWriter.h"
 #include "llvm/Object/Binary.h"
+#include "llvm/Object/ELFObjectFile.h"
+#include "llvm/Object/IRObjectFile.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Object/OffloadBinary.h"
 #include "llvm/Support/CommandLine.h"
@@ -23,6 +30,7 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Signals.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
 
@@ -54,6 +62,12 @@
   OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n';
 }
 
+/// A single OffloadBinary and its associated memory.
+using OffloadFile = OwningBinary<OffloadBinary>;
+
+Error extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
+                        SmallVectorImpl<OffloadFile> &DeviceFiles);
+
 // Get a map containing all the arguments for the image. Repeated arguments will
 // be placed in a comma separated list.
 static DenseMap<StringRef, StringRef> getImageArguments(StringRef Image,
@@ -123,12 +137,11 @@
   return Error::success();
 }
 
-static Expected<SmallVector<std::unique_ptr<OffloadBinary>>>
-extractOffloadFiles(MemoryBufferRef Contents) {
-  if (identify_magic(Contents.getBuffer()) != file_magic::offload_binary)
-    return createStringError(inconvertibleErrorCode(),
-                             "Input buffer not an offloading binary");
-  SmallVector<std::unique_ptr<OffloadBinary>> Binaries;
+/// 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()) {
@@ -138,12 +151,136 @@
     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();
+}
 
-    Offset += (*BinaryOrErr)->getSize();
-    Binaries.emplace_back(std::move(*BinaryOrErr));
+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;
   }
 
-  return std::move(Binaries);
+  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);
+  }
+  case file_magic::offload_binary:
+    return extractOffloadFiles(*Buffer, DeviceFiles);
+  default:
+    return Error::success();
+  }
 }
 
 static Error unbundleImages() {
@@ -159,9 +296,9 @@
     Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(),
                                             Buffer->getBufferIdentifier());
 
-  auto BinariesOrErr = extractOffloadFiles(*Buffer);
-  if (!BinariesOrErr)
-    return BinariesOrErr.takeError();
+  SmallVector<OffloadFile> Binaries;
+  if (Error Err = extractFromBuffer(std::move(Buffer), Binaries))
+    return Err;
 
   // Try to extract each device image specified by the user from the input file.
   for (StringRef Image : DeviceImages) {
@@ -169,8 +306,8 @@
     StringSaver Saver(Alloc);
     auto Args = getImageArguments(Image, Saver);
 
-    for (uint64_t I = 0, E = BinariesOrErr->size(); I != E; ++I) {
-      const auto &Binary = (*BinariesOrErr)[I];
+    for (uint64_t I = 0, E = Binaries.size(); I != E; ++I) {
+      const auto *Binary = Binaries[I].getBinary();
       // We handle the 'file' and 'kind' identifiers differently.
       bool Match = llvm::all_of(Args, [&](auto &Arg) {
         const auto [Key, Value] = Arg;
Index: clang/tools/clang-offload-packager/CMakeLists.txt
===================================================================
--- clang/tools/clang-offload-packager/CMakeLists.txt
+++ clang/tools/clang-offload-packager/CMakeLists.txt
@@ -1,6 +1,9 @@
 set(LLVM_LINK_COMPONENTS 
   ${LLVM_TARGETS_TO_BUILD}
   BinaryFormat
+  BitWriter
+  Core
+  IRReader
   Object
   Support)
 
Index: clang/test/Driver/offload-packager.c
===================================================================
--- clang/test/Driver/offload-packager.c
+++ clang/test/Driver/offload-packager.c
@@ -29,3 +29,25 @@
 // RUN: diff *-amdgcn-amd-amdhsa-gfx908.2.o %S/Inputs/dummy-elf.o; rm *-amdgcn-amd-amdhsa-gfx908.2.o
 // RUN: diff *-amdgcn-amd-amdhsa-gfx90a.3.o %S/Inputs/dummy-elf.o; rm *-amdgcn-amd-amdhsa-gfx90a.3.o
 // RUN: not diff *-amdgcn-amd-amdhsa-gfx90c.4.o %S/Inputs/dummy-elf.o
+
+// Check that we can extract from an ELF object file
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \
+// RUN:   --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.out
+// RUN: clang-offload-packager %t.out \
+// RUN:   --image=file=%t-sm_70.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t-gfx908.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908
+// RUN: diff %t-sm_70.o %S/Inputs/dummy-elf.o
+// RUN: diff %t-gfx908.o %S/Inputs/dummy-elf.o
+
+// Check that we can extract from a bitcode file
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \
+// RUN:   --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-llvm -o %t.bc -fembed-offload-object=%t.out
+// RUN: clang-offload-packager %t.out \
+// RUN:   --image=file=%t-sm_70.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t-gfx908.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908
+// RUN: diff %t-sm_70.o %S/Inputs/dummy-elf.o
+// RUN: diff %t-gfx908.o %S/Inputs/dummy-elf.o
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to