[clang-offload-bundler] Add new unbundling mode 'a' for archives (#2840)

This patch adds support for unbundling archives into 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.

Signed-off-by: Sergey Dmitriev <[email protected]>

Author: sndmitriev

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

@@ -46,6 +46,7 @@
 // CK-HELP: {{.*}}o {{.*}}- object
 // CK-HELP: {{.*}}gch {{.*}}- precompiled-header
 // CK-HELP: {{.*}}ast {{.*}}- clang AST file
+// CK-HELP: {{.*}}a {{.*}}- archive of objects
 // CK-HELP: {{.*}}ao {{.*}}- archive with one object; output is an unbundled object
 // CK-HELP: {{.*}}aoo {{.*}}- archive; output file is a list of unbundled objects
 // CK-HELP: {{.*}}-unbundle {{.*}}- Unbundle bundled file into several output files.
@@ -341,6 +342,17 @@
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
 
+// Check archive mode.
+// 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-TGT1-LIST
+// RUN: llvm-ar t %t.tgt2.a | FileCheck %s --check-prefix=CHECK-AR-TGT2-LIST
+
+// CHECK-AR-TGT1-LIST: openmp-powerpc64le-ibm-linux-gnu.{{.+}}.bundle3.o
+// CHECK-AR-TGT1-LIST: openmp-powerpc64le-ibm-linux-gnu.{{.+}}.bundle4.o
+// CHECK-AR-TGT2-LIST: openmp-x86_64-pc-linux-gnu.{{.+}}.bundle3.o
+// CHECK-AR-TGT2-LIST: openmp-x86_64-pc-linux-gnu.{{.+}}.bundle4.o
+
 // Some code so that we can create a binary out of this file.
 int A = 0;
 void test_func(void) {

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

@@ -24,6 +24,7 @@
 #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"
@@ -88,6 +89,7 @@ static cl::opt<std::string> FilesType(
              "  oo  - object; output file is a list of unbundled objects\n"
              "  gch - precompiled-header\n"
              "  ast - clang AST file\n"
+             "  a   - archive of objects\n"
              "  ao  - archive with one object; output is an unbundled object\n"
              "  aoo - archive; output file is a list of unbundled objects\n"),
     cl::cat(ClangOffloadBundlerCategory));
@@ -161,7 +163,7 @@ class FileHandler {
   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.
@@ -342,7 +344,7 @@ class BinaryFileHandler final : public FileHandler {
     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,
@@ -624,7 +626,7 @@ class ObjectFileHandler final : public FileHandler {
 
   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 {
     assert(CurBundle != TripleToBundleInfo.end() &&
            "all bundles have been read already");
 
@@ -899,7 +901,7 @@ class TextFileHandler final : public FileHandler {
     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;
 
@@ -955,11 +957,26 @@ class ArchiveFileHandler final : public FileHandler {
   StringMap<unsigned>::iterator CurrBundle = Bundles.end();
   StringMap<unsigned>::iterator NextBundle = Bundles.end();
 
+  /// Output mode for the archive unbundler.
+  enum class OutputType {
+    Unknown,
+    FileList, // Output is a list file with extracted object file names
+    Object,   // Output is a single object file
+    Archive   // Output is an archive with extracted objects
+  };
+  const OutputType Mode = StringSwitch<OutputType>(FilesType)
+                              .Case("aoo", OutputType::FileList)
+                              .Case("ao", OutputType::Object)
+                              .Case("a", OutputType::Archive)
+                              .Default(OutputType::Unknown);
+
 public:
   ArchiveFileHandler() = default;
   ~ArchiveFileHandler() = default;
 
   Error ReadHeader(MemoryBuffer &Input) override {
+    assert(Mode != OutputType::Unknown && "unknown output mode");
+
     // Create archive instance for the given input.
     auto ArOrErr = Archive::create(Input);
     if (!ArOrErr)
@@ -1014,18 +1031,27 @@ class ArchiveFileHandler final : public FileHandler {
 
   Error ReadBundleEnd(MemoryBuffer &Input) override { return Error::success(); }
 
-  Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) override {
+  Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) override {
     assert(CurrBundle->second && "attempt to extract nonexistent bundle");
 
-    bool FileListMode = FilesType == "aoo";
-
     // In single-file mode we do not expect to see bundle more than once.
-    if (!FileListMode && CurrBundle->second > 1)
+    if (Mode == OutputType::Object && CurrBundle->second > 1)
       return createStringError(
           errc::invalid_argument,
           "'ao' file type is requested, but the archive contains multiple "
           "device objects; use 'aoo' instead");
 
+    // For 'host' archive bundle just copy input data to the output stream.
+    if (Mode == OutputType::Archive && 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 children.
     Error Err = Error::success();
     for (auto &C : Ar->children(Err)) {
@@ -1043,6 +1069,10 @@ class ArchiveFileHandler final : public FileHandler {
       auto Obj = std::unique_ptr<ObjectFile>(cast<ObjectFile>(Bin.release()));
       auto Buf = MemoryBuffer::getMemBuffer(Obj->getMemoryBufferRef(), false);
 
+      auto ChildNameOrErr = C.getName();
+      if (!ChildNameOrErr)
+        return ChildNameOrErr.takeError();
+
       ObjectFileHandler OFH(std::move(Obj));
       if (Error Err = OFH.ReadHeader(*Buf))
         return Err;
@@ -1052,10 +1082,9 @@ class ArchiveFileHandler final : public FileHandler {
       while (*NameOrErr) {
         auto TT = **NameOrErr;
         if (TT == CurrBundle->first()) {
-          // This is the bundle we are looking for. Create temporary file where
-          // the device part will be extracted if we are in the file-list mode,
-          // or write directly to the output file otherwise.
-          if (FileListMode) {
+          // This is the bundle we are looking for.
+          if (Mode == OutputType::FileList) {
+            // Create temporary file where the device part will be extracted to.
             SmallString<128u> ChildFileName;
             auto EC = sys::fs::createTemporaryFile(TempFileNameBase, "o",
                                                    ChildFileName);
@@ -1075,8 +1104,23 @@ class ArchiveFileHandler final : public FileHandler {
             // Add temporary file name with the device part to the output file
             // list.
             OS << ChildFileName << "\n";
-          } else if (Error Err = OFH.ReadBundle(OS, *Buf))
-            return Err;
+          } else if (Mode == OutputType::Object) {
+            // Extract the bundle to the output file in single file mode.
+            if (Error Err = OFH.ReadBundle(OS, *Buf))
+              return Err;
+          } else if (Mode == OutputType::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;
         }
@@ -1087,6 +1131,23 @@ class ArchiveFileHandler final : public FileHandler {
     }
     if (Err)
       return Err;
+
+    if (Mode == OutputType::Archive) {
+      // Determine archive kind for the 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, /*WriteSymtab=*/true, ArKind,
+                               /*Deterministic=*/true, /*Thin=*/false);
+      if (!NewAr)
+        return NewAr.takeError();
+      OS << NewAr.get()->getBuffer();
+    }
     return Error::success();
   }
 
@@ -1152,7 +1213,7 @@ CreateFileHandler(MemoryBuffer &FirstInput) {
     return std::make_unique<BinaryFileHandler>();
   if (FilesType == "ast")
     return std::make_unique<BinaryFileHandler>();
-  if (FilesType == "ao" || FilesType == "aoo")
+  if (FilesType == "a" || FilesType == "ao" || FilesType == "aoo")
     return std::make_unique<ArchiveFileHandler>();
 
   return createStringError(errc::invalid_argument,
@@ -1163,7 +1224,7 @@ CreateFileHandler(MemoryBuffer &FirstInput) {
 static Error BundleFiles() {
   std::error_code EC;
 
-  if (FilesType == "ao" || FilesType == "aoo")
+  if (FilesType == "a" || FilesType == "ao" || FilesType == "aoo")
     return createStringError(errc::invalid_argument,
                              "bundling is not supported for archives");