[Clang][Bundler] Do not require host triple for extracting device bundles (#571)

* [Clang][Bundler] Do not require host triple for extracting device bundles

Bundler currently requires host triple to be provided no matter if you are performing bundling or unbundling, but for unbundling operation such requirement is too restrictive. You may for example want to examine device part of the object for a particular offload target, but you have to extract host part as well even though you do not need it. Host triple isn't really needed for unbundling, so this patch removes that requirement.

This is a cherry-pick of two llvm.org patches
  https://reviews.llvm.org/D66598 (prerequisite for the second patch)
  https://reviews.llvm.org/D66601

I addition to that this patch includes a fix for a bundler assertion on debug build

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

Author: sndmitriev

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

@@ -170,14 +170,20 @@
 // RUN: diff %t.i %t.res.i
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
+// RUN: clang-offload-bundler -type=i -targets=openmp-powerpc64le-ibm-linux-gnu -outputs=%t.res.tgt1 -inputs=%t.bundle3.i -unbundle
+// RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: clang-offload-bundler -type=ii -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.ii,%t.res.tgt1,%t.res.tgt2 -inputs=%t.bundle3.ii -unbundle
 // RUN: diff %t.ii %t.res.ii
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
+// RUN: clang-offload-bundler -type=ii -targets=openmp-x86_64-pc-linux-gnu -outputs=%t.res.tgt2 -inputs=%t.bundle3.ii -unbundle
+// RUN: diff %t.tgt2 %t.res.tgt2
 // RUN: clang-offload-bundler -type=ll -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.ll,%t.res.tgt1,%t.res.tgt2 -inputs=%t.bundle3.ll -unbundle
 // RUN: diff %t.ll %t.res.ll
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
+// RUN: clang-offload-bundler -type=ll -targets=openmp-powerpc64le-ibm-linux-gnu -outputs=%t.res.tgt1 -inputs=%t.bundle3.ll -unbundle
+// RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: clang-offload-bundler -type=s -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.s,%t.res.tgt1,%t.res.tgt2 -inputs=%t.bundle3.s -unbundle
 // RUN: diff %t.s %t.res.s
 // RUN: diff %t.tgt1 %t.res.tgt1
@@ -186,6 +192,8 @@
 // RUN: diff %t.s %t.res.s
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
+// RUN: clang-offload-bundler -type=s -targets=openmp-x86_64-pc-linux-gnu -outputs=%t.res.tgt2 -inputs=%t.bundle3.s -unbundle
+// RUN: diff %t.tgt2 %t.res.tgt2
 
 // Check if we can unbundle a file with no magic strings.
 // RUN: clang-offload-bundler -type=s -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.s,%t.res.tgt1,%t.res.tgt2 -inputs=%t.s -unbundle
@@ -197,6 +205,10 @@
 // RUN: diff %t.empty %t.res.tgt1
 // RUN: diff %t.empty %t.res.tgt2
 
+// Check that bindler prints an error if given host bundle does not exist in the fat binary.
+// RUN: not clang-offload-bundler -type=s -targets=host-x86_64-xxx-linux-gnu,openmp-powerpc64le-ibm-linux-gnu -outputs=%t.res.s,%t.res.tgt1 -inputs=%t.bundle3.s -unbundle 2>&1 | FileCheck %s --check-prefix CK-NO-HOST-BUNDLE
+// CK-NO-HOST-BUNDLE: error: Can't find bundle for the host target
+
 //
 // Check binary bundle/unbundle. The content that we have before bundling must be the same we have after unbundling.
 //
@@ -208,10 +220,14 @@
 // RUN: diff %t.bc %t.res.bc
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
+// RUN: clang-offload-bundler -type=bc -targets=openmp-powerpc64le-ibm-linux-gnu -outputs=%t.res.tgt1 -inputs=%t.bundle3.bc -unbundle
+// RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: clang-offload-bundler -type=gch -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.gch,%t.res.tgt1,%t.res.tgt2 -inputs=%t.bundle3.gch -unbundle
 // RUN: diff %t.ast %t.res.gch
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
+// RUN: clang-offload-bundler -type=gch -targets=openmp-x86_64-pc-linux-gnu -outputs=%t.res.tgt2 -inputs=%t.bundle3.gch -unbundle
+// RUN: diff %t.tgt2 %t.res.tgt2
 // RUN: clang-offload-bundler -type=ast -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.ast,%t.res.tgt1,%t.res.tgt2 -inputs=%t.bundle3.ast -unbundle
 // RUN: diff %t.ast %t.res.ast
 // RUN: diff %t.tgt1 %t.res.tgt1
@@ -224,6 +240,8 @@
 // RUN: diff %t.ast %t.res.ast
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
+// RUN: clang-offload-bundler -type=ast -targets=openmp-powerpc64le-ibm-linux-gnu -outputs=%t.res.tgt1 -inputs=%t.bundle3.ast -unbundle
+// RUN: diff %t.tgt1 %t.res.tgt1
 
 // Check if we can unbundle a file with no magic strings.
 // RUN: clang-offload-bundler -type=bc -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.bc,%t.res.tgt1,%t.res.tgt2 -inputs=%t.bc -unbundle
@@ -235,6 +253,11 @@
 // RUN: diff %t.empty %t.res.tgt1
 // RUN: diff %t.empty %t.res.tgt2
 
+// Check that we do not have to unbundle all available bundles from the fat binary.
+// RUN: clang-offload-bundler -type=ast -targets=host-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.ast,%t.res.tgt2 -inputs=%t.bundle3.unordered.ast -unbundle
+// RUN: diff %t.ast %t.res.ast
+// RUN: diff %t.tgt2 %t.res.tgt2
+
 //
 // Check object bundle/unbundle. The content should be bundled into an ELF
 // section (we are using a PowerPC little-endian host which uses ELF). We
@@ -255,6 +278,8 @@
 // RUN: diff %t.2.o %t.res.o
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
+// RUN: clang-offload-bundler -type=o -targets=openmp-powerpc64le-ibm-linux-gnu -outputs=%t.res.tgt1 -inputs=%t.2.o -unbundle
+// RUN: diff %t.tgt1 %t.res.tgt1
 
 // Check if we can unbundle a file with no magic strings.
 // RUN: clang-offload-bundler -type=o -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.o,%t.res.tgt1,%t.res.tgt2 -inputs=%t.o -unbundle

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

@@ -257,6 +257,7 @@ class BinaryFileHandler final : public FileHandler {
 
   /// Iterator for the bundle information that is being read.
   StringMap<BundleInfo>::iterator CurBundleInfo;
+  StringMap<BundleInfo>::iterator NextBundleInfo;
 
 public:
   BinaryFileHandler() : FileHandler() {}
@@ -326,19 +327,19 @@ class BinaryFileHandler final : public FileHandler {
       BundlesInfo[Triple] = BundleInfo(Size, Offset);
     }
     // Set the iterator to where we will start to read.
-    CurBundleInfo = BundlesInfo.begin();
+    CurBundleInfo = BundlesInfo.end();
+    NextBundleInfo = BundlesInfo.begin();
   }
 
   StringRef ReadBundleStart(MemoryBuffer &Input) final {
-    if (CurBundleInfo == BundlesInfo.end())
+    if (NextBundleInfo == BundlesInfo.end())
       return StringRef();
-
+    CurBundleInfo = NextBundleInfo++;
     return CurBundleInfo->first();
   }
 
   void ReadBundleEnd(MemoryBuffer &Input) final {
     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
-    ++CurBundleInfo;
   }
 
   using FileHandler::ReadBundle; // to avoid hiding via the overload below
@@ -1247,15 +1248,16 @@ static bool UnbundleFiles() {
       break;
 
     auto Output = Worklist.find(CurTriple);
-
-    // Read the bundle if triple is included in targets
-    if (Output != Worklist.end()) {
-      // Check if the output file can be opened and copy the bundle to it.
-      FH->ReadBundle(Output->second, Input);
-      Worklist.erase(Output);
+    // The file may have more bundles for other targets, that we don't care
+    // about. Therefore, move on to the next triple
+    if (Output == Worklist.end()) {
+      continue;
     }
-    
+
+    // Check if the output file can be opened and copy the bundle to it.
+    FH->ReadBundle(Output->second, Input);
     FH->ReadBundleEnd(Input);
+    Worklist.erase(Output);
 
     // Record if we found the host bundle.
     if (hasHostKind(CurTriple))
@@ -1282,8 +1284,9 @@ static bool UnbundleFiles() {
     return false;
   }
 
-  // If we found elements, we emit an error if none of those were for the host.
-  if (!FoundHostBundle) {
+  // If we found elements, we emit an error if none of those were for the host
+  // in case host bundle name was provided in command line.
+  if (!FoundHostBundle && HostInputIndex != ~0u) {
     errs() << "error: Can't find bundle for the host target\n";
     return true;
   }
@@ -1340,8 +1343,8 @@ static bool CheckBundledSection() {
 
     if(CurTriple == triple) {
       found = true;
+      break;
     }
-    FH->ReadBundleEnd(Input);
   }
   return found;
 }
@@ -1457,7 +1460,10 @@ int main(int argc, const char **argv) {
     ++Index;
   }
 
-  if (!CheckSection && HostTargetNum != 1) {
+  // Host triple is not really needed for unbundling operation, so do not
+  // treat missing host triple as error if we do unbundling.
+  if (!CheckSection &&
+      ((Unbundle && HostTargetNum > 1) || (!Unbundle && HostTargetNum != 1))) {
     Error = true;
     errs() << "error: expecting exactly one host target but got "
            << HostTargetNum << ".\n";