[ClangOffloadBundler] fix unbundling archive (#84195)

When unbundling an archive, need to save the content of each object file
to a temporary file before passing it to llvm-objcopy, instead of
passing the original input archive file to llvm-objcopy.

Also allows extracting host bundles for archives.

Fixes: #83509

Author: yxsamliu

clang/lib/Driver/OffloadBundler.cpp

@@ -590,7 +590,8 @@ class ObjectFileHandler final : public FileHandler {
     // Copy fat object contents to the output when extracting host bundle.
     std::string ModifiedContent;
     if (Content.size() == 1u && Content.front() == 0) {
-      auto HostBundleOrErr = getHostBundle();
+      auto HostBundleOrErr = getHostBundle(
+          StringRef(Input.getBufferStart(), Input.getBufferSize()));
       if (!HostBundleOrErr)
         return HostBundleOrErr.takeError();
 
@@ -700,7 +701,7 @@ class ObjectFileHandler final : public FileHandler {
     return Error::success();
   }
 
-  Expected<std::string> getHostBundle() {
+  Expected<std::string> getHostBundle(StringRef Input) {
     TempFileHandlerRAII TempFiles;
 
     auto ModifiedObjPathOrErr = TempFiles.Create(std::nullopt);
@@ -715,7 +716,24 @@ class ObjectFileHandler final : public FileHandler {
     ObjcopyArgs.push_back("--regex");
     ObjcopyArgs.push_back("--remove-section=__CLANG_OFFLOAD_BUNDLE__.*");
     ObjcopyArgs.push_back("--");
-    ObjcopyArgs.push_back(BundlerConfig.InputFileNames.front());
+
+    StringRef ObjcopyInputFileName;
+    // When unbundling an archive, the content of each object file in the
+    // archive is passed to this function by parameter Input, which is different
+    // from the content of the original input archive file, therefore it needs
+    // to be saved to a temporary file before passed to llvm-objcopy. Otherwise,
+    // Input is the same as the content of the original input file, therefore
+    // temporary file is not needed.
+    if (StringRef(BundlerConfig.FilesType).starts_with("a")) {
+      auto InputFileOrErr =
+          TempFiles.Create(ArrayRef<char>(Input.data(), Input.size()));
+      if (!InputFileOrErr)
+        return InputFileOrErr.takeError();
+      ObjcopyInputFileName = *InputFileOrErr;
+    } else
+      ObjcopyInputFileName = BundlerConfig.InputFileNames.front();
+
+    ObjcopyArgs.push_back(ObjcopyInputFileName);
     ObjcopyArgs.push_back(ModifiedObjPath);
 
     if (Error Err = executeObjcopy(BundlerConfig.ObjcopyPath, ObjcopyArgs))
@@ -1628,10 +1646,8 @@ Error OffloadBundler::UnbundleArchive() {
     while (!CodeObject.empty()) {
       SmallVector<StringRef> CompatibleTargets;
       auto CodeObjectInfo = OffloadTargetInfo(CodeObject, BundlerConfig);
-      if (CodeObjectInfo.hasHostKind()) {
-        // Do nothing, we don't extract host code yet.
-      } else if (getCompatibleOffloadTargets(CodeObjectInfo, CompatibleTargets,
-                                             BundlerConfig)) {
+      if (getCompatibleOffloadTargets(CodeObjectInfo, CompatibleTargets,
+                                      BundlerConfig)) {
         std::string BundleData;
         raw_string_ostream DataStream(BundleData);
         if (Error Err = FileHandler->ReadBundle(DataStream, CodeObjectBuffer))

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

@@ -13,6 +13,19 @@
 // RUN: obj2yaml %t.o > %t.o.yaml
 // RUN: %clang -O0 -target %itanium_abi_triple %s -emit-ast -o %t.ast
 
+// RUN: echo 'void a() {}' >%t.a.cpp
+// RUN: echo 'void b() {}' >%t.b.cpp
+// RUN: %clang -target %itanium_abi_triple %t.a.cpp -c -o %t.a.o
+// RUN: %clang -target %itanium_abi_triple %t.b.cpp -c -o %t.b.o
+//
+// Remove .llvm_addrsig section since its offset changes after llvm-objcopy
+// removes clang-offload-bundler sections, therefore not good for comparison.
+//
+// RUN: llvm-objcopy --remove-section=.llvm_addrsig %t.a.o
+// RUN: llvm-objcopy --remove-section=.llvm_addrsig %t.b.o
+// RUN: obj2yaml %t.a.o > %t.a.yaml
+// RUN: obj2yaml %t.b.o > %t.b.yaml
+
 //
 // Generate an empty file to help with the checks of empty files.
 //
@@ -414,6 +427,25 @@
 // HIP-AR-906-DAG: hip_bundle1-hip-amdgcn-amd-amdhsa--gfx906
 // HIP-AR-906-DAG: hip_bundle2-hip-amdgcn-amd-amdhsa--gfx906
 
+//
+// Check unbundling archive for host target
+//
+// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,hip-amdgcn-amd-amdhsa--gfx900 \
+// RUN:   -input=%t.a.o -input=%t.tgt1 -output=%t.a.bundled.o
+// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,hip-amdgcn-amd-amdhsa--gfx900 \
+// RUN:   -input=%t.b.o -input=%t.tgt1 -output=%t.b.bundled.o
+// RUN: rm -f %t.bundled.a
+// RUN: llvm-ar cr %t.bundled.a %t.a.bundled.o %t.b.bundled.o
+// RUN: cp %t.bundled.a %t.bundled.a.bak
+// RUN: clang-offload-bundler -unbundle --targets=host-%itanium_abi_triple -type=a -input=%t.bundled.a -output=%t.host.a
+// RUN: rm -f *%itanium_abi_triple*.a.bundled.o *%itanium_abi_triple*.b.bundled.o
+// RUN: llvm-ar -x %t.host.a
+// RUN: diff %t.bundled.a %t.bundled.a.bak
+// RUN: obj2yaml *%itanium_abi_triple*.a.bundled.o > %t.a.unbundled.yaml
+// RUN: diff %t.a.unbundled.yaml %t.a.yaml
+// RUN: obj2yaml *%itanium_abi_triple*.b.bundled.o > %t.b.unbundled.yaml
+// RUN: diff %t.b.unbundled.yaml %t.b.yaml
+//
 // Check clang-offload-bundler reporting an error when trying to unbundle an archive but
 // the input file is not an archive.
 //