[Driver][OpenMP][CUDA] Add capability to bundle object files in sections of the host binary format.

[sfantao]

Summary

This patch adds the capability to bundle object files in sections of the host binary using a designated naming convention for these sections. This patch uses the functionality of the object reader already in the LLVM library to read bundled files, and invokes clang with the incremental linking options to create bundle files.

Bundling files involves creating an IR file with the contents of the bundle assigned as initializers of globals binded to the designated sections. This way the bundling implementation is agnostic of the host object format.

The features added by this patch were requested in the RFC discussion in http://lists.llvm.org/pipermail/cfe-dev/2016-February/047547.html.

[ABataev]

tools/clang-offload-bundler/ClangOffloadBundler.cpp

337–347

3 slashes, please

348

'final'?

350

No '\brief's

358

'const'? or 'static'?

463

Not sure that this a good solution to crash the compiler in this case. I think it must exit gracefully.

481

Again, maybe just emit an error an exit?

[sfantao]

• Remove \brief and return errors instead of crashing when temporary tiles can't be created.

[sfantao]

Hi Alexey,

Thanks for the review!

tools/clang-offload-bundler/ClangOffloadBundler.cpp

358

It can be static.

463

Ok, replaced the unreachable code with logic to emit an error and exist the tool.

481

Emitting an error now.

[sfantao]

• Start staic function name with caps.

[Hahnfeld]

tools/clang-offload-bundler/ClangOffloadBundler.cpp

724–737

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

/..//bin/ld: unrecognised emulation mode: elf64lppc
Supported emulations: elf_x86_64 elf32_x86_64 elf_i386 i386linux elf_l1om elf_k1om

and therefore fails.

