1 files changed, 54 insertions, 51 deletions
diff --git a/contrib/llvm/lib/CodeGen/ParallelCG.cpp b/contrib/llvm/lib/CodeGen/ParallelCG.cpp
index e73ba02..ccdaec1 100644
--- a/contrib/llvm/lib/CodeGen/ParallelCG.cpp
+++ b/contrib/llvm/lib/CodeGen/ParallelCG.cpp
@@ -19,78 +19,81 @@
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/TargetRegistry.h"
-#include "llvm/Support/thread.h"
+#include "llvm/Support/ThreadPool.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
 
 using namespace llvm;
 
 static void codegen(Module *M, llvm::raw_pwrite_stream &OS,
-                    const Target *TheTarget, StringRef CPU, StringRef Features,
-                    const TargetOptions &Options, Reloc::Model RM,
-                    CodeModel::Model CM, CodeGenOpt::Level OL,
+                    function_ref<std::unique_ptr<TargetMachine>()> TMFactory,
                     TargetMachine::CodeGenFileType FileType) {
-  std::unique_ptr<TargetMachine> TM(TheTarget->createTargetMachine(
-      M->getTargetTriple(), CPU, Features, Options, RM, CM, OL));
-
+  std::unique_ptr<TargetMachine> TM = TMFactory();
   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) {
-  StringRef TripleStr = M->getTargetTriple();
-  std::string ErrMsg;
-  const Target *TheTarget = TargetRegistry::lookupTarget(TripleStr, ErrMsg);
-  if (!TheTarget)
-    report_fatal_error(Twine("Target not found: ") + ErrMsg);
+std::unique_ptr<Module> llvm::splitCodeGen(
+    std::unique_ptr<Module> M, ArrayRef<llvm::raw_pwrite_stream *> OSs,
+    ArrayRef<llvm::raw_pwrite_stream *> BCOSs,
+    const std::function<std::unique_ptr<TargetMachine>()> &TMFactory,
+    TargetMachine::CodeGenFileType FileType, bool PreserveLocals) {
+  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], TMFactory, FileType);
     return M;
   }
 
-  std::vector<thread> Threads;
-  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);
+  // Create ThreadPool in nested scope so that threads will be joined
+  // on destruction.
+  {
+    ThreadPool CodegenThreadPool(OSs.size());
+    int ThreadCount = 0;
 
-    llvm::raw_pwrite_stream *ThreadOS = OSs[Threads.size()];
-    Threads.emplace_back(
-        [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()),
-                                               "<split-module>"),
-                               Ctx);
-          if (!MOrErr)
-            report_fatal_error("Failed to read bitcode");
-          std::unique_ptr<Module> MPartInCtx = std::move(MOrErr.get());
+    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.
+          SmallString<0> BC;
+          raw_svector_ostream BCOS(BC);
+          WriteBitcodeToFile(MPart.get(), BCOS);
 
-          codegen(MPartInCtx.get(), *ThreadOS, TheTarget, CPU, Features,
-                  Options, RM, CM, OL, FileType);
-        },
-        // Pass BC using std::move to ensure that it get moved rather than
-        // copied into the thread's context.
-        std::move(BC));
-  });
+          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(
+              [TMFactory, FileType, ThreadOS](const SmallString<0> &BC) {
+                LLVMContext Ctx;
+                ErrorOr<std::unique_ptr<Module>> MOrErr = parseBitcodeFile(
+                    MemoryBufferRef(StringRef(BC.data(), BC.size()),
+                                    "<split-module>"),
+                    Ctx);
+                if (!MOrErr)
+                  report_fatal_error("Failed to read bitcode");
+                std::unique_ptr<Module> MPartInCtx = std::move(MOrErr.get());
 
-  for (thread &T : Threads)
-    T.join();
+                codegen(MPartInCtx.get(), *ThreadOS, TMFactory, FileType);
+              },
+              // Pass BC using std::move to ensure that it get moved rather than
+              // copied into the thread's context.
+              std::move(BC));
+        },
+        PreserveLocals);
+  }
 
   return {};
 }