[gold] Save bitcode for module partitions (save-temps + split codegen).

[eugenis]

Summary

When doing parallel codegen (-plugin-opt,save-temps, jobs=2), save bitcode for module partitions.
Currently we save object files for module partitions, but not bitcode files.

[• rafael]

LGTM

[eugenis]

Thanks, r265583

include/llvm/CodeGen/ParallelCG.h

 class TargetOptions;
 class raw_pwrite_stream;
 
 /// Split M into OSs.size() partitions, and generate code for each. Writes
 /// OSs.size() output files to the output streams in OSs. The resulting output
 /// files if linked together are intended to be equivalent to the single output
 /// file that would have been code generated from M.
 ///
+/// Writes bitcode for individual partitions into output streams in BCOSs, if
+/// BCOSs is not empty.
+///
 /// \returns M if OSs.size() == 1, otherwise returns std::unique_ptr<Module>().
 std::unique_ptr<Module>
 splitCodeGen(std::unique_ptr<Module> M, ArrayRef<raw_pwrite_stream *> OSs,
-             StringRef CPU, StringRef Features, const TargetOptions &Options,
+             ArrayRef<llvm::raw_pwrite_stream *> BCOSs, StringRef CPU,
+             StringRef Features, const TargetOptions &Options,
              Reloc::Model RM = Reloc::Default,
              CodeModel::Model CM = CodeModel::Default,
              CodeGenOpt::Level OL = CodeGenOpt::Default,
              TargetMachine::CodeGenFileType FT = TargetMachine::CGFT_ObjectFile,
              bool PreserveLocals = false);
 
 } // namespace llvm
 
 #endif

lib/CodeGen/ParallelCG.cpp

std::unique_ptr<TargetMachine> TM(TheTarget->createTargetMachine(
       M->getTargetTriple(), CPU, Features, Options, RM, CM, OL));
 
   legacy::PassManager CodeGenPasses;
   if (TM->addPassesToEmitFile(CodeGenPasses, OS, FileType))
     report_fatal_error("Failed to setup codegen");
   CodeGenPasses.run(*M);
 }
 
