1 //===----- CompileOnDemandLayer.cpp - Lazily emit IR on first call --------===//
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/CompileOnDemandLayer.h"
10 #include "llvm/IR/Mangler.h"
11 #include "llvm/IR/Module.h"
12
13 using namespace llvm;
14 using namespace llvm::orc;
15
extractSubModule(ThreadSafeModule & TSM,StringRef Suffix,GVPredicate ShouldExtract)16 static 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>(GV)) {
26 auto &F = cast<Function>(GV);
27 F.deleteBody();
28 F.setPersonalityFn(nullptr);
29 } else if (isa<GlobalVariable>(GV)) {
30 cast<GlobalVariable>(GV).setInitializer(nullptr);
31 } else if (isa<GlobalAlias>(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>(GV);
35 Constant *Aliasee = A.getAliasee();
36 assert(A.hasName() && "Anonymous alias?");
37 assert(Aliasee->hasName() && "Anonymous aliasee");
38 std::string AliasName = A.getName();
39
40 if (isa<Function>(Aliasee)) {
41 auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee));
42 A.replaceAllUsesWith(F);
43 A.eraseFromParent();
44 F->setName(AliasName);
45 } else if (isa<GlobalVariable>(Aliasee)) {
46 auto *G = cloneGlobalVariableDecl(*A.getParent(),
47 *cast<GlobalVariable>(Aliasee));
48 A.replaceAllUsesWith(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(TSM, ShouldExtract, DeleteExtractedDefs);
58 NewTSM.withModuleDo([&](Module &M) {
59 M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
60 });
61
62 return NewTSM;
63 }
64
65 namespace llvm {
66 namespace orc {
67
68 class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
69 public:
PartitioningIRMaterializationUnit(ExecutionSession & ES,const ManglingOptions & MO,ThreadSafeModule TSM,VModuleKey K,CompileOnDemandLayer & Parent)70 PartitioningIRMaterializationUnit(ExecutionSession &ES,
71 const ManglingOptions &MO,
72 ThreadSafeModule TSM, VModuleKey K,
73 CompileOnDemandLayer &Parent)
74 : IRMaterializationUnit(ES, MO, std::move(TSM), std::move(K)),
75 Parent(Parent) {}
76
PartitioningIRMaterializationUnit(ThreadSafeModule TSM,SymbolFlagsMap SymbolFlags,SymbolNameToDefinitionMap SymbolToDefinition,CompileOnDemandLayer & Parent)77 PartitioningIRMaterializationUnit(
78 ThreadSafeModule TSM, SymbolFlagsMap SymbolFlags,
79 SymbolNameToDefinitionMap SymbolToDefinition,
80 CompileOnDemandLayer &Parent)
81 : IRMaterializationUnit(std::move(TSM), std::move(K),
82 std::move(SymbolFlags),
83 std::move(SymbolToDefinition)),
84 Parent(Parent) {}
85
86 private:
materialize(MaterializationResponsibility R)87 void materialize(MaterializationResponsibility R) override {
88 Parent.emitPartition(std::move(R), std::move(TSM),
89 std::move(SymbolToDefinition));
90 }
91
discard(const JITDylib & V,const SymbolStringPtr & Name)92 void discard(const JITDylib &V, const SymbolStringPtr &Name) override {
93 // All original symbols were materialized by the CODLayer and should be
94 // final. The function bodies provided by M should never be overridden.
95 llvm_unreachable("Discard should never be called on an "
96 "ExtractingIRMaterializationUnit");
97 }
98
99 mutable std::mutex SourceModuleMutex;
100 CompileOnDemandLayer &Parent;
101 };
102
103 Optional<CompileOnDemandLayer::GlobalValueSet>
compileRequested(GlobalValueSet Requested)104 CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) {
105 return std::move(Requested);
106 }
107
108 Optional<CompileOnDemandLayer::GlobalValueSet>
compileWholeModule(GlobalValueSet Requested)109 CompileOnDemandLayer::compileWholeModule(GlobalValueSet Requested) {
110 return None;
111 }
112
CompileOnDemandLayer(ExecutionSession & ES,IRLayer & BaseLayer,LazyCallThroughManager & LCTMgr,IndirectStubsManagerBuilder BuildIndirectStubsManager)113 CompileOnDemandLayer::CompileOnDemandLayer(
114 ExecutionSession &ES, IRLayer &BaseLayer, LazyCallThroughManager &LCTMgr,
115 IndirectStubsManagerBuilder BuildIndirectStubsManager)
116 : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer),
117 LCTMgr(LCTMgr),
118 BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {}
119
setPartitionFunction(PartitionFunction Partition)120 void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) {
121 this->Partition = std::move(Partition);
122 }
123
setImplMap(ImplSymbolMap * Imp)124 void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) {
125 this->AliaseeImpls = Imp;
126 }
emit(MaterializationResponsibility R,ThreadSafeModule TSM)127 void CompileOnDemandLayer::emit(MaterializationResponsibility R,
128 ThreadSafeModule TSM) {
129 assert(TSM && "Null module");
130
131 auto &ES = getExecutionSession();
132
133 // Sort the callables and non-callables, build re-exports and lodge the
134 // actual module with the implementation dylib.
135 auto &PDR = getPerDylibResources(R.getTargetJITDylib());
136
137 SymbolAliasMap NonCallables;
138 SymbolAliasMap Callables;
139 TSM.withModuleDo([&](Module &M) {
140 // First, do some cleanup on the module:
141 cleanUpModule(M);
142 });
143
144 for (auto &KV : R.getSymbols()) {
145 auto &Name = KV.first;
146 auto &Flags = KV.second;
147 if (Flags.isCallable())
148 Callables[Name] = SymbolAliasMapEntry(Name, Flags);
149 else
150 NonCallables[Name] = SymbolAliasMapEntry(Name, Flags);
151 }
152
153 // Create a partitioning materialization unit and lodge it with the
154 // implementation dylib.
155 if (auto Err = PDR.getImplDylib().define(
156 std::make_unique<PartitioningIRMaterializationUnit>(
157 ES, *getManglingOptions(), std::move(TSM), R.getVModuleKey(),
158 *this))) {
159 ES.reportError(std::move(Err));
160 R.failMaterialization();
161 return;
162 }
163
164 R.replace(reexports(PDR.getImplDylib(), std::move(NonCallables),
165 JITDylibLookupFlags::MatchAllSymbols));
166 R.replace(lazyReexports(LCTMgr, PDR.getISManager(), PDR.getImplDylib(),
167 std::move(Callables), AliaseeImpls));
168 }
169
170 CompileOnDemandLayer::PerDylibResources &
getPerDylibResources(JITDylib & TargetD)171 CompileOnDemandLayer::getPerDylibResources(JITDylib &TargetD) {
172 auto I = DylibResources.find(&TargetD);
173 if (I == DylibResources.end()) {
174 auto &ImplD =
175 getExecutionSession().createJITDylib(TargetD.getName() + ".impl");
176 TargetD.withSearchOrderDo(
177 [&](const JITDylibSearchOrder &TargetSearchOrder) {
178 auto NewSearchOrder = TargetSearchOrder;
179 assert(
180 !NewSearchOrder.empty() &&
181 NewSearchOrder.front().first == &TargetD &&
182 NewSearchOrder.front().second ==
183 JITDylibLookupFlags::MatchAllSymbols &&
184 "TargetD must be at the front of its own search order and match "
185 "non-exported symbol");
186 NewSearchOrder.insert(std::next(NewSearchOrder.begin()),
187 {&ImplD, JITDylibLookupFlags::MatchAllSymbols});
188 ImplD.setSearchOrder(std::move(NewSearchOrder), false);
189 });
190 PerDylibResources PDR(ImplD, BuildIndirectStubsManager());
191 I = DylibResources.insert(std::make_pair(&TargetD, std::move(PDR))).first;
192 }
193
194 return I->second;
195 }
196
cleanUpModule(Module & M)197 void CompileOnDemandLayer::cleanUpModule(Module &M) {
198 for (auto &F : M.functions()) {
199 if (F.isDeclaration())
200 continue;
201
202 if (F.hasAvailableExternallyLinkage()) {
203 F.deleteBody();
204 F.setPersonalityFn(nullptr);
205 continue;
206 }
207 }
208 }
209
expandPartition(GlobalValueSet & Partition)210 void CompileOnDemandLayer::expandPartition(GlobalValueSet &Partition) {
211 // Expands the partition to ensure the following rules hold:
212 // (1) If any alias is in the partition, its aliasee is also in the partition.
213 // (2) If any aliasee is in the partition, its aliases are also in the
214 // partiton.
215 // (3) If any global variable is in the partition then all global variables
216 // are in the partition.
217 assert(!Partition.empty() && "Unexpected empty partition");
218
219 const Module &M = *(*Partition.begin())->getParent();
220 bool ContainsGlobalVariables = false;
221 std::vector<const GlobalValue *> GVsToAdd;
222
223 for (auto *GV : Partition)
224 if (isa<GlobalAlias>(GV))
225 GVsToAdd.push_back(
226 cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));
227 else if (isa<GlobalVariable>(GV))
228 ContainsGlobalVariables = true;
229
230 for (auto &A : M.aliases())
231 if (Partition.count(cast<GlobalValue>(A.getAliasee())))
232 GVsToAdd.push_back(&A);
233
234 if (ContainsGlobalVariables)
235 for (auto &G : M.globals())
236 GVsToAdd.push_back(&G);
237
238 for (auto *GV : GVsToAdd)
239 Partition.insert(GV);
240 }
241
emitPartition(MaterializationResponsibility R,ThreadSafeModule TSM,IRMaterializationUnit::SymbolNameToDefinitionMap Defs)242 void CompileOnDemandLayer::emitPartition(
243 MaterializationResponsibility R, ThreadSafeModule TSM,
244 IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {
245
246 // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
247 // extracted module key, extracted module, and source module key
248 // together. This could be used, for example, to provide a specific
249 // memory manager instance to the linking layer.
250
251 auto &ES = getExecutionSession();
252 GlobalValueSet RequestedGVs;
253 for (auto &Name : R.getRequestedSymbols()) {
254 assert(Defs.count(Name) && "No definition for symbol");
255 RequestedGVs.insert(Defs[Name]);
256 }
257
258 /// Perform partitioning with the context lock held, since the partition
259 /// function is allowed to access the globals to compute the partition.
260 auto GVsToExtract =
261 TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); });
262
263 // Take a 'None' partition to mean the whole module (as opposed to an empty
264 // partition, which means "materialize nothing"). Emit the whole module
265 // unmodified to the base layer.
266 if (GVsToExtract == None) {
267 Defs.clear();
268 BaseLayer.emit(std::move(R), std::move(TSM));
269 return;
270 }
271
272 // If the partition is empty, return the whole module to the symbol table.
273 if (GVsToExtract->empty()) {
274 R.replace(std::make_unique<PartitioningIRMaterializationUnit>(
275 std::move(TSM), R.getSymbols(), std::move(Defs), *this));
276 return;
277 }
278
279 // Ok -- we actually need to partition the symbols. Promote the symbol
280 // linkages/names, expand the partition to include any required symbols
281 // (i.e. symbols that can't be separated from our partition), and
282 // then extract the partition.
283 //
284 // FIXME: We apply this promotion once per partitioning. It's safe, but
285 // overkill.
286
287 auto ExtractedTSM =
288 TSM.withModuleDo([&](Module &M) -> Expected<ThreadSafeModule> {
289 auto PromotedGlobals = PromoteSymbols(M);
290 if (!PromotedGlobals.empty()) {
291 MangleAndInterner Mangle(ES, M.getDataLayout());
292 SymbolFlagsMap SymbolFlags;
293 for (auto &GV : PromotedGlobals)
294 SymbolFlags[Mangle(GV->getName())] =
295 JITSymbolFlags::fromGlobalValue(*GV);
296 if (auto Err = R.defineMaterializing(SymbolFlags))
297 return std::move(Err);
298 }
299
300 expandPartition(*GVsToExtract);
301
302 // Extract the requested partiton (plus any necessary aliases) and
303 // put the rest back into the impl dylib.
304 auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
305 return GVsToExtract->count(&GV);
306 };
307
308 return extractSubModule(TSM, ".submodule", ShouldExtract);
309 });
310
311 if (!ExtractedTSM) {
312 ES.reportError(ExtractedTSM.takeError());
313 R.failMaterialization();
314 return;
315 }
316
317 R.replace(std::make_unique<PartitioningIRMaterializationUnit>(
318 ES, *getManglingOptions(), std::move(TSM), R.getVModuleKey(), *this));
319 BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));
320 }
321
322 } // end namespace orc
323 } // end namespace llvm
324