[clang-offload-bundler] Add support for unbundling archives with fat objects

[sdmitriev]

Summary

This patch adds support for unbundling archives under a new file type 'a'.
Input file for this mode is expected to be an archive with fat object files
and outputs (one per offload target) will be archives with extracted device
specific parts from the input's objects.

[sdmitriev]

Fixed warnings from pre-merge checks.

[JonChesterfield]

If this splits an archive of bundles into N archives, it presumably doesn't suffer the same implicit whole-archive semantics that result from passing code through llvm-link. That seems great.

[sdmitriev]

In D94005#2477031, @JonChesterfield wrote:

If this splits an archive of bundles into N archives, it presumably doesn't suffer the same implicit whole-archive semantics that result from passing code through llvm-link. That seems great.

This patch implements similar functionality to D80816/D93525 (except adding '-<device>' suffix to the target to be extracted) but it does not restrict users to specify only one target and one output file for unbundling. So, if you need to unbundle more than one target from input archive you will be able to do it in one clang-offload-bundler invocation as opposed to running clang-offload-bundler once per each target with D80816/D93525. That is probably unrelated to the whole-archive semantics when linking archives with llvm-link.

[sdmitriev]

But regarding llvm-link behavior wrt archives, I assume it will be possible to link bitcode archives as libraries (i.e. link only required symbols) with –only-needed option after applying D92535.

clang/test/Driver/clang-offload-bundler.c

 // CK-HELP: {{.*}}i {{.*}}- cpp-output
 // CK-HELP: {{.*}}ii {{.*}}- c++-cpp-output
 // CK-HELP: {{.*}}ll {{.*}}- llvm
 // CK-HELP: {{.*}}bc {{.*}}- llvm-bc
 // CK-HELP: {{.*}}s {{.*}}- assembler
 // CK-HELP: {{.*}}o {{.*}}- object
 // CK-HELP: {{.*}}gch {{.*}}- precompiled-header
 // CK-HELP: {{.*}}ast {{.*}}- clang AST file
+// CK-HELP: {{.*}}a {{.*}}- archive of objects
 // CK-HELP: {{.*}}-unbundle {{.*}}- Unbundle bundled file into several output files.
 
 //
 // Check errors.
 //
 // RUN: not clang-offload-bundler -type=i -targets=host-%itanium_abi_triple,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -inputs=%t.i,%t.tgt1,%t.tgt2 -outputs=%t.bundle.i -unbundle 2>&1 | FileCheck %s --check-prefix CK-ERR1
 // CK-ERR1: error: only one input file supported in unbundling mode
 // CK-ERR1: error: number of output files and targets should match in unbundling mode

 // CK-ERR8: error: invalid target 'openxp-pxxerpc64le-ibm-linux-gnu', unknown offloading kind 'openxp', unknown target triple 'pxxerpc64le-ibm-linux-gnu'
 // CK-ERR8: error: invalid target 'xpenmp-x86_xx-pc-linux-gnu', unknown offloading kind 'xpenmp', unknown target triple 'x86_xx-pc-linux-gnu'
 
 // RUN: not clang-offload-bundler -type=i -targets=openmp-powerpc64le-linux,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -inputs=%t.i,%t.tgt1,%t.tgt2 -outputs=%t.bundle.i 2>&1 | FileCheck %s --check-prefix CK-ERR9A
 // RUN: not clang-offload-bundler -type=i -targets=host-%itanium_abi_triple,host-%itanium_abi_triple,openmp-x86_64-pc-linux-gnu -inputs=%t.i,%t.tgt1,%t.tgt2 -outputs=%t.bundle.i 2>&1 | FileCheck %s --check-prefix CK-ERR9B
 // CK-ERR9A: error: expecting exactly one host target but got 0
 // CK-ERR9B: error: Duplicate targets are not allowed
 
