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                     CodeGenFileType FileType) {
30   std::unique_ptr<TargetMachine> TM = TMFactory();
31   assert(TM && "Failed to create target machine!");
32 
33   legacy::PassManager CodeGenPasses;
34   if (TM->addPassesToEmitFile(CodeGenPasses, OS, nullptr, FileType))
35     report_fatal_error("Failed to setup codegen");
36   CodeGenPasses.run(*M);
37 }
38 
39 void llvm::splitCodeGen(
40     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     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, *BCOSs[0]);
49     codegen(&M, *OSs[0], TMFactory, FileType);
50     return;
51   }
52 
53   // Create ThreadPool in nested scope so that threads will be joined
54   // on destruction.
55   {
56     ThreadPool CodegenThreadPool(hardware_concurrency(OSs.size()));
57     int ThreadCount = 0;
58 
59     SplitModule(
60         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, 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