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

[jhuber6]

Summary

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.

[tra]

Should these be merged into a public interface via Object/OffloadBinary.h?

I'm all for consolidating relevant code.

[jhuber6]

In D132607#3747394, @tra wrote:

Should these be merged into a public interface via Object/OffloadBinary.h?

I'm all for consolidating relevant code.

Basically it would be a free function doing that extractFromBuffer does here. I would also need to define a wrapper around OwningBinary<OffloadBinary> as that's what I use in the linker wrapper.

I'm not sure if it would be easier to consolidate this before or after this patch.

[tra]

before or after this patch.

That would be your call. I personally would be biased towards doing refactoring early. Once something is in place, the temptation not to fix what already works might win.

[jhuber6]

Updating after refactoring extraction code.

[jhuber6]

CMake changes no longer needed

[yaxunl]

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

17–21

Are these include files necessary? I do not see code changes that need new include files. Or they were missing before?

[jhuber6]

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

17–21

You're right, they were included previously but I since rolled that logic into D132689. I'll remove them.

[jhuber6]

Removing unused headers.

[saiislam]

LGTM. Thanks!

clang/test/Driver/offload-packager.c

 // RUN:   --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx90a \
 // RUN:   --image=file=%S/Inputs/dummy-elf.o,kind=hip,triple=amdgcn-amd-amdhsa,arch=gfx90c 
 // RUN: cd $(dirname "%t") && clang-offload-packager %t --image=kind=openmp
 // RUN: diff *-nvptx64-nvidia-cuda-sm_70.0.o %S/Inputs/dummy-elf.o; rm *-nvptx64-nvidia-cuda-sm_70.0.o
 // RUN: diff *-nvptx64-nvidia-cuda-sm_80.1.o %S/Inputs/dummy-elf.o; rm *-nvptx64-nvidia-cuda-sm_80.1.o
 // 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

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

 //===-- clang-offload-packager/ClangOffloadPackager.cpp - file bundler ---===//
 //
 // 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
 //
 //===---------------------------------------------------------------------===//
 //
 // This tool takes several device object files and bundles them into a single
 // binary image using a custom binary format. This is intended to be used to
 // embed many device files into an application to create a fat binary.
 //
 //===---------------------------------------------------------------------===//
 
 #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"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
 
 using namespace llvm;
 using namespace llvm::object;
if (!OutputOrErr)
     return OutputOrErr.takeError();
   std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
   std::copy(BinaryData.begin(), BinaryData.end(), Output->getBufferStart());
   if (Error E = Output->commit())
     return E;
   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);
   if (std::error_code EC = BufferOrErr.getError())
     return createFileError(InputFile, EC);
   std::unique_ptr<MemoryBuffer> Buffer = std::move(*BufferOrErr);
 
   // This data can be misaligned if extracted from an archive.
   if (!isAddrAligned(Align(OffloadBinary::getAlignment()),
                      Buffer->getBufferStart()))
     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) {
     BumpPtrAllocator Alloc;
     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;
         if (Key == "file")
           return true;
         if (Key == "kind")
           return Binary->getOffloadKind() == getOffloadKind(Value);
         return Binary->getString(Key) == Value;