-std::unique_ptr<Module>
-llvm::splitCodeGen(std::unique_ptr<Module> M,
-                   ArrayRef<llvm::raw_pwrite_stream *> OSs, StringRef CPU,
-                   StringRef Features, const TargetOptions &Options,
-                   Reloc::Model RM, CodeModel::Model CM, CodeGenOpt::Level OL,
-                   TargetMachine::CodeGenFileType FileType,
-                   bool PreserveLocals) {
+std::unique_ptr<Module> llvm::splitCodeGen(
+    std::unique_ptr<Module> M, ArrayRef<llvm::raw_pwrite_stream *> OSs,
+    ArrayRef<llvm::raw_pwrite_stream *> BCOSs, StringRef CPU,
+    StringRef Features, const TargetOptions &Options, Reloc::Model RM,
+    CodeModel::Model CM, CodeGenOpt::Level OL,
+    TargetMachine::CodeGenFileType FileType, bool PreserveLocals) {
   StringRef TripleStr = M->getTargetTriple();
   std::string ErrMsg;
   const Target *TheTarget = TargetRegistry::lookupTarget(TripleStr, ErrMsg);
   if (!TheTarget)
     report_fatal_error(Twine("Target not found: ") + ErrMsg);
 
+  assert(BCOSs.empty() || BCOSs.size() == OSs.size());
+
   if (OSs.size() == 1) {
-    codegen(M.get(), *OSs[0], TheTarget, CPU, Features, Options, RM, CM,
-            OL, FileType);
+    if (!BCOSs.empty())
+      WriteBitcodeToFile(M.get(), *BCOSs[0]);
+    codegen(M.get(), *OSs[0], TheTarget, CPU, Features, Options, RM, CM, OL,
+            FileType);
     return M;
   }
 
   // Create ThreadPool in nested scope so that threads will be joined
   // on destruction.
   {
     ThreadPool CodegenThreadPool(OSs.size());
     int ThreadCount = 0;
 
     SplitModule(
         std::move(M), OSs.size(),
         [&](std::unique_ptr<Module> MPart) {
           // We want to clone the module in a new context to multi-thread the
           // codegen. We do it by serializing partition modules to bitcode
           // (while still on the main thread, in order to avoid data races) and
           // spinning up new threads which deserialize the partitions into
           // separate contexts.
           // FIXME: Provide a more direct way to do this in LLVM.
           SmallVector<char, 0> BC;
           raw_svector_ostream BCOS(BC);
           WriteBitcodeToFile(MPart.get(), BCOS);
 
+          if (!BCOSs.empty()) {
+            BCOSs[ThreadCount]->write(BC.begin(), BC.size());
+            BCOSs[ThreadCount]->flush();
+          }
+
           llvm::raw_pwrite_stream *ThreadOS = OSs[ThreadCount++];
           // Enqueue the task
           CodegenThreadPool.async(
               [TheTarget, CPU, Features, Options, RM, CM, OL, FileType,
                ThreadOS](const SmallVector<char, 0> &BC) {
                 LLVMContext Ctx;
                 ErrorOr<std::unique_ptr<Module>> MOrErr = parseBitcodeFile(
                     MemoryBufferRef(StringRef(BC.data(), BC.size()),

lib/LTO/LTOCodeGenerator.cpp

bool LTOCodeGenerator::compileOptimized(ArrayRef<raw_pwrite_stream *> Out) {
   // for splitting
   restoreLinkageForExternals();
 
   // Do code generation. We need to preserve the module in case the client calls
   // writeMergedModules() after compilation, but we only need to allow this at
   // parallelism level 1. This is achieved by having splitCodeGen return the
   // original module at parallelism level 1 which we then assign back to
   // MergedModule.
-  MergedModule =
-      splitCodeGen(std::move(MergedModule), Out, MCpu, FeatureStr, Options,
-                   RelocModel, CodeModel::Default, CGOptLevel, FileType,
-                   ShouldRestoreGlobalsLinkage);
+  MergedModule = splitCodeGen(
+      std::move(MergedModule), Out, {}, MCpu, FeatureStr, Options, RelocModel,
+      CodeModel::Default, CGOptLevel, FileType, ShouldRestoreGlobalsLinkage);
 
   // If statistics were requested, print them out after codegen.
   if (llvm::AreStatisticsEnabled())
     llvm::PrintStatistics();
 
   return true;
 }
 

test/tools/gold/X86/parallel.ll

 ; RUN: llvm-as -o %t.bc %s
 ; RUN: env LD_PRELOAD=%llvmshlibdir/LLVMgold.so %gold -plugin %llvmshlibdir/LLVMgold.so -u foo -u bar -plugin-opt jobs=2 -plugin-opt save-temps -m elf_x86_64 -o %t %t.bc
+; RUN: llvm-dis %t.opt.bc0 -o - | FileCheck --check-prefix=CHECK-BC0 %s
+; RUN: llvm-dis %t.opt.bc1 -o - | FileCheck --check-prefix=CHECK-BC1 %s
 ; RUN: llvm-nm %t.o0 | FileCheck --check-prefix=CHECK0 %s
 ; RUN: llvm-nm %t.o1 | FileCheck --check-prefix=CHECK1 %s
 
 target triple = "x86_64-unknown-linux-gnu"
 
+; CHECK-BC0: define void @foo
+; CHECK-BC0: declare void @bar
 ; CHECK0-NOT: bar
 ; CHECK0: T foo
 ; CHECK0-NOT: bar
 define void @foo() {
   call void @bar()
   ret void
 }
 
+; CHECK-BC1: declare void @foo
+; CHECK-BC1: define void @bar
 ; CHECK1-NOT: foo
 ; CHECK1: T bar
 ; CHECK1-NOT: foo
 define void @bar() {
   call void @foo()
   ret void
 }

tools/gold/gold-plugin.cpp

private:
   /// module/task.
   void initTargetMachine();
 
   /// Run all LTO passes on the module.
   void runLTOPasses();
 
   /// Sets up output files necessary to perform optional multi-threaded
   /// split code generation, and invokes the code generation implementation.
-  void runSplitCodeGen();
+  /// If BCFileName is not empty, saves bitcode for module partitions into
+  /// {BCFileName}0 .. {BCFileName}N.
+  void runSplitCodeGen(const SmallString<128> &BCFilename);
 };
 }
 
 static SubtargetFeatures getFeatures(Triple &TheTriple) {
   SubtargetFeatures Features;
   Features.getDefaultSubtargetFeatures(TheTriple);
   for (const std::string &A : MAttrs)
     Features.AddFeature(A);
void CodeGen::runCodegenPasses() {
   assert(OS && "Output stream must be set before emitting to file");
   legacy::PassManager CodeGenPasses;
   if (TM->addPassesToEmitFile(CodeGenPasses, *OS,
                               TargetMachine::CGFT_ObjectFile))
     report_fatal_error("Failed to setup codegen");
   CodeGenPasses.run(*M);
 }
 
-void CodeGen::runSplitCodeGen() {
+void CodeGen::runSplitCodeGen(const SmallString<128> &BCFilename) {
   const std::string &TripleStr = M->getTargetTriple();
   Triple TheTriple(TripleStr);
 
   SubtargetFeatures Features = getFeatures(TheTriple);
 
   TargetOptions Options = InitTargetOptionsFromCodeGenFlags();
   CodeGenOpt::Level CGOptLevel = getCGOptLevel();
 
   SmallString<128> Filename;
   // Note that openOutputFile will append a unique ID for each task
   if (!options::obj_path.empty())
     Filename = options::obj_path;
   else if (options::TheOutputType == options::OT_SAVE_TEMPS)
     Filename = output_name + ".o";
 
   // Note that the default parallelism is 1 instead of the
   // hardware_concurrency, as there are behavioral differences between
   // parallelism levels (e.g. symbol ordering will be different, and some uses
   // of inline asm currently have issues with parallelism >1).
   unsigned int MaxThreads = options::Parallelism ? options::Parallelism : 1;
 
   std::vector<SmallString<128>> Filenames(MaxThreads);
+  std::vector<SmallString<128>> BCFilenames(MaxThreads);
   bool TempOutFile = Filename.empty();
   {
     // Open a file descriptor for each backend task. This is done in a block
     // so that the output file descriptors are closed before gold opens them.
     std::list<llvm::raw_fd_ostream> OSs;
     std::vector<llvm::raw_pwrite_stream *> OSPtrs(MaxThreads);
     for (unsigned I = 0; I != MaxThreads; ++I) {
       int FD = openOutputFile(Filename, TempOutFile, Filenames[I],
                               // Only append ID if there are multiple tasks.
                               MaxThreads > 1 ? I : -1);
       OSs.emplace_back(FD, true);
       OSPtrs[I] = &OSs.back();
     }
 
+    std::list<llvm::raw_fd_ostream> BCOSs;
+    std::vector<llvm::raw_pwrite_stream *> BCOSPtrs;
+    if (!BCFilename.empty() && MaxThreads > 1) {
+      for (unsigned I = 0; I != MaxThreads; ++I) {
+        int FD = openOutputFile(BCFilename, false, BCFilenames[I], I);
+        BCOSs.emplace_back(FD, true);
+        BCOSPtrs.push_back(&BCOSs.back());
+      }
+    }
+
     // Run backend tasks.
-    splitCodeGen(std::move(M), OSPtrs, options::mcpu, Features.getString(),
+    splitCodeGen(std::move(M), OSPtrs, BCOSPtrs, options::mcpu, Features.getString(),
                  Options, RelocationModel, CodeModel::Default, CGOptLevel);
   }
 
   for (auto &Filename : Filenames)
     recordFile(Filename.c_str(), TempOutFile);
 }
 
 void CodeGen::runAll() {
   runLTOPasses();
 
+  SmallString<128> OptFilename;
   if (options::TheOutputType == options::OT_SAVE_TEMPS) {
-    std::string OptFilename = output_name;
+    OptFilename = output_name;
     // If the CodeGen client provided a filename, use it. Always expect
     // a provided filename if we are in a task (i.e. ThinLTO backend).
     assert(!SaveTempsFilename.empty() || TaskID == -1);
     if (!SaveTempsFilename.empty())
       OptFilename = SaveTempsFilename;
-    saveBCFile(OptFilename + ".opt.bc", *M);
+    OptFilename += ".opt.bc";
+    saveBCFile(OptFilename, *M);
   }
 
   // If we are already in a thread (i.e. ThinLTO), just perform
   // codegen passes directly.
   if (TaskID >= 0)
     runCodegenPasses();
   // Otherwise attempt split code gen.
   else
-    runSplitCodeGen();
+    runSplitCodeGen(OptFilename);
 }
 
 /// Links the module in \p View from file \p F into the combined module
 /// saved in the IRMover \p L. Returns true on error, false on success.
 static bool linkInModule(LLVMContext &Context, IRMover &L, claimed_file &F,
                          const void *View, ld_plugin_input_file &File,
                          raw_fd_ostream *ApiFile, StringSet<> &Internalize,
                          StringSet<> &Maybe) {