lib/CodeGen/ParallelCG.cpp

   1 //===-- ParallelCG.cpp ----------------------------------------------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This file defines functions that can be used for parallel code generation.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/CodeGen/ParallelCG.h"
  15 #include "llvm/Bitcode/BitcodeReader.h"
  16 #include "llvm/Bitcode/BitcodeWriter.h"
  17 #include "llvm/IR/LLVMContext.h"
  18 #include "llvm/IR/LegacyPassManager.h"
  19 #include "llvm/IR/Module.h"
  20 #include "llvm/Support/ErrorOr.h"
  21 #include "llvm/Support/MemoryBuffer.h"
  22 #include "llvm/Support/TargetRegistry.h"
  23 #include "llvm/Support/ThreadPool.h"
  24 #include "llvm/Target/TargetMachine.h"
  25 #include "llvm/Transforms/Utils/SplitModule.h"
  26
  27 using namespace llvm;
  28
  29 static void codegen(Module *M, llvm::raw_pwrite_stream &OS,
  30                     function_ref<std::unique_ptr<TargetMachine>()> TMFactory,
  31                     TargetMachine::CodeGenFileType FileType) {
  32   std::unique_ptr<TargetMachine> TM = TMFactory();
  33   legacy::PassManager CodeGenPasses;
  34   if (TM->addPassesToEmitFile(CodeGenPasses, OS, FileType))
  35     report_fatal_error("Failed to setup codegen");
  36   CodeGenPasses.run(*M);
  37 }
  38
  39 std::unique_ptr<Module> llvm::splitCodeGen(
  40     std::unique_ptr<Module> M, ArrayRef<llvm::raw_pwrite_stream *> OSs,
  41     ArrayRef<llvm::raw_pwrite_stream *> BCOSs,
  42     const std::function<std::unique_ptr<TargetMachine>()> &TMFactory,
  43     TargetMachine::CodeGenFileType FileType, bool PreserveLocals) {
  44   assert(BCOSs.empty() || BCOSs.size() == OSs.size());
  45
  46   if (OSs.size() == 1) {
  47     if (!BCOSs.empty())
  48       WriteBitcodeToFile(M.get(), *BCOSs[0]);
  49     codegen(M.get(), *OSs[0], TMFactory, FileType);
  50     return M;
  51   }
  52
  53   // Create ThreadPool in nested scope so that threads will be joined
  54   // on destruction.
  55   {
  56     ThreadPool CodegenThreadPool(OSs.size());
  57     int ThreadCount = 0;
  58
  59     SplitModule(
  60         std::move(M), OSs.size(),
  61         [&](std::unique_ptr<Module> MPart) {
  62           // We want to clone the module in a new context to multi-thread the
  63           // codegen. We do it by serializing partition modules to bitcode
  64           // (while still on the main thread, in order to avoid data races) and
  65           // spinning up new threads which deserialize the partitions into
  66           // separate contexts.
  67           // FIXME: Provide a more direct way to do this in LLVM.
  68           SmallString<0> BC;
  69           raw_svector_ostream BCOS(BC);
  70           WriteBitcodeToFile(MPart.get(), BCOS);
  71
  72           if (!BCOSs.empty()) {
  73             BCOSs[ThreadCount]->write(BC.begin(), BC.size());
  74             BCOSs[ThreadCount]->flush();
  75           }
  76
  77           llvm::raw_pwrite_stream *ThreadOS = OSs[ThreadCount++];
  78           // Enqueue the task
  79           CodegenThreadPool.async(
  80               [TMFactory, FileType, ThreadOS](const SmallString<0> &BC) {
  81                 LLVMContext Ctx;
  82                 Expected<std::unique_ptr<Module>> MOrErr = parseBitcodeFile(
  83                     MemoryBufferRef(StringRef(BC.data(), BC.size()),
  84                                     "<split-module>"),
  85                     Ctx);
  86                 if (!MOrErr)
  87                   report_fatal_error("Failed to read bitcode");
  88                 std::unique_ptr<Module> MPartInCtx = std::move(MOrErr.get());
  89
  90                 codegen(MPartInCtx.get(), *ThreadOS, TMFactory, FileType);
  91               },
  92               // Pass BC using std::move to ensure that it get moved rather than
  93               // copied into the thread's context.
  94               std::move(BC));
  95         },
  96         PreserveLocals);
  97   }
  98
  99   return {};
 100 }