1 | //===- SplitModule.cpp - Split a module into partitions -------------------===// |
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 the function llvm::SplitModule, which splits a module |
10 | // into multiple linkable partitions. It can be used to implement parallel code |
11 | // generation for link-time optimization. |
12 | // |
13 | //===----------------------------------------------------------------------===// |
14 | |
15 | #include "llvm/Transforms/Utils/SplitModule.h" |
16 | #include "llvm/ADT/DenseMap.h" |
17 | #include "llvm/ADT/EquivalenceClasses.h" |
18 | #include "llvm/ADT/SmallPtrSet.h" |
19 | #include "llvm/ADT/SmallVector.h" |
20 | #include "llvm/ADT/StringRef.h" |
21 | #include "llvm/IR/Comdat.h" |
22 | #include "llvm/IR/Constant.h" |
23 | #include "llvm/IR/Constants.h" |
24 | #include "llvm/IR/Function.h" |
25 | #include "llvm/IR/GlobalAlias.h" |
26 | #include "llvm/IR/GlobalObject.h" |
27 | #include "llvm/IR/GlobalValue.h" |
28 | #include "llvm/IR/GlobalVariable.h" |
29 | #include "llvm/IR/Instruction.h" |
30 | #include "llvm/IR/Module.h" |
31 | #include "llvm/IR/User.h" |
32 | #include "llvm/IR/Value.h" |
33 | #include "llvm/Support/Casting.h" |
34 | #include "llvm/Support/Debug.h" |
35 | #include "llvm/Support/ErrorHandling.h" |
36 | #include "llvm/Support/MD5.h" |
37 | #include "llvm/Support/raw_ostream.h" |
38 | #include "llvm/Transforms/Utils/Cloning.h" |
39 | #include "llvm/Transforms/Utils/ValueMapper.h" |
40 | #include <algorithm> |
41 | #include <cassert> |
42 | #include <iterator> |
43 | #include <memory> |
44 | #include <queue> |
45 | #include <utility> |
46 | #include <vector> |
47 | |
48 | using namespace llvm; |
49 | |
50 | #define DEBUG_TYPE "split-module" |
51 | |
52 | namespace { |
53 | |
54 | using ClusterMapType = EquivalenceClasses<const GlobalValue *>; |
55 | using ComdatMembersType = DenseMap<const Comdat *, const GlobalValue *>; |
56 | using ClusterIDMapType = DenseMap<const GlobalValue *, unsigned>; |
57 | |
58 | } // end anonymous namespace |
59 | |
60 | static void addNonConstUser(ClusterMapType &GVtoClusterMap, |
61 | const GlobalValue *GV, const User *U) { |
62 | assert((!isa<Constant>(U) || isa<GlobalValue>(U)) && "Bad user" ); |
63 | |
64 | if (const Instruction *I = dyn_cast<Instruction>(Val: U)) { |
65 | const GlobalValue *F = I->getParent()->getParent(); |
66 | GVtoClusterMap.unionSets(V1: GV, V2: F); |
67 | } else if (const GlobalValue *GVU = dyn_cast<GlobalValue>(Val: U)) { |
68 | GVtoClusterMap.unionSets(V1: GV, V2: GVU); |
69 | } else { |
70 | llvm_unreachable("Underimplemented use case" ); |
71 | } |
72 | } |
73 | |
74 | // Adds all GlobalValue users of V to the same cluster as GV. |
75 | static void addAllGlobalValueUsers(ClusterMapType &GVtoClusterMap, |
76 | const GlobalValue *GV, const Value *V) { |
77 | for (const auto *U : V->users()) { |
78 | SmallVector<const User *, 4> Worklist; |
79 | Worklist.push_back(Elt: U); |
80 | while (!Worklist.empty()) { |
81 | const User *UU = Worklist.pop_back_val(); |
82 | // For each constant that is not a GV (a pure const) recurse. |
83 | if (isa<Constant>(Val: UU) && !isa<GlobalValue>(Val: UU)) { |
84 | Worklist.append(in_start: UU->user_begin(), in_end: UU->user_end()); |
85 | continue; |
86 | } |
87 | addNonConstUser(GVtoClusterMap, GV, U: UU); |
88 | } |
89 | } |
90 | } |
91 | |
92 | static const GlobalObject *getGVPartitioningRoot(const GlobalValue *GV) { |
93 | const GlobalObject *GO = GV->getAliaseeObject(); |
94 | if (const auto *GI = dyn_cast_or_null<GlobalIFunc>(Val: GO)) |
95 | GO = GI->getResolverFunction(); |
96 | return GO; |
97 | } |
98 | |
99 | // Find partitions for module in the way that no locals need to be |
100 | // globalized. |
101 | // Try to balance pack those partitions into N files since this roughly equals |
102 | // thread balancing for the backend codegen step. |
103 | static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap, |
104 | unsigned N) { |
105 | // At this point module should have the proper mix of globals and locals. |
106 | // As we attempt to partition this module, we must not change any |
107 | // locals to globals. |
108 | LLVM_DEBUG(dbgs() << "Partition module with (" << M.size() << ")functions\n" ); |
109 | ClusterMapType GVtoClusterMap; |
110 | ComdatMembersType ComdatMembers; |
111 | |
112 | auto recordGVSet = [&GVtoClusterMap, &ComdatMembers](GlobalValue &GV) { |
113 | if (GV.isDeclaration()) |
114 | return; |
115 | |
116 | if (!GV.hasName()) |
117 | GV.setName("__llvmsplit_unnamed" ); |
118 | |
119 | // Comdat groups must not be partitioned. For comdat groups that contain |
120 | // locals, record all their members here so we can keep them together. |
121 | // Comdat groups that only contain external globals are already handled by |
122 | // the MD5-based partitioning. |
123 | if (const Comdat *C = GV.getComdat()) { |
124 | auto &Member = ComdatMembers[C]; |
125 | if (Member) |
126 | GVtoClusterMap.unionSets(V1: Member, V2: &GV); |
127 | else |
128 | Member = &GV; |
129 | } |
130 | |
131 | // Aliases should not be separated from their aliasees and ifuncs should |
132 | // not be separated from their resolvers regardless of linkage. |
133 | if (const GlobalObject *Root = getGVPartitioningRoot(GV: &GV)) |
134 | if (&GV != Root) |
135 | GVtoClusterMap.unionSets(V1: &GV, V2: Root); |
136 | |
137 | if (const Function *F = dyn_cast<Function>(Val: &GV)) { |
138 | for (const BasicBlock &BB : *F) { |
139 | BlockAddress *BA = BlockAddress::lookup(BB: &BB); |
140 | if (!BA || !BA->isConstantUsed()) |
141 | continue; |
142 | addAllGlobalValueUsers(GVtoClusterMap, GV: F, V: BA); |
143 | } |
144 | } |
145 | |
146 | if (GV.hasLocalLinkage()) |
147 | addAllGlobalValueUsers(GVtoClusterMap, GV: &GV, V: &GV); |
148 | }; |
149 | |
150 | llvm::for_each(Range: M.functions(), F: recordGVSet); |
151 | llvm::for_each(Range: M.globals(), F: recordGVSet); |
152 | llvm::for_each(Range: M.aliases(), F: recordGVSet); |
153 | |
154 | // Assigned all GVs to merged clusters while balancing number of objects in |
155 | // each. |
156 | auto CompareClusters = [](const std::pair<unsigned, unsigned> &a, |
157 | const std::pair<unsigned, unsigned> &b) { |
158 | if (a.second || b.second) |
159 | return a.second > b.second; |
160 | else |
161 | return a.first > b.first; |
162 | }; |
163 | |
164 | std::priority_queue<std::pair<unsigned, unsigned>, |
165 | std::vector<std::pair<unsigned, unsigned>>, |
166 | decltype(CompareClusters)> |
167 | BalancinQueue(CompareClusters); |
168 | // Pre-populate priority queue with N slot blanks. |
169 | for (unsigned i = 0; i < N; ++i) |
170 | BalancinQueue.push(x: std::make_pair(x&: i, y: 0)); |
171 | |
172 | using SortType = std::pair<unsigned, ClusterMapType::iterator>; |
173 | |
174 | SmallVector<SortType, 64> Sets; |
175 | SmallPtrSet<const GlobalValue *, 32> Visited; |
176 | |
177 | // To guarantee determinism, we have to sort SCC according to size. |
178 | // When size is the same, use leader's name. |
179 | for (ClusterMapType::iterator I = GVtoClusterMap.begin(), |
180 | E = GVtoClusterMap.end(); I != E; ++I) |
181 | if (I->isLeader()) |
182 | Sets.push_back( |
183 | Elt: std::make_pair(x: std::distance(first: GVtoClusterMap.member_begin(I), |
184 | last: GVtoClusterMap.member_end()), y&: I)); |
185 | |
186 | llvm::sort(C&: Sets, Comp: [](const SortType &a, const SortType &b) { |
187 | if (a.first == b.first) |
188 | return a.second->getData()->getName() > b.second->getData()->getName(); |
189 | else |
190 | return a.first > b.first; |
191 | }); |
192 | |
193 | for (auto &I : Sets) { |
194 | unsigned CurrentClusterID = BalancinQueue.top().first; |
195 | unsigned CurrentClusterSize = BalancinQueue.top().second; |
196 | BalancinQueue.pop(); |
197 | |
198 | LLVM_DEBUG(dbgs() << "Root[" << CurrentClusterID << "] cluster_size(" |
199 | << I.first << ") ----> " << I.second->getData()->getName() |
200 | << "\n" ); |
201 | |
202 | for (ClusterMapType::member_iterator MI = |
203 | GVtoClusterMap.findLeader(I: I.second); |
204 | MI != GVtoClusterMap.member_end(); ++MI) { |
205 | if (!Visited.insert(Ptr: *MI).second) |
206 | continue; |
207 | LLVM_DEBUG(dbgs() << "----> " << (*MI)->getName() |
208 | << ((*MI)->hasLocalLinkage() ? " l " : " e " ) << "\n" ); |
209 | Visited.insert(Ptr: *MI); |
210 | ClusterIDMap[*MI] = CurrentClusterID; |
211 | CurrentClusterSize++; |
212 | } |
213 | // Add this set size to the number of entries in this cluster. |
214 | BalancinQueue.push(x: std::make_pair(x&: CurrentClusterID, y&: CurrentClusterSize)); |
215 | } |
216 | } |
217 | |
218 | static void externalize(GlobalValue *GV) { |
219 | if (GV->hasLocalLinkage()) { |
220 | GV->setLinkage(GlobalValue::ExternalLinkage); |
221 | GV->setVisibility(GlobalValue::HiddenVisibility); |
222 | } |
223 | |
224 | // Unnamed entities must be named consistently between modules. setName will |
225 | // give a distinct name to each such entity. |
226 | if (!GV->hasName()) |
227 | GV->setName("__llvmsplit_unnamed" ); |
228 | } |
229 | |
230 | // Returns whether GV should be in partition (0-based) I of N. |
231 | static bool isInPartition(const GlobalValue *GV, unsigned I, unsigned N) { |
232 | if (const GlobalObject *Root = getGVPartitioningRoot(GV)) |
233 | GV = Root; |
234 | |
235 | StringRef Name; |
236 | if (const Comdat *C = GV->getComdat()) |
237 | Name = C->getName(); |
238 | else |
239 | Name = GV->getName(); |
240 | |
241 | // Partition by MD5 hash. We only need a few bits for evenness as the number |
242 | // of partitions will generally be in the 1-2 figure range; the low 16 bits |
243 | // are enough. |
244 | MD5 H; |
245 | MD5::MD5Result R; |
246 | H.update(Str: Name); |
247 | H.final(Result&: R); |
248 | return (R[0] | (R[1] << 8)) % N == I; |
249 | } |
250 | |
251 | void llvm::SplitModule( |
252 | Module &M, unsigned N, |
253 | function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback, |
254 | bool PreserveLocals) { |
255 | if (!PreserveLocals) { |
256 | for (Function &F : M) |
257 | externalize(GV: &F); |
258 | for (GlobalVariable &GV : M.globals()) |
259 | externalize(GV: &GV); |
260 | for (GlobalAlias &GA : M.aliases()) |
261 | externalize(GV: &GA); |
262 | for (GlobalIFunc &GIF : M.ifuncs()) |
263 | externalize(GV: &GIF); |
264 | } |
265 | |
266 | // This performs splitting without a need for externalization, which might not |
267 | // always be possible. |
268 | ClusterIDMapType ClusterIDMap; |
269 | findPartitions(M, ClusterIDMap, N); |
270 | |
271 | // FIXME: We should be able to reuse M as the last partition instead of |
272 | // cloning it. Note that the callers at the moment expect the module to |
273 | // be preserved, so will need some adjustments as well. |
274 | for (unsigned I = 0; I < N; ++I) { |
275 | ValueToValueMapTy VMap; |
276 | std::unique_ptr<Module> MPart( |
277 | CloneModule(M, VMap, ShouldCloneDefinition: [&](const GlobalValue *GV) { |
278 | if (ClusterIDMap.count(Val: GV)) |
279 | return (ClusterIDMap[GV] == I); |
280 | else |
281 | return isInPartition(GV, I, N); |
282 | })); |
283 | if (I != 0) |
284 | MPart->setModuleInlineAsm("" ); |
285 | ModuleCallback(std::move(MPart)); |
286 | } |
287 | } |
288 | |