+// RUN: not clang-offload-bundler -type=a -targets=host-%itanium_abi_triple,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -inputs=%t.i,%t.tgt1,%t.tgt2 -outputs=%t.bundle.a 2>&1 | FileCheck %s --check-prefix CK-ERR10
+// CK-ERR10: error: bundling is not supported for archives
+
 //
 // Check text bundle. This is a readable format, so we check for the format we expect to find.
 //
 // RUN: clang-offload-bundler -type=i -targets=host-%itanium_abi_triple,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -inputs=%t.i,%t.tgt1,%t.tgt2 -outputs=%t.bundle3.i
 // RUN: clang-offload-bundler -type=ii -targets=host-%itanium_abi_triple,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -inputs=%t.ii,%t.tgt1,%t.tgt2 -outputs=%t.bundle3.ii
 // RUN: clang-offload-bundler -type=ll -targets=host-%itanium_abi_triple,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -inputs=%t.ll,%t.tgt1,%t.tgt2 -outputs=%t.bundle3.ll
 // RUN: clang-offload-bundler -type=s -targets=host-%itanium_abi_triple,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -inputs=%t.s,%t.tgt1,%t.tgt2 -outputs=%t.bundle3.s
 // RUN: clang-offload-bundler -type=s -targets=openmp-powerpc64le-ibm-linux-gnu,host-%itanium_abi_triple,openmp-x86_64-pc-linux-gnu -inputs=%t.tgt1,%t.s,%t.tgt2 -outputs=%t.bundle3.unordered.s

 //
 // RUN: not clang-offload-bundler -type=bc -targets=host-%itanium_abi_triple,hip-amdgcn-amd-amdhsa-gfx900,hip-amdgcn-amd-amdhsa-gfx900 \
 // RUN:   -inputs=%t.bc,%t.tgt1,%t.tgt1 -outputs=%t.hip.bundle.bc 2>&1 | FileCheck -check-prefix=DUP %s
 // RUN: not clang-offload-bundler -type=bc -inputs=%t.hip.bundle.bc -outputs=%t.tmp.bc,%t.tmp2.bc -unbundle \
 // RUN:   -targets=hip-amdgcn-amd-amdhsa-gfx906,hip-amdgcn-amd-amdhsa-gfx906 \
 // RUN:   2>&1 | FileCheck -check-prefix=DUP %s
 // DUP: error: Duplicate targets are not allowed
 
+//
+// Check archive unbundling.
+//
+// RUN: rm -f %t.a
+// RUN: llvm-ar crv %t.a %t.bundle3.o
+// RUN: clang-offload-bundler -type=a -targets=host-%itanium_abi_triple,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.host.a,%t.tgt1.a,%t.tgt2.a -inputs=%t.a -unbundle
+// RUN: cmp %t.host.a %t.a
+// RUN: llvm-ar t %t.tgt1.a | FileCheck %s --check-prefix=CHECK-AR-LIST1
+// RUN: llvm-ar t %t.tgt2.a | FileCheck %s --check-prefix=CHECK-AR-LIST2
+//
+// CHECK-AR-LIST1: openmp-powerpc64le-ibm-linux-gnu.{{.+}}.bundle3.o
+// CHECK-AR-LIST2: openmp-x86_64-pc-linux-gnu.{{.+}}.bundle3.o
+
 // Some code so that we can create a binary out of this file.
 int A = 0;
 void test_func(void) {
   ++A;
 }

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

 //===----------------------------------------------------------------------===//
 
 #include "clang/Basic/Version.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/Object/Archive.h"
