[ClangPackager] Add an option to extract inputs to an archive

[jhuber6]

Summary

Currently we simply overwrite the output file if we get muliple matches
in the fatbinary. This patch introduces the --archive option which
allows us to combine all of the files into a static archive instead.
This is usefuly for creating a device specific static archive library
from a fatbinary.

[JonChesterfield]

LGTM. As a meta level comment, we have far too many of these binary munging sorts of tools.

[jhuber6]

In D153568#4445133, @JonChesterfield wrote:

LGTM. As a meta level comment, we have far too many of these binary munging sorts of tools.

Definitely agree. This is supposed to be a direct repalcement for clang-offload-bundler, but that is still in use by HIP and I don't know if it's something we could remove in a reasonable amount of time.

clang/test/Driver/offload-packager.c

 // RUN:   --image=file=%t.elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \
 // RUN:   --image=file=%t.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 %t.elf.o
 // RUN: diff %t-gfx908.o %t.elf.o
+
+// Check that we can extract from an archive file to an archive file.
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \
+// RUN:   --image=file=%t.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: llvm-ar rcs %t.a %t.o
+// RUN: clang-offload-packager %t.a --archive --image=file=%t-gfx908.a,arch=gfx908
+// RUN: llvm-ar t %t-gfx908.a 2>&1 | FileCheck %s
+// CHECK: {{.*}}.o

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

 // 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/BinaryFormat/Magic.h"
+#include "llvm/Object/ArchiveWriter.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"

 
 static cl::list<std::string>
     DeviceImages("image",
                  cl::desc("List of key and value arguments. Required keywords "
                           "are 'file' and 'triple'."),
                  cl::value_desc("<key>=<value>,..."),
                  cl::cat(ClangOffloadPackagerCategory));
 
+static cl::opt<bool>
+    CreateArchive("archive",
+                  cl::desc("Write extracted files to a static archive"),
+                  cl::cat(ClangOffloadPackagerCategory));
+
+/// Path of the current binary.
+static const char *PackagerExecutable;
+
 static void PrintVersion(raw_ostream &OS) {
   OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n';
 }
 
 // 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,
                                                         StringSaver &Saver) {
   DenseMap<StringRef, StringRef> Args;
   for (StringRef Arg : llvm::split(Image, ",")) {
     auto [Key, Value] = Arg.split("=");
     if (Args.count(Key))
       Args[Key] = Saver.save(Args[Key] + "," + Value);
     else
       Args[Key] = Value;
   }
 
   return Args;
 }
 
+static Error writeFile(StringRef Filename, StringRef Data) {
+  Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
+      FileOutputBuffer::create(Filename, Data.size());
+  if (!OutputOrErr)
+    return OutputOrErr.takeError();
+  std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
+  llvm::copy(Data, Output->getBufferStart());
+  if (Error E = Output->commit())
+    return E;
+  return Error::success();
+}
+
 static Error bundleImages() {
   SmallVector<char, 1024> BinaryData;
   raw_svector_ostream OS(BinaryData);
   for (StringRef Image : DeviceImages) {
     BumpPtrAllocator Alloc;
     StringSaver Saver(Alloc);
     DenseMap<StringRef, StringRef> Args = getImageArguments(Image, Saver);
 
for (StringRef Image : DeviceImages) {
     }
     std::unique_ptr<MemoryBuffer> Buffer = OffloadBinary::write(ImageBinary);
     if (Buffer->getBufferSize() % OffloadBinary::getAlignment() != 0)
       return createStringError(inconvertibleErrorCode(),
                                "Offload binary has invalid size alignment");
     OS << Buffer->getBuffer();
   }
 
-  Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
-      FileOutputBuffer::create(OutputFile, BinaryData.size());
-  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())
+  if (Error E = writeFile(OutputFile,
+                          StringRef(BinaryData.begin(), BinaryData.size())))
     return E;
   return Error::success();
 }
 
 static Error unbundleImages() {
   ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
       MemoryBuffer::getFileOrSTDIN(InputFile);
   if (std::error_code EC = BufferOrErr.getError())
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 = Binaries.size(); I != E; ++I) {
-      const auto *Binary = Binaries[I].getBinary();
+    SmallVector<const OffloadBinary *> Extracted;
+    for (const OffloadFile &File : Binaries) {
+      const auto *Binary = File.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;
       });
-      if (!Match)
-        continue;
+      if (Match)
+        Extracted.push_back(Binary);
+    }
+
+    if (Extracted.empty())
 
-      // If the user did not provide a filename derive one from the input and
-      // image.
-      StringRef Filename =
-          !Args.count("file")
-              ? Saver.save(sys::path::stem(InputFile) + "-" +
-                           Binary->getTriple() + "-" + Binary->getArch() + "." +
-                           std::to_string(I) + "." +
-                           getImageKindName(Binary->getImageKind()))
-              : Args["file"];
+    if (CreateArchive) {
+      if (!Args.count("file"))
+        return createStringError(inconvertibleErrorCode(),
+                                 "Image must have a 'file' argument.");
 
-      Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
-          FileOutputBuffer::create(Filename, Binary->getImage().size());
-      if (!OutputOrErr)
-        return OutputOrErr.takeError();
-      std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
-      llvm::copy(Binary->getImage(), Output->getBufferStart());
-      if (Error E = Output->commit())
+      SmallVector<NewArchiveMember> Members;
+      for (const OffloadBinary *Binary : Extracted)
+        Members.emplace_back(MemoryBufferRef(
+            Binary->getImage(),
+            Binary->getMemoryBufferRef().getBufferIdentifier()));
+
+      if (Error E = writeArchive(Args["file"], Members, true,
+                                 Archive::getDefaultKindForHost(), true, false,
+                                 nullptr))
         return E;
+    } else if (Args.count("file")) {
+      if (Extracted.size() > 1)
+        WithColor::warning(errs(), PackagerExecutable)
+            << "Multiple inputs match to a single file, '" << Args["file"]
+            << "'\n";
+      if (Error E = writeFile(Args["file"], Extracted.back()->getImage()))
+        return E;
+    } else {
+      uint64_t Idx = 0;
+      for (const OffloadBinary *Binary : Extracted) {
+        StringRef Filename =
+            Saver.save(sys::path::stem(InputFile) + "-" + Binary->getTriple() +
+                       "-" + Binary->getArch() + "." + std::to_string(Idx++) +
+                       "." + getImageKindName(Binary->getImageKind()));
+        if (Error E = writeFile(Filename, Binary->getImage()))
+          return E;
+      }
     }
   }
 
   return Error::success();
 }
 
 int main(int argc, const char **argv) {
   sys::PrintStackTraceOnErrorSignal(argv[0]);
   cl::HideUnrelatedOptions(ClangOffloadPackagerCategory);
   cl::SetVersionPrinter(PrintVersion);
   cl::ParseCommandLineOptions(
       argc, argv,
       "A utility for bundling several object files into a single binary.\n"
       "The output binary can then be embedded into the host section table\n"
       "to create a fatbinary containing offloading code.\n");
 
   if (Help) {
     cl::PrintHelpMessage();
     return EXIT_SUCCESS;
   }
 
+  PackagerExecutable = argv[0];
   auto reportError = [argv](Error E) {
     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
     return EXIT_FAILURE;
   };
 
   if (!InputFile.empty() && !OutputFile.empty())
     return reportError(
         createStringError(inconvertibleErrorCode(),