1//===----- IRPartitionLayer.cpp - Partition IR module into submodules -----===//
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#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
10#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
11#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
12
13using namespace llvm;
14using namespace llvm::orc;
15
16static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM,
17 StringRef Suffix,
18 GVPredicate ShouldExtract) {
19
20 auto DeleteExtractedDefs = [](GlobalValue &GV) {
21 // Bump the linkage: this global will be provided by the external module.
22 GV.setLinkage(GlobalValue::ExternalLinkage);
23
24 // Delete the definition in the source module.
25 if (isa<Function>(Val: GV)) {
26 auto &F = cast<Function>(Val&: GV);
27 F.deleteBody();
28 F.setPersonalityFn(nullptr);
29 } else if (isa<GlobalVariable>(Val: GV)) {
30 cast<GlobalVariable>(Val&: GV).setInitializer(nullptr);
31 } else if (isa<GlobalAlias>(Val: GV)) {
32 // We need to turn deleted aliases into function or variable decls based
33 // on the type of their aliasee.
34 auto &A = cast<GlobalAlias>(Val&: GV);
35 Constant *Aliasee = A.getAliasee();
36 assert(A.hasName() && "Anonymous alias?");
37 assert(Aliasee->hasName() && "Anonymous aliasee");
38 std::string AliasName = std::string(A.getName());
39
40 if (isa<Function>(Val: Aliasee)) {
41 auto *F = cloneFunctionDecl(Dst&: *A.getParent(), F: *cast<Function>(Val: Aliasee));
42 A.replaceAllUsesWith(V: F);
43 A.eraseFromParent();
44 F->setName(AliasName);
45 } else if (isa<GlobalVariable>(Val: Aliasee)) {
46 auto *G = cloneGlobalVariableDecl(Dst&: *A.getParent(),
47 GV: *cast<GlobalVariable>(Val: Aliasee));
48 A.replaceAllUsesWith(V: G);
49 A.eraseFromParent();
50 G->setName(AliasName);
51 } else
52 llvm_unreachable("Alias to unsupported type");
53 } else
54 llvm_unreachable("Unsupported global type");
55 };
56
57 auto NewTSM = cloneToNewContext(TSMW: TSM, ShouldCloneDef: ShouldExtract, UpdateClonedDefSource: DeleteExtractedDefs);
58 NewTSM.withModuleDo(F: [&](Module &M) {
59 M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
60 });
61
62 return NewTSM;
63}
64
65namespace llvm {
66namespace orc {
67
68class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
69public:
70 PartitioningIRMaterializationUnit(ExecutionSession &ES,
71 const IRSymbolMapper::ManglingOptions &MO,
72 ThreadSafeModule TSM,
73 IRPartitionLayer &Parent)
74 : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {}
75
76 PartitioningIRMaterializationUnit(
77 ThreadSafeModule TSM, Interface I,
78 SymbolNameToDefinitionMap SymbolToDefinition, IRPartitionLayer &Parent)
79 : IRMaterializationUnit(std::move(TSM), std::move(I),
80 std::move(SymbolToDefinition)),
81 Parent(Parent) {}
82
83private:
84 void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
85 Parent.emitPartition(R: std::move(R), TSM: std::move(TSM),
86 Defs: std::move(SymbolToDefinition));
87 }
88
89 void discard(const JITDylib &V, const SymbolStringPtr &Name) override {
90 // All original symbols were materialized by the CODLayer and should be
91 // final. The function bodies provided by M should never be overridden.
92 llvm_unreachable("Discard should never be called on an "
93 "ExtractingIRMaterializationUnit");
94 }
95
96 IRPartitionLayer &Parent;
97};
98
99} // namespace orc
100} // namespace llvm
101
102IRPartitionLayer::IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer)
103 : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer) {}
104
105void IRPartitionLayer::setPartitionFunction(PartitionFunction Partition) {
106 this->Partition = Partition;
107}
108
109std::optional<IRPartitionLayer::GlobalValueSet>
110IRPartitionLayer::compileRequested(GlobalValueSet Requested) {
111 return std::move(Requested);
112}
113
114std::optional<IRPartitionLayer::GlobalValueSet>
115IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) {
116 return std::nullopt;
117}
118
119void IRPartitionLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
120 ThreadSafeModule TSM) {
121 assert(TSM && "Null module");
122
123 auto &ES = getExecutionSession();
124 TSM.withModuleDo(F: [&](Module &M) {
125 // First, do some cleanup on the module:
126 cleanUpModule(M);
127 });
128
129 // Create a partitioning materialization unit and pass the responsibility.
130 if (auto Err = R->replace(MU: std::make_unique<PartitioningIRMaterializationUnit>(
131 args&: ES, args: *getManglingOptions(), args: std::move(TSM), args&: *this))) {
132 ES.reportError(Err: std::move(Err));
133 R->failMaterialization();
134 return;
135 }
136}
137
138void IRPartitionLayer::cleanUpModule(Module &M) {
139 for (auto &F : M.functions()) {
140 if (F.isDeclaration())
141 continue;
142
143 if (F.hasAvailableExternallyLinkage()) {
144 F.deleteBody();
145 F.setPersonalityFn(nullptr);
146 continue;
147 }
148 }
149}
150
151void IRPartitionLayer::expandPartition(GlobalValueSet &Partition) {
152 // Expands the partition to ensure the following rules hold:
153 // (1) If any alias is in the partition, its aliasee is also in the partition.
154 // (2) If any aliasee is in the partition, its aliases are also in the
155 // partiton.
156 // (3) If any global variable is in the partition then all global variables
157 // are in the partition.
158 assert(!Partition.empty() && "Unexpected empty partition");
159
160 const Module &M = *(*Partition.begin())->getParent();
161 bool ContainsGlobalVariables = false;
162 std::vector<const GlobalValue *> GVsToAdd;
163
164 for (const auto *GV : Partition)
165 if (isa<GlobalAlias>(Val: GV))
166 GVsToAdd.push_back(
167 x: cast<GlobalValue>(Val: cast<GlobalAlias>(Val: GV)->getAliasee()));
168 else if (isa<GlobalVariable>(Val: GV))
169 ContainsGlobalVariables = true;
170
171 for (auto &A : M.aliases())
172 if (Partition.count(x: cast<GlobalValue>(Val: A.getAliasee())))
173 GVsToAdd.push_back(x: &A);
174
175 if (ContainsGlobalVariables)
176 for (auto &G : M.globals())
177 GVsToAdd.push_back(x: &G);
178
179 for (const auto *GV : GVsToAdd)
180 Partition.insert(x: GV);
181}
182
183void IRPartitionLayer::emitPartition(
184 std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM,
185 IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {
186
187 // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
188 // extracted module key, extracted module, and source module key
189 // together. This could be used, for example, to provide a specific
190 // memory manager instance to the linking layer.
191
192 auto &ES = getExecutionSession();
193 GlobalValueSet RequestedGVs;
194 for (auto &Name : R->getRequestedSymbols()) {
195 if (Name == R->getInitializerSymbol())
196 TSM.withModuleDo(F: [&](Module &M) {
197 for (auto &GV : getStaticInitGVs(M))
198 RequestedGVs.insert(x: &GV);
199 });
200 else {
201 assert(Defs.count(Name) && "No definition for symbol");
202 RequestedGVs.insert(x: Defs[Name]);
203 }
204 }
205
206 /// Perform partitioning with the context lock held, since the partition
207 /// function is allowed to access the globals to compute the partition.
208 auto GVsToExtract =
209 TSM.withModuleDo(F: [&](Module &M) { return Partition(RequestedGVs); });
210
211 // Take a 'None' partition to mean the whole module (as opposed to an empty
212 // partition, which means "materialize nothing"). Emit the whole module
213 // unmodified to the base layer.
214 if (GVsToExtract == std::nullopt) {
215 Defs.clear();
216 BaseLayer.emit(R: std::move(R), TSM: std::move(TSM));
217 return;
218 }
219
220 // If the partition is empty, return the whole module to the symbol table.
221 if (GVsToExtract->empty()) {
222 if (auto Err =
223 R->replace(MU: std::make_unique<PartitioningIRMaterializationUnit>(
224 args: std::move(TSM),
225 args: MaterializationUnit::Interface(R->getSymbols(),
226 R->getInitializerSymbol()),
227 args: std::move(Defs), args&: *this))) {
228 getExecutionSession().reportError(Err: std::move(Err));
229 R->failMaterialization();
230 return;
231 }
232 return;
233 }
234
235 // Ok -- we actually need to partition the symbols. Promote the symbol
236 // linkages/names, expand the partition to include any required symbols
237 // (i.e. symbols that can't be separated from our partition), and
238 // then extract the partition.
239 //
240 // FIXME: We apply this promotion once per partitioning. It's safe, but
241 // overkill.
242 auto ExtractedTSM = TSM.withModuleDo(F: [&](Module &M)
243 -> Expected<ThreadSafeModule> {
244 auto PromotedGlobals = PromoteSymbols(M);
245 if (!PromotedGlobals.empty()) {
246
247 MangleAndInterner Mangle(ES, M.getDataLayout());
248 SymbolFlagsMap SymbolFlags;
249 IRSymbolMapper::add(ES, MO: *getManglingOptions(), GVs: PromotedGlobals,
250 SymbolFlags);
251
252 if (auto Err = R->defineMaterializing(SymbolFlags))
253 return std::move(Err);
254 }
255
256 expandPartition(Partition&: *GVsToExtract);
257
258 // Submodule name is given by hashing the names of the globals.
259 std::string SubModuleName;
260 {
261 std::vector<const GlobalValue *> HashGVs;
262 HashGVs.reserve(n: GVsToExtract->size());
263 llvm::append_range(C&: HashGVs, R&: *GVsToExtract);
264 llvm::sort(C&: HashGVs, Comp: [](const GlobalValue *LHS, const GlobalValue *RHS) {
265 return LHS->getName() < RHS->getName();
266 });
267 hash_code HC(0);
268 for (const auto *GV : HashGVs) {
269 assert(GV->hasName() && "All GVs to extract should be named by now");
270 auto GVName = GV->getName();
271 HC = hash_combine(args: HC, args: hash_combine_range(R&: GVName));
272 }
273 raw_string_ostream(SubModuleName)
274 << ".submodule."
275 << formatv(Fmt: sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}",
276 Vals: static_cast<size_t>(HC))
277 << ".ll";
278 }
279
280 // Extract the requested partiton (plus any necessary aliases) and
281 // put the rest back into the impl dylib.
282 auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
283 return GVsToExtract->count(x: &GV);
284 };
285
286 return extractSubModule(TSM, Suffix: SubModuleName, ShouldExtract);
287 });
288
289 if (!ExtractedTSM) {
290 ES.reportError(Err: ExtractedTSM.takeError());
291 R->failMaterialization();
292 return;
293 }
294
295 if (auto Err = R->replace(MU: std::make_unique<PartitioningIRMaterializationUnit>(
296 args&: ES, args: *getManglingOptions(), args: std::move(TSM), args&: *this))) {
297 ES.reportError(Err: std::move(Err));
298 R->failMaterialization();
299 return;
300 }
301 BaseLayer.emit(R: std::move(R), TSM: std::move(*ExtractedTSM));
302}
303