+#include "llvm/Object/ArchiveWriter.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/FileSystem.h"
FilesType("type", cl::Required,
                        "  ii  - c++-cpp-output\n"
                        "  cui - cuda/hip-output\n"
                        "  d   - dependency\n"
                        "  ll  - llvm\n"
                        "  bc  - llvm-bc\n"
                        "  s   - assembler\n"
                        "  o   - object\n"
                        "  gch - precompiled-header\n"
-                       "  ast - clang AST file"),
+                       "  ast - clang AST file\n"
+                       "  a   - archive of objects"),
               cl::cat(ClangOffloadBundlerCategory));
 static cl::opt<bool>
     Unbundle("unbundle",
              cl::desc("Unbundle bundled file into several output files.\n"),
              cl::init(false), cl::cat(ClangOffloadBundlerCategory));
 
 static cl::opt<bool> PrintExternalCommands(
     "###",
public:
   /// more bundles to be read.
   virtual Expected<Optional<StringRef>>
   ReadBundleStart(MemoryBuffer &Input) = 0;
 
   /// Read the marker that closes the current bundle.
   virtual Error ReadBundleEnd(MemoryBuffer &Input) = 0;
 
   /// Read the current bundle and write the result into the stream \a OS.
-  virtual Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) = 0;
+  virtual Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) = 0;
 
   /// Write the header of the bundled file to \a OS based on the information
   /// gathered from \a Inputs.
   virtual Error WriteHeader(raw_fd_ostream &OS,
                             ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) = 0;
 
   /// Write the marker that initiates a bundle for the triple \a TargetTriple to
   /// \a OS.
Expected<Optional<StringRef>> ReadBundleStart(MemoryBuffer &Input) final {
     return CurBundleInfo->first();
   }
 
   Error ReadBundleEnd(MemoryBuffer &Input) final {
     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
     return Error::success();
   }
 
-  Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
+  Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
     StringRef FC = Input.getBuffer();
     OS.write(FC.data() + CurBundleInfo->second.Offset,
              CurBundleInfo->second.Size);
     return Error::success();
   }
 
   Error WriteHeader(raw_fd_ostream &OS,
while (NextSection != Obj->section_end()) {
       if (*TripleOrErr)
         return **TripleOrErr;
     }
     return None;
   }
 
   Error ReadBundleEnd(MemoryBuffer &Input) final { return Error::success(); }
 
-  Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
+  Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
     Expected<StringRef> ContentOrErr = CurrentSection->getContents();
     if (!ContentOrErr)
       return ContentOrErr.takeError();
     StringRef Content = *ContentOrErr;
 
     // Copy fat object contents to the output when extracting host bundle.
     if (Content.size() == 1u && Content.front() == 0)
       Content = StringRef(Input.getBufferStart(), Input.getBufferSize());
Error ReadBundleEnd(MemoryBuffer &Input) final {
     size_t TripleEnd = ReadChars = FC.find("\n", ReadChars + 1);
     if (TripleEnd != FC.npos)
       // Next time we read after the new line.
       ++ReadChars;
 
     return Error::success();
   }
 
-  Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
+  Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
     StringRef FC = Input.getBuffer();
     size_t BundleStart = ReadChars;
 
     // Find end of the bundle.
     size_t BundleEnd = ReadChars = FC.find(BundleEndString, ReadChars);
 
     StringRef Bundle(&FC.data()[BundleStart], BundleEnd - BundleStart);
     OS << Bundle;
TextFileHandler(StringRef Comment)
       : FileHandler(), Comment(Comment), ReadChars(0) {
     BundleStartString =
         "\n" + Comment.str() + " " OFFLOAD_BUNDLER_MAGIC_STR "__START__ ";
     BundleEndString =
         "\n" + Comment.str() + " " OFFLOAD_BUNDLER_MAGIC_STR "__END__ ";
   }
 };
 