I'm on an x86_64 Linux and obviously my GNU ld version 2.23.52.0.1-55.el7 20130226 doesn't support Power :-(

[sfantao]

Hi Jonas,

Thanks for trying that out!

tools/clang-offload-bundler/ClangOffloadBundler.cpp

724–737

Oh, right... I cannot run the bundler in the regression tests to tests the bundler. I guess I need some sort of dry run option to check the commands are correct without actually running them.

I'll fix that.

[Hahnfeld]

Another overall question: Back in February you said that it would be possible to have a "default" object that can be taken without knowledge of the bundler (see http://lists.llvm.org/pipermail/cfe-dev/2016-February/047555.html).

In my tests with the patch, this has not worked yet - am I missing something or is this not yet implemented?

tools/clang-offload-bundler/ClangOffloadBundler.cpp

724–737

Yes, comparable to -### in clang?

Another option (that I don't really prefer, just for completeness) would be to have separate tests that have appropriate REQUIRES...

[sfantao]

In D21851#477905, @Hahnfeld wrote:

Another overall question: Back in February you said that it would be possible to have a "default" object that can be taken without knowledge of the bundler (see http://lists.llvm.org/pipermail/cfe-dev/2016-February/047555.html).

In my tests with the patch, this has not worked yet - am I missing something or is this not yet implemented?

What I have proposed in the patch should be able to produce a bundled object file that can be read by other tools (the device image is embedded in designated sections, so these other tools may just ignore these sections). Also, if you feed the bundled an object file that does not have device sections, it should produce empty device files instead of just failing. That's what I was proposing there. Did you have something different in mind?

That's actually tested in the regression tests that goes with the patch. It won't work for you as is because of the ppc64le specific problem that I have to fix, but if you change that to x86_64 you may be able to test it.

tools/clang-offload-bundler/ClangOffloadBundler.cpp

724–737

Yes, -### is exactly what I was aiming at.

[sfantao]

That was probably not the reason why you were getting the failure, but the bundler was not prepared to bundle files if the host was not the first input. I fixed that.

I also replaced the test that was running the linker by something that checks the command and temporary files using -#. However, I'm still producing temporary files with -#. This problem is similar to the linker script test in http://reviews.llvm.org/D21847. Once we figure out how to tests the linker script I'll replicate the solution here too.

[sfantao]

• Allow host object to come in any order when generating the binary. Add testing for that.
• Add save temps and command print option.
• Add bundled object for testing.

[Hahnfeld]

In D21851#479042, @sfantao wrote:

That was probably not the reason why you were getting the failure, but the bundler was not prepared to bundle files if the host was not the first input. I fixed that.

This was indeed the problem, this now looks better and has the host object as default (though I don't yet understand why the order in -targets= is different for me...).

Thanks,
Jonas

[sfantao]

• Refactor code to dump contents of temporary file instead of creating the actual file in dry-run mode.
• Rebase.

[Hahnfeld]

LGTM

tools/clang-offload-bundler/ClangOffloadBundler.cpp

432

Unnecessary return, please remove as Alexey requested on the other patch

[sfantao]

• Remove redundant return statement.

test/Driver/clang-offload-bundler.c

 //
 // Generate all the types of files we can bundle.
 //
 // RUN: %clang -O0 -target powerpc64le-ibm-linux-gnu %s -E -o %t.i
 // RUN: %clangxx -O0 -target powerpc64le-ibm-linux-gnu -x c++ %s -E -o %t.ii
 // RUN: %clang -O0 -target powerpc64le-ibm-linux-gnu %s -S -emit-llvm -o %t.ll
 // RUN: %clang -O0 -target powerpc64le-ibm-linux-gnu %s -c -emit-llvm -o %t.bc
 // RUN: %clang -O0 -target powerpc64le-ibm-linux-gnu %s -S -o %t.s
+// RUN: %clang -O0 -target powerpc64le-ibm-linux-gnu %s -c -o %t.o
 // RUN: %clang -O0 -target powerpc64le-ibm-linux-gnu %s -emit-ast -o %t.ast
 
 //
 // Generate an empty file to help with the checks of empty files.
 //
 // RUN: touch %t.empty
 
 //

 // 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-powerpc64le-ibm-linux-gnu,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 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.
+//
+
+// RUN: clang-offload-bundler -type=o -targets=host-powerpc64le-ibm-linux-gnu,openmp-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -inputs=%t.o,%t.tgt1,%t.tgt2 -outputs=%t.bundle3.o -### -dump-temporary-files 2>&1 \
+// RUN: | FileCheck %s --check-prefix CK-OBJ-CMD
+// CK-OBJ-CMD: private constant [1 x i8] zeroinitializer, section "__CLANG_OFFLOAD_BUNDLE__host-powerpc64le-ibm-linux-gnu"
+// CK-OBJ-CMD: private constant [25 x i8] c"Content of device file 1\0A", section "__CLANG_OFFLOAD_BUNDLE__openmp-powerpc64le-ibm-linux-gnu"
+// CK-OBJ-CMD: private constant [25 x i8] c"Content of device file 2\0A", section "__CLANG_OFFLOAD_BUNDLE__openmp-x86_64-pc-linux-gnu"
+// CK-OBJ-CMD: clang" "-r" "-target" "powerpc64le-ibm-linux-gnu" "-o" "{{.+}}.o" "{{.+}}.o" "{{.+}}.bc" "-nostdlib"
+
+// 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=%s.o -unbundle
+// RUN: diff %s.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,host-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.tgt1,%t.res.o,%t.res.tgt2 -inputs=%s.o -unbundle
+// RUN: diff %s.o %t.res.o
+// RUN: diff %t.tgt1 %t.res.tgt1
+// RUN: diff %t.tgt2 %t.res.tgt2
+
+// 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
+// RUN: diff %t.o %t.res.o
+// RUN: diff %t.empty %t.res.tgt1
+// RUN: diff %t.empty %t.res.tgt2
+// RUN: clang-offload-bundler -type=o -targets=openmp-powerpc64le-ibm-linux-gnu,host-powerpc64le-ibm-linux-gnu,openmp-x86_64-pc-linux-gnu -outputs=%t.res.tgt1,%t.res.o,%t.res.tgt2 -inputs=%t.o -unbundle
+// RUN: diff %t.o %t.res.o
+// RUN: diff %t.empty %t.res.tgt1
+// RUN: diff %t.empty %t.res.tgt2
+
 // Some code so that we can create a binary out of this file.
 int A = 0;
 void test_func(void) {
   ++A;
 }

tools/clang-offload-bundler/ClangOffloadBundler.cpp

FilesType("type", cl::Required,
                        "  gch - precompiled-header\n"
                        "  ast - clang AST file"),
               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(
+    "###",
+    cl::desc("Print any external commands that are to be executed "
+             "instead of actually executing them - for testing purposes.\n"),
+    cl::init(false), cl::cat(ClangOffloadBundlerCategory));
+
+static cl::opt<bool> DumpTemporaryFiles(
+    "dump-temporary-files",
+    cl::desc("Dumps any temporary files created - for testing purposes.\n"),
+    cl::init(false), cl::cat(ClangOffloadBundlerCategory));
+
 /// Magic string that marks the existence of offloading data.
 #define OFFLOAD_BUNDLER_MAGIC_STR "__CLANG_OFFLOAD_BUNDLE__"
 
 /// The index of the host input in the list of inputs.
 static unsigned HostInputIndex = ~0u;
 
+/// Path to the current binary.
+static std::string BundlerExecutable;
+
 /// Obtain the offload kind and real machine triple out of the target
 /// information specified by the user.
 static void getOffloadKindAndTriple(StringRef Target, StringRef &OffloadKind,
                                     StringRef &Triple) {
   auto KindTriplePair = Target.split('-');
   OffloadKind = KindTriplePair.first;
   Triple = KindTriplePair.second;
 }
+static StringRef getTriple(StringRef Target) {
+  StringRef OffloadKind;
+  StringRef Triple;
+  getOffloadKindAndTriple(Target, OffloadKind, Triple);
+  return Triple;
+}
 static bool hasHostKind(StringRef Target) {
   StringRef OffloadKind;
   StringRef Triple;
   getOffloadKindAndTriple(Target, OffloadKind, Triple);
   return OffloadKind == "host";
 }
 
 /// Generic file handler interface.
public:
   /// Write the header of the bundled file to \a OS based on the information
   /// gathered from \a Inputs.
   virtual void 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.
   virtual void WriteBundleStart(raw_fd_ostream &OS, StringRef TargetTriple) = 0;
   /// Write the marker that closes a bundle for the triple \a TargetTriple to \a
-  /// OS.
-  virtual void WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) = 0;
+  /// OS. Return true if any error was found.
+  virtual bool WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) = 0;
   /// Write the bundle from \a Input into \a OS.
   virtual void WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) = 0;
 
   FileHandler() {}
   virtual ~FileHandler() {}
 };
 
 /// Handler for binary files. The bundled file will have the following format
for (auto &T : TargetNames) {
       HeaderSize += MB.getBufferSize();
       // Size of the triple
       Write8byteIntegerToBuffer(OS, T.size());
       // Triple
       OS << T;
     }
   }
   void WriteBundleStart(raw_fd_ostream &OS, StringRef TargetTriple) {}
