[clang] a69404c - [OffloadPackager] Add ability to extract images from other file types

[Joseph Huber via cfe-commits]

Author: Joseph Huber
Date: 2022-09-06T08:55:17-05:00
New Revision: a69404c0a294ce65432ce67d5f3e7dce28106496

URL: https://github.com/llvm/llvm-project/commit/a69404c0a294ce65432ce67d5f3e7dce28106496
DIFF: https://github.com/llvm/llvm-project/commit/a69404c0a294ce65432ce67d5f3e7dce28106496.diff

LOG: [OffloadPackager] Add ability to extract images from other file types

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.

Reviewed By: saiislam

Differential Revision: https://reviews.llvm.org/D132607

Added: 
    

Modified: 
    clang/test/Driver/offload-packager.c
    clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

Removed: 
    


################################################################################
diff  --git a/clang/test/Driver/offload-packager.c b/clang/test/Driver/offload-packager.c
index c4617d06e93d3..8d6ee50f2a190 100644
--- a/clang/test/Driver/offload-packager.c
+++ b/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

diff  --git a/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp b/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
index c9c722e0a5b5c..47ef155ef2783 100644
--- a/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
+++ b/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp
@@ -14,8 +14,7 @@
 
 #include "clang/Basic/Version.h"
 
-#include "llvm/Object/Binary.h"
-#include "llvm/Object/ObjectFile.h"
+#include "llvm/BinaryFormat/Magic.h"
 #include "llvm/Object/OffloadBinary.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileOutputBuffer.h"
@@ -123,29 +122,6 @@ static Error bundleImages() {
   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;
-  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();
-
-    Offset += (*BinaryOrErr)->getSize();
-    Binaries.emplace_back(std::move(*BinaryOrErr));
-  }
-
-  return std::move(Binaries);
-}
-
 static Error unbundleImages() {
   ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
       MemoryBuffer::getFileOrSTDIN(InputFile);
@@ -159,9 +135,9 @@ static Error unbundleImages() {
     Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(),
                                             Buffer->getBufferIdentifier());
 
-  auto BinariesOrErr = extractOffloadFiles(*Buffer);
-  if (!BinariesOrErr)
-    return BinariesOrErr.takeError();
+  SmallVector<OffloadFile> Binaries;
+  if (Error Err = extractOffloadBinaries(*Buffer, Binaries))
+    return Err;
 
   // Try to extract each device image specified by the user from the input file.
   for (StringRef Image : DeviceImages) {
@@ -169,8 +145,8 @@ static Error unbundleImages() {
     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 
diff erently.
       bool Match = llvm::all_of(Args, [&](auto &Arg) {
         const auto [Key, Value] = Arg;