+/// Archive file handler. Only unbundling is supported so far.
+class ArchiveFileHandler final : public FileHandler {
+  /// Archive we are dealing with.
+  std::unique_ptr<Archive> Ar;
+
+  /// Set of bundle names from all archive objects.
+  StringSet<> Bundles;
+
+  /// Iterators over the bundle names.
+  StringSet<>::iterator CurrBundle = Bundles.end();
+  StringSet<>::iterator NextBundle = Bundles.end();
+
+public:
+  ArchiveFileHandler() = default;
+  ~ArchiveFileHandler() = default;
+
+  Error ReadHeader(MemoryBuffer &Input) override {
+    // Create archive instance from the given input.
+    auto ArOrErr = Archive::create(Input);
+    if (!ArOrErr)
+      return ArOrErr.takeError();
+    Ar = std::move(*ArOrErr);
+
+    // Read all archive children.
+    Error Err = Error::success();
+    for (auto &C : Ar->children(Err)) {
+      auto BinOrErr = C.getAsBinary();
+      if (!BinOrErr) {
+        if (auto Err = isNotObjectErrorInvalidFileType(BinOrErr.takeError()))
+          return Err;
+        continue;
+      }
+
+      // We are interested in object files only, so ignore the child if it is
+      // not an object.
+      auto &Bin = BinOrErr.get();
+      if (!Bin->isObject())
+        continue;
+
+      auto Obj = std::unique_ptr<ObjectFile>(cast<ObjectFile>(Bin.release()));
+      auto Buf = MemoryBuffer::getMemBuffer(Obj->getMemoryBufferRef(), false);
+
+      // Create object file handler for the child object and collect the list of
+      // bundles from it.
+      ObjectFileHandler OFH(std::move(Obj));
+      if (Error Err = OFH.ReadHeader(*Buf))
+        return Err;
+      Expected<Optional<StringRef>> NameOrErr = OFH.ReadBundleStart(*Buf);
+      if (!NameOrErr)
+        return NameOrErr.takeError();
+      while (*NameOrErr) {
+        Bundles.insert(**NameOrErr);
+        NameOrErr = OFH.ReadBundleStart(*Buf);
+        if (!NameOrErr)
+          return NameOrErr.takeError();
+      }
+    }
+    if (Err)
+      return Err;
+
+    CurrBundle = Bundles.end();
+    NextBundle = Bundles.begin();
+    return Error::success();
+  }
+
+  Expected<Optional<StringRef>> ReadBundleStart(MemoryBuffer &Input) override {
+    if (NextBundle == Bundles.end())
+      return None;
+    CurrBundle = NextBundle++;
+    return CurrBundle->first();
+  }
+
+  Error ReadBundleEnd(MemoryBuffer &Input) override { return Error::success(); }
+
+  Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) override {
+    // For 'host' archive bundle just copy input data to the output stream.
+    if (hasHostKind(CurrBundle->first())) {
+      OS << Input.getBuffer();
+      return Error::success();
+    }
+
+    // Extracted objects data for archive mode.
+    SmallVector<std::string, 8u> ArNames;
+    SmallVector<SmallVector<char, 0u>, 8u> ArData;
+    SmallVector<NewArchiveMember, 8u> ArMembers;
+
+    // Read all archive children.
+    Error Err = Error::success();
+    for (auto &C : Ar->children(Err)) {
+      auto BinOrErr = C.getAsBinary();
+      if (!BinOrErr) {
+        if (auto Err = isNotObjectErrorInvalidFileType(BinOrErr.takeError()))
+          return Err;
+        continue;
+      }
+
+      // We are interested in object files only, so ignore the child if it is
+      // not an object.
+      auto &Bin = BinOrErr.get();
+      if (!Bin->isObject())
+        continue;
+
+      auto Obj = std::unique_ptr<ObjectFile>(cast<ObjectFile>(Bin.release()));
+      auto Buf = MemoryBuffer::getMemBuffer(Obj->getMemoryBufferRef(), false);
+
+      // Get the child name. We will use it for the target archive member.
+      auto ChildNameOrErr = C.getName();
+      if (!ChildNameOrErr)
+        return ChildNameOrErr.takeError();
+
+      ObjectFileHandler OFH(std::move(Obj));
+      if (Error Err = OFH.ReadHeader(*Buf))
+        return Err;
+      Expected<Optional<StringRef>> NameOrErr = OFH.ReadBundleStart(*Buf);
+      if (!NameOrErr)
+        return NameOrErr.takeError();
+      while (*NameOrErr) {
+        auto TT = **NameOrErr;
+        if (TT == CurrBundle->first()) {
+          // This is the bundle we are looking for. Extract target bundle from
+          // the child and add it to the list of target archive members. Use
+          // bundle name as a prefix for the child name is the target archive.
+          auto &Name = ArNames.emplace_back((TT + "." + *ChildNameOrErr).str());
+          auto &Data = ArData.emplace_back();
+          raw_svector_ostream ChildOS{Data};
+
+          // Extract the bundle.
+          if (Error Err = OFH.ReadBundle(ChildOS, *Buf))
+            return Err;
+
+          ArMembers.emplace_back(
+              MemoryBufferRef{StringRef(Data.data(), Data.size()), Name});
+
+          if (Error Err = OFH.ReadBundleEnd(*Buf))
+            return Err;
+        }
+
+        NameOrErr = OFH.ReadBundleStart(*Buf);
+        if (!NameOrErr)
+          return NameOrErr.takeError();
+      }
+    }
+    if (Err)
+      return Err;
+
+    // Determine archive kind to use for this offload target.
+    StringRef TargetKind;
+    StringRef TargetTriple;
+    getOffloadKindAndTriple(CurrBundle->first(), TargetKind, TargetTriple);
+    auto ArKind =
+        Triple(TargetTriple).isOSDarwin() ? Archive::K_DARWIN : Archive::K_GNU;
+
+    // And write archive to the output.
+    Expected<std::unique_ptr<MemoryBuffer>> NewAr =
+        writeArchiveToBuffer(ArMembers, true, ArKind, true, false);
+    if (!NewAr)
+      return NewAr.takeError();
+    OS << NewAr.get()->getBuffer();
+
+    return Error::success();
+  }
+
+  Error WriteHeader(raw_fd_ostream &OS,
+                    ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) override {
+    llvm_unreachable("unsupported for the ArchiveFileHandler");
+  }
+
+  Error WriteBundleStart(raw_fd_ostream &OS, StringRef TargetTriple) override {
+    llvm_unreachable("unsupported for the ArchiveFileHandler");
+  }
+
+  Error WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) override {
+    llvm_unreachable("unsupported for the ArchiveFileHandler");
+  }
+
+  Error WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) override {
+    llvm_unreachable("unsupported for the ArchiveFileHandler");
+  }
+};
+
 /// Return an appropriate object file handler. We use the specific object
 /// handler if we know how to deal with that format, otherwise we use a default
 /// binary file handler.
 static std::unique_ptr<FileHandler>
 CreateObjectFileHandler(MemoryBuffer &FirstInput) {
   // Check if the input file format is one that we know how to deal with.
   Expected<std::unique_ptr<Binary>> BinaryOrErr = createBinary(FirstInput);
 
CreateFileHandler(MemoryBuffer &FirstInput) {
   if (FilesType == "s")
     return std::make_unique<TextFileHandler>(/*Comment=*/"#");
   if (FilesType == "o")
     return CreateObjectFileHandler(FirstInput);
   if (FilesType == "gch")
     return std::make_unique<BinaryFileHandler>();
   if (FilesType == "ast")
     return std::make_unique<BinaryFileHandler>();
+  if (FilesType == "a")
+    return std::make_unique<ArchiveFileHandler>();
 
   return createStringError(errc::invalid_argument,
                            "'" + FilesType + "': invalid file type specified");
 }
 
 /// Bundle the files. Return true if an error was found.
 static Error BundleFiles() {
   std::error_code EC;
 
+  if (FilesType == "a")
+    return createStringError(errc::invalid_argument,
+                             "bundling is not supported for archives");
+
   // Create output file.
   raw_fd_ostream OutputFile(OutputFileNames.front(), EC, sys::fs::OF_None);
   if (EC)
     return createFileError(OutputFileNames.front(), EC);
 
   // Open input files.
   SmallVector<std::unique_ptr<MemoryBuffer>, 8u> InputBuffers;
   InputBuffers.reserve(InputFileNames.size());