lib/CodeGen/ParallelCG.cpp

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