-  void WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) {}
+  bool WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) {
+    return false;
+  }
   void WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) {
     OS.write(Input.getBufferStart(), Input.getBufferSize());
   }
 
   BinaryFileHandler() : FileHandler() {}
   ~BinaryFileHandler() {}
 };
 
+/// Handler for object files. The bundles are organized by sections with a
+/// designated name.
+///
+/// In order to bundle we create an IR file with the content of each section and
+/// use incremental linking to produce the resulting object. We also add section
+/// with a single byte to state the name of the component the main object file
+/// (the one we are bundling into) refers to.
+///
+/// To unbundle, we use just copy the contents of the designated section. If the
+/// requested bundle refer to the main object file, we just copy it with no
+/// changes.
+class ObjectFileHandler final : public FileHandler {
+
+  /// The object file we are currently dealing with.
+  ObjectFile &Obj;
+
+  /// Return the input file contents.
+  StringRef getInputFileContents() const { return Obj.getData(); }
+
+  /// Return true if the provided section is an offload section and return the
+  /// triple by reference.
+  static bool IsOffloadSection(SectionRef CurSection,
+                               StringRef &OffloadTriple) {
+    StringRef SectionName;
+    CurSection.getName(SectionName);
+
+    if (SectionName.empty())
+      return false;
+
+    // If it does not start with the reserved suffix, just skip this section.
+    if (!SectionName.startswith(OFFLOAD_BUNDLER_MAGIC_STR))
+      return false;
+
+    // Return the triple that is right after the reserved prefix.
+    OffloadTriple = SectionName.substr(sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1);
+    return true;
+  }
+
+  /// Total number of inputs.
+  unsigned NumberOfInputs = 0;
+
+  /// Total number of processed inputs, i.e, inputs that were already
+  /// read from the buffers.
+  unsigned NumberOfProcessedInputs = 0;
+
+  /// LLVM context used to to create the auxiliar modules.
+  LLVMContext VMContext;
+
+  /// LLVM module used to create an object with all the bundle
+  /// components.
+  std::unique_ptr<Module> AuxModule;
+
+  /// The current triple we are working with.
+  StringRef CurrentTriple;
+
+  /// The name of the main input file.
+  StringRef MainInputFileName;
+
+  /// Iterator of the current and next section.
+  section_iterator CurrentSection;
+  section_iterator NextSection;
+
+public:
+  void ReadHeader(MemoryBuffer &Input) {}
+  StringRef ReadBundleStart(MemoryBuffer &Input) {
+
+    while (NextSection != Obj.section_end()) {
+      CurrentSection = NextSection;
+      ++NextSection;
+
+      StringRef OffloadTriple;
+      // Check if the current section name starts with the reserved prefix. If
+      // so, return the triple.
+      if (IsOffloadSection(*CurrentSection, OffloadTriple))
+        return OffloadTriple;
+    }
+    return StringRef();
+  }
+  void ReadBundleEnd(MemoryBuffer &Input) {}
+  void ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) {
+    // If the current section has size one, that means that the content we are
+    // interested in is the file itself. Otherwise it is the content of the
+    // section.
+    //
+    // TODO: Instead of copying the input file as is, deactivate the section
+    // that is no longer needed.
+
+    StringRef Content;
+    CurrentSection->getContents(Content);
+
+    if (Content.size() < 2)
+      OS.write(Input.getBufferStart(), Input.getBufferSize());
+    else
+      OS.write(Content.data(), Content.size());
+  }
+
+  void WriteHeader(raw_fd_ostream &OS,
+                   ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) {
+    assert(HostInputIndex != ~0u && "Host input index not defined.");
+
+    // Record number of inputs.
+    NumberOfInputs = Inputs.size();
+
+    // Create an LLVM module to have the content we need to bundle.
+    auto *M = new Module("clang-offload-bundle", VMContext);
+    M->setTargetTriple(getTriple(TargetNames[HostInputIndex]));
+    AuxModule.reset(M);
+  }
+  void WriteBundleStart(raw_fd_ostream &OS, StringRef TargetTriple) {
+    ++NumberOfProcessedInputs;
+
+    // Record the triple we are using, that will be used to name the section we
+    // will create.
+    CurrentTriple = TargetTriple;
+  }
+  bool WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) {
+    assert(NumberOfProcessedInputs <= NumberOfInputs &&
+           "Processing more inputs that actually exist!");
+    assert(HostInputIndex != ~0u && "Host input index not defined.");
+
+    // If this is not the last output, we don't have to do anything.
+    if (NumberOfProcessedInputs != NumberOfInputs)
+      return false;
+
+    // Create the bitcode file name to write the resulting code to. Keep it if
+    // save-temps is active.
+    SmallString<128> BitcodeFileName;
+    if (sys::fs::createTemporaryFile("clang-offload-bundler", "bc",
+                                     BitcodeFileName)) {
+      llvm::errs() << "error: unable to create temporary file.\n";
+      return true;
+    }
+
+    // Dump the contents of the temporary file if that was requested.
+    if (DumpTemporaryFiles) {
+      llvm::errs() << ";\n; Object file bundler IR file.\n;\n";
+      AuxModule.get()->dump();
+    }
+
+    // Find clang in order to create the bundle binary.
+    StringRef Dir = llvm::sys::path::parent_path(BundlerExecutable);
+
+    auto ClangBinary = sys::findProgramByName("clang", Dir);
+    if (ClangBinary.getError()) {
+      // Remove bitcode file.
+      sys::fs::remove(BitcodeFileName);
+
+      llvm::errs() << "error: unable to find 'clang' in path.\n";
+      return true;
+    }
+
+    // Do the incremental linking. We write to the output file directly. So, we
+    // close it and use the name to pass down to clang.
+    OS.close();
+    SmallString<128> TargetName = getTriple(TargetNames[HostInputIndex]);
+    const char *ClangArgs[] = {"clang",
+                               "-r",
+                               "-target",
+                               TargetName.c_str(),
+                               "-o",
+                               OutputFileNames.front().c_str(),
+                               InputFileNames[HostInputIndex].c_str(),
+                               BitcodeFileName.c_str(),
+                               "-nostdlib",
+                               nullptr};
+
+    // If the user asked for the commands to be printed out, we do that instead
+    // of executing it.
+    if (PrintExternalCommands) {
+      llvm::errs() << "\"" << ClangBinary.get() << "\"";
+      for (unsigned I = 1; ClangArgs[I]; ++I)
+        llvm::errs() << " \"" << ClangArgs[I] << "\"";
+      llvm::errs() << "\n";
+    } else {
+      // Write the bitcode contents to the temporary file.
+      {
+        std::error_code EC;
+        raw_fd_ostream BitcodeFile(BitcodeFileName, EC, sys::fs::F_None);
+        if (EC) {
+          llvm::errs() << "error: unable to open temporary file.\n";
+          return true;
+        }
+        WriteBitcodeToFile(AuxModule.get(), BitcodeFile);
+      }
+
+      bool Failed = sys::ExecuteAndWait(ClangBinary.get(), ClangArgs);
+
+      // Remove bitcode file.
+      sys::fs::remove(BitcodeFileName);
+
+      if (Failed) {
+        llvm::errs() << "error: incremental linking by external tool failed.\n";
+        return true;
+      }
+    }
+
+    return false;
+  }
+  void WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) {
+    Module *M = AuxModule.get();
+
+    // Create the new section name, it will consist of the reserved prefix
+    // concatenated with the triple.
+    std::string SectionName = OFFLOAD_BUNDLER_MAGIC_STR;
+    SectionName += CurrentTriple;
+
+    // Create the constant with the content of the section. For the input we are
+    // bundling into (the host input), this is just a place-holder, so a single
+    // byte is sufficient.
+    assert(HostInputIndex != ~0u && "Host input index undefined??");
+    Constant *Content;
+    if (NumberOfProcessedInputs == HostInputIndex + 1) {
+      uint8_t Byte[] = {0};
+      Content = ConstantDataArray::get(VMContext, Byte);
+    } else
+      Content = ConstantDataArray::get(
+          VMContext, ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(
+                                           Input.getBufferStart()),
+                                       Input.getBufferSize()));
+
+    // Create the global in the desired section. We don't want these globals in
+    // the symbol table, so we mark them private.
+    auto *GV = new GlobalVariable(*M, Content->getType(), /*IsConstant=*/true,
+                                  GlobalVariable::PrivateLinkage, Content);
+    GV->setSection(SectionName);
+  }
+
+  ObjectFileHandler(ObjectFile &Obj)
+      : FileHandler(), Obj(Obj), CurrentSection(Obj.section_begin()),
+        NextSection(Obj.section_begin()) {}
+  ~ObjectFileHandler() {}
+};
+
 /// Handler for text files. The bundled file will have the following format.
 ///
 /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__START__ triple"
 /// Bundle 1
 /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__END__ triple"
 /// ...
 /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__START__ triple"
 /// Bundle N
void ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) {
     OS << Bundle;
   }
 
   void WriteHeader(raw_fd_ostream &OS,
                    ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) {}
   void WriteBundleStart(raw_fd_ostream &OS, StringRef TargetTriple) {
     OS << BundleStartString << TargetTriple << "\n";
   }
-  void WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) {
+  bool WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) {
     OS << BundleEndString << TargetTriple << "\n";
+    return false;
   }
   void WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) {
     OS << Input.getBuffer();
   }
 
 public:
   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__ ";
   }
 };
 
+/// 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 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);
+
+  // Failed to open the input as a known binary. Use the default binary handler.
+  if (!BinaryOrErr) {
+    // We don't really care about the error (we just consume it), if we could
+    // not get a valid device binary object we use the default binary handler.
+    consumeError(BinaryOrErr.takeError());
+    return new BinaryFileHandler();
+  }
+
+  // We only support regular object files. If this is not an object file,
+  // default to the binary handler. The handler will be owned by the client of
+  // this function.
+  ObjectFile *Obj = dyn_cast<ObjectFile>(BinaryOrErr.get().release());
+
+  if (!Obj)
+    return new BinaryFileHandler();
+
+  return new ObjectFileHandler(*Obj);
+}
+
 /// Return an appropriate handler given the input files and options.
 static FileHandler *CreateFileHandler(MemoryBuffer &FirstInput) {
   if (FilesType == "i")
     return new TextFileHandler(/*Comment=*/"//");
   if (FilesType == "ii")
     return new TextFileHandler(/*Comment=*/"//");
   if (FilesType == "ll")
     return new TextFileHandler(/*Comment=*/";");
   if (FilesType == "bc")
     return new BinaryFileHandler();
   if (FilesType == "s")
     return new TextFileHandler(/*Comment=*/"#");
   if (FilesType == "o")
-    return new BinaryFileHandler();
+    return CreateObjectFileHandler(FirstInput);
   if (FilesType == "gch")
     return new BinaryFileHandler();
   if (FilesType == "ast")
     return new BinaryFileHandler();
 
   llvm::errs() << "error: invalid file type specified.\n";
   return nullptr;
 }
 
 /// Bundle the files. Return true if an error was found.
 static bool BundleFiles() {
   std::error_code EC;
 
   // Create output file.
   raw_fd_ostream OutputFile(OutputFileNames.front(), EC, sys::fs::F_None);
 
   if (EC) {
     llvm::errs() << "error: Can't open file " << OutputFileNames.front()
                  << ".\n";
     return true;
   }
 
   // Open input files.
   std::vector<std::unique_ptr<MemoryBuffer>> InputBuffers(
       InputFileNames.size());
 
   unsigned Idx = 0;
   for (auto &I : InputFileNames) {
     ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
         MemoryBuffer::getFileOrSTDIN(I);
     if (std::error_code EC = CodeOrErr.getError()) {
       llvm::errs() << "error: Can't open file " << I << ": " << EC.message()
                    << "\n";
       return true;
     }
     InputBuffers[Idx++] = std::move(CodeOrErr.get());
   }
 
   // Get the file handler. We use the host buffer as reference.
   assert(HostInputIndex != ~0u && "Host input index undefined??");
   std::unique_ptr<FileHandler> FH;
   FH.reset(CreateFileHandler(*InputBuffers[HostInputIndex].get()));
 
   // Quit if we don't have a handler.
   if (!FH.get())
     return true;
 
   // Write header.
   FH.get()->WriteHeader(OutputFile, InputBuffers);
 
-  // Write all bundles along with the start/end markers.
+  // Write all bundles along with the start/end markers. If an error was found
+  // writing the end of the bundle component, abort the bundle writing.
   auto Input = InputBuffers.begin();
   for (auto &Triple : TargetNames) {
     FH.get()->WriteBundleStart(OutputFile, Triple);
     FH.get()->WriteBundle(OutputFile, *Input->get());
-    FH.get()->WriteBundleEnd(OutputFile, Triple);
+    if (FH.get()->WriteBundleEnd(OutputFile, Triple))
+      return true;
     ++Input;
   }
   return false;
 }
 
 // Unbundle the files. Return true if an error was found.
 static bool UnbundleFiles() {
   // Open Input file.
if (HostTargetNum != 1) {
     Error = true;
     llvm::errs() << "error: expecting exactly one host target but got "
                  << HostTargetNum << ".\n";
   }
 
   if (Error)
     return 1;
 
+  // Save the current executable directory as it will be useful to find other
+  // tools.
+  BundlerExecutable =
+      llvm::sys::fs::getMainExecutable(argv[0], &BundlerExecutable);
+
   return Unbundle ? UnbundleFiles() : BundleFiles();
 }