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