[Clang][Bundler] Fix for a hang when unbundling fat binary

[sdmitriev]

Summary

clang-offload-bundler tool may hang under certain conditions when it extracts a subset of all available device bundles from the fat binary that is handled by the BinaryFileHandler. This patch fixes this problem.

[sdmitriev]

Rebase.

[ABataev]

Could you provide some more details about the problem?

[sdmitriev]

Sure. The main unbundling loop looks as follows (see loop on line 745)

while (!Worklist.empty()) {
  StringRef CurTriple = FH.get()->ReadBundleStart(Input);

  if (CurTriple.empty())
    break;

  auto Output = Worklist.find(CurTriple);
  if (Output == Worklist.end()) {
    continue;
  } 
 ...
}

Here Worklist is a collection of bundles that need to be extracted, and FH is the object implementing FileHandler interface. FileHandler::ReadBundleStart() returns the name of the bundle FH currently points to. As you can see in the loop above if the name returned by that call is not in the set of bundles that need to be extracted, we just call ReadBundleStart next time assuming that FileHandler would advance to the next bundle on each ReadBundleStart call. That assumption is correct for all FileHandler implementations except BinaryFileHandler which does not advance to the next bundle when this method is called. As a result we are going into an infinite loop if we need to skip a bundle for the file handled by the BinaryFileHandler . This patch just fixes this problem in the BinaryFileHandler implementation.

[sdmitriev]

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

301 ↗

(On Diff #217232)

Just noticed that this line now looks redundant and should be removed.

[ABataev]

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

295 ↗

(On Diff #217485)

Maybe just:

const BundleInfo &Bundle = *CurBundleInfo;
std::advance(CurBundleInfo, 1);
return Bundle.first();

instead of all these changes?

[sdmitriev]

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

295 ↗

(On Diff #217485)

It could be done like this if CurBundleInfo was not used in ReadBundle.

[ABataev]

LG

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

 // RUN: diff %t.bc %t.res.bc
 // RUN: diff %t.empty %t.res.tgt1
 // RUN: diff %t.empty %t.res.tgt2
 // RUN: clang-offload-bundler -type=bc -targets=openmp-powerpc64le-ibm-linux-gnu,host-%itanium_abi_triple,openmp-x86_64-pc-linux-gnu -outputs=%t.res.tgt1,%t.res.bc,%t.res.tgt2 -inputs=%t.bc -unbundle
 // RUN: diff %t.bc %t.res.bc
 // 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-%itanium_abi_triple,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
 // have an already bundled file to check the unbundle and do a dry run on the
 // bundling as it cannot be tested in all host platforms that will run these
 // tests.
 //
 

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

struct BundleInfo final {
     BundleInfo(uint64_t Size, uint64_t Offset) : Size(Size), Offset(Offset) {}
   };
 
   /// Map between a triple and the corresponding bundle information.
   StringMap<BundleInfo> BundlesInfo;
 
   /// Iterator for the bundle information that is being read.
   StringMap<BundleInfo>::iterator CurBundleInfo;
+  StringMap<BundleInfo>::iterator NextBundleInfo;
 
 public:
   BinaryFileHandler() : FileHandler() {}
 
   ~BinaryFileHandler() final {}
 
   void ReadHeader(MemoryBuffer &Input) final {
     StringRef FC = Input.getBuffer();
for (uint64_t i = 0; i < NumberOfBundles; ++i) {
       if (!Offset || Offset + Size > FC.size())
         return;
 
       assert(BundlesInfo.find(Triple) == BundlesInfo.end() &&
              "Triple is duplicated??");
       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;
   }
 
   void ReadBundle(raw_fd_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);
   }