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/MemoryBufferRef.h" |
20 | #include "llvm/Support/ThreadPool.h" |
21 | #include "llvm/Target/TargetMachine.h" |
22 | #include "llvm/Transforms/Utils/SplitModule.h" |
23 | |
24 | using namespace llvm; |
25 | |
26 | static void codegen(Module *M, llvm::raw_pwrite_stream &OS, |
27 | function_ref<std::unique_ptr<TargetMachine>()> TMFactory, |
28 | CodeGenFileType FileType) { |
29 | std::unique_ptr<TargetMachine> TM = TMFactory(); |
30 | assert(TM && "Failed to create target machine!" ); |
31 | |
32 | legacy::PassManager CodeGenPasses; |
33 | if (TM->addPassesToEmitFile(CodeGenPasses, OS, nullptr, FileType)) |
34 | report_fatal_error(reason: "Failed to setup codegen" ); |
35 | CodeGenPasses.run(M&: *M); |
36 | } |
37 | |
38 | void llvm::splitCodeGen( |
39 | Module &M, ArrayRef<llvm::raw_pwrite_stream *> OSs, |
40 | ArrayRef<llvm::raw_pwrite_stream *> BCOSs, |
41 | const std::function<std::unique_ptr<TargetMachine>()> &TMFactory, |
42 | CodeGenFileType FileType, bool PreserveLocals) { |
43 | assert(BCOSs.empty() || BCOSs.size() == OSs.size()); |
44 | |
45 | if (OSs.size() == 1) { |
46 | if (!BCOSs.empty()) |
47 | WriteBitcodeToFile(M, Out&: *BCOSs[0]); |
48 | codegen(M: &M, OS&: *OSs[0], TMFactory, FileType); |
49 | return; |
50 | } |
51 | |
52 | // Create ThreadPool in nested scope so that threads will be joined |
53 | // on destruction. |
54 | { |
55 | ThreadPool CodegenThreadPool(hardware_concurrency(ThreadCount: OSs.size())); |
56 | int ThreadCount = 0; |
57 | |
58 | SplitModule( |
59 | M, N: OSs.size(), |
60 | ModuleCallback: [&](std::unique_ptr<Module> MPart) { |
61 | // We want to clone the module in a new context to multi-thread the |
62 | // codegen. We do it by serializing partition modules to bitcode |
63 | // (while still on the main thread, in order to avoid data races) and |
64 | // spinning up new threads which deserialize the partitions into |
65 | // separate contexts. |
66 | // FIXME: Provide a more direct way to do this in LLVM. |
67 | SmallString<0> BC; |
68 | raw_svector_ostream BCOS(BC); |
69 | WriteBitcodeToFile(M: *MPart, Out&: BCOS); |
70 | |
71 | if (!BCOSs.empty()) { |
72 | BCOSs[ThreadCount]->write(Ptr: BC.begin(), Size: BC.size()); |
73 | BCOSs[ThreadCount]->flush(); |
74 | } |
75 | |
76 | llvm::raw_pwrite_stream *ThreadOS = OSs[ThreadCount++]; |
77 | // Enqueue the task |
78 | CodegenThreadPool.async( |
79 | F: [TMFactory, FileType, ThreadOS](const SmallString<0> &BC) { |
80 | LLVMContext Ctx; |
81 | Expected<std::unique_ptr<Module>> MOrErr = parseBitcodeFile( |
82 | Buffer: MemoryBufferRef(StringRef(BC.data(), BC.size()), |
83 | "<split-module>" ), |
84 | Context&: Ctx); |
85 | if (!MOrErr) |
86 | report_fatal_error(reason: "Failed to read bitcode" ); |
87 | std::unique_ptr<Module> MPartInCtx = std::move(MOrErr.get()); |
88 | |
89 | codegen(M: MPartInCtx.get(), OS&: *ThreadOS, TMFactory, FileType); |
90 | }, |
91 | // Pass BC using std::move to ensure that it get moved rather than |
92 | // copied into the thread's context. |
93 | ArgList: std::move(BC)); |
94 | }, |
95 | PreserveLocals); |
96 | } |
97 | } |
98 | |