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 #include "llvm/Support/raw_ostream.h"
14 
15 using namespace llvm;
16 using namespace llvm::orc;
17 
18 namespace {
19 
20 template <typename MaterializerFtor>
21 class LambdaValueMaterializer final : public ValueMaterializer {
22 public:
23   LambdaValueMaterializer(MaterializerFtor M) : M(std::move(M)) {}
24 
25   Value *materialize(Value *V) final { return M(V); }
26 
27 private:
28   MaterializerFtor M;
29 };
30 
31 template <typename MaterializerFtor>
32 LambdaValueMaterializer<MaterializerFtor>
33 createLambdaValueMaterializer(MaterializerFtor M) {
34   return LambdaValueMaterializer<MaterializerFtor>(std::move(M));
35 }
36 } // namespace
37 
38 static void extractAliases(MaterializationResponsibility &R, Module &M,
39                            MangleAndInterner &Mangle) {
40   SymbolAliasMap Aliases;
41 
42   std::vector<GlobalAlias *> ModAliases;
43   for (auto &A : M.aliases())
44     ModAliases.push_back(&A);
45 
46   for (auto *A : ModAliases) {
47     Constant *Aliasee = A->getAliasee();
48     assert(A->hasName() && "Anonymous alias?");
49     assert(Aliasee->hasName() && "Anonymous aliasee");
50     std::string AliasName = A->getName();
51 
52     Aliases[Mangle(AliasName)] = SymbolAliasMapEntry(
53         {Mangle(Aliasee->getName()), JITSymbolFlags::fromGlobalValue(*A)});
54 
55     if (isa<Function>(Aliasee)) {
56       auto *F = cloneFunctionDecl(M, *cast<Function>(Aliasee));
57       A->replaceAllUsesWith(F);
58       A->eraseFromParent();
59       F->setName(AliasName);
60     } else if (isa<GlobalValue>(Aliasee)) {
61       auto *G = cloneGlobalVariableDecl(M, *cast<GlobalVariable>(Aliasee));
62       A->replaceAllUsesWith(G);
63       A->eraseFromParent();
64       G->setName(AliasName);
65     }
66   }
67 
68   R.delegate(symbolAliases(std::move(Aliases)));
69 }
70 
71 static std::unique_ptr<Module> extractGlobals(Module &M) {
72   // FIXME: Add alias support.
73 
74   auto GlobalsModule = llvm::make_unique<Module>(
75       (M.getName() + ".globals").str(), M.getContext());
76   GlobalsModule->setDataLayout(M.getDataLayout());
77 
78   ValueToValueMapTy VMap;
79 
80   for (auto &GV : M.globals())
81     if (!GV.isDeclaration() && !VMap.count(&GV))
82       cloneGlobalVariableDecl(*GlobalsModule, GV, &VMap);
83 
84   // Clone the module flags.
85   cloneModuleFlagsMetadata(*GlobalsModule, M, VMap);
86 
87   auto Materializer = createLambdaValueMaterializer([&](Value *V) -> Value * {
88     if (auto *F = dyn_cast<Function>(V))
89       return cloneFunctionDecl(*GlobalsModule, *F);
90     return nullptr;
91   });
92 
93   // Move the global variable initializers.
94   for (auto &GV : M.globals()) {
95     if (!GV.isDeclaration())
96       moveGlobalVariableInitializer(GV, VMap, &Materializer);
97     GV.setInitializer(nullptr);
98   }
99 
100   return GlobalsModule;
101 }
102 
103 namespace llvm {
104 namespace orc {
105 
106 class ExtractingIRMaterializationUnit : public IRMaterializationUnit {
107 public:
108   ExtractingIRMaterializationUnit(
109       ExecutionSession &ES, CompileOnDemandLayer2 &Parent,
110       std::unique_ptr<Module> M,
111       std::shared_ptr<SymbolResolver> BackingResolver)
112       : IRMaterializationUnit(ES, std::move(M)), Parent(Parent),
113         BackingResolver(std::move(BackingResolver)) {}
114 
115   ExtractingIRMaterializationUnit(
116       std::unique_ptr<Module> M, SymbolFlagsMap SymbolFlags,
117       SymbolNameToDefinitionMap SymbolToDefinition,
118       CompileOnDemandLayer2 &Parent,
119       std::shared_ptr<SymbolResolver> BackingResolver)
120       : IRMaterializationUnit(std::move(M), std::move(SymbolFlags),
121                               std::move(SymbolToDefinition)),
122         Parent(Parent), BackingResolver(std::move(BackingResolver)) {}
123 
124 private:
125   void materialize(MaterializationResponsibility R) override {
126     // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
127     //        extracted module key, extracted module, and source module key
128     //        together. This could be used, for example, to provide a specific
129     //        memory manager instance to the linking layer.
130 
131     // FIXME: The derived constructor should *only* look for the names of
132     //        original function definitions in the target VSO. All other
133     //        symbols should be looked up in the backing resolver.
134 
135     // Find the functions that have been requested.
136     auto RequestedSymbols = R.getRequestedSymbols();
137 
138     // Extract them into a new module.
139     auto ExtractedFunctionsModule =
140         Parent.extractFunctions(*M, RequestedSymbols, SymbolToDefinition);
141 
142     // Build a new ExtractingIRMaterializationUnit to delegate the unrequested
143     // symbols to.
144     SymbolFlagsMap DelegatedSymbolFlags;
145     IRMaterializationUnit::SymbolNameToDefinitionMap
146         DelegatedSymbolToDefinition;
147     for (auto &KV : SymbolToDefinition) {
148       if (RequestedSymbols.count(KV.first))
149         continue;
150       DelegatedSymbolFlags[KV.first] =
151           JITSymbolFlags::fromGlobalValue(*KV.second);
152       DelegatedSymbolToDefinition[KV.first] = KV.second;
153     }
154 
155     if (!DelegatedSymbolFlags.empty()) {
156       assert(DelegatedSymbolFlags.size() ==
157                  DelegatedSymbolToDefinition.size() &&
158              "SymbolFlags and SymbolToDefinition should have the same number "
159              "of entries");
160       R.delegate(llvm::make_unique<ExtractingIRMaterializationUnit>(
161           std::move(M), std::move(DelegatedSymbolFlags),
162           std::move(DelegatedSymbolToDefinition), Parent, BackingResolver));
163     }
164 
165     Parent.emitExtractedFunctionsModule(
166         std::move(R), std::move(ExtractedFunctionsModule), BackingResolver);
167   }
168 
169   void discard(const VSO &V, SymbolStringPtr Name) override {
170     // All original symbols were materialized by the CODLayer and should be
171     // final. The function bodies provided by M should never be overridden.
172     llvm_unreachable("Discard should never be called on an "
173                      "ExtractingIRMaterializationUnit");
174   }
175 
176   CompileOnDemandLayer2 &Parent;
177   std::shared_ptr<SymbolResolver> BackingResolver;
178 };
179 
180 CompileOnDemandLayer2::CompileOnDemandLayer2(
181     ExecutionSession &ES, IRLayer &BaseLayer, JITCompileCallbackManager &CCMgr,
182     IndirectStubsManagerBuilder BuildIndirectStubsManager,
183     GetSymbolResolverFunction GetSymbolResolver,
184     SetSymbolResolverFunction SetSymbolResolver,
185     GetAvailableContextFunction GetAvailableContext)
186     : IRLayer(ES), BaseLayer(BaseLayer), CCMgr(CCMgr),
187       BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)),
188       GetSymbolResolver(std::move(GetSymbolResolver)),
189       SetSymbolResolver(std::move(SetSymbolResolver)),
190       GetAvailableContext(std::move(GetAvailableContext)) {}
191 
192 Error CompileOnDemandLayer2::add(VSO &V, VModuleKey K,
193                                  std::unique_ptr<Module> M) {
194   return IRLayer::add(V, K, std::move(M));
195 }
196 
197 void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
198                                  std::unique_ptr<Module> M) {
199   auto &ES = getExecutionSession();
200   assert(M && "M should not be null");
201 
202   for (auto &GV : M->global_values())
203     if (GV.hasWeakLinkage())
204       GV.setLinkage(GlobalValue::ExternalLinkage);
205 
206   MangleAndInterner Mangle(ES, M->getDataLayout());
207 
208   extractAliases(R, *M, Mangle);
209 
210   auto GlobalsModule = extractGlobals(*M);
211 
212   // Delete the bodies of any available externally functions, rename the
213   // rest, and build the compile callbacks.
214   std::map<SymbolStringPtr, std::pair<JITTargetAddress, JITSymbolFlags>>
215       StubCallbacksAndLinkages;
216   auto &TargetVSO = R.getTargetVSO();
217 
218   for (auto &F : M->functions()) {
219     if (F.isDeclaration())
220       continue;
221 
222     if (F.hasAvailableExternallyLinkage()) {
223       F.deleteBody();
224       continue;
225     }
226 
227     assert(F.hasName() && "Function should have a name");
228     std::string StubUnmangledName = F.getName();
229     F.setName(F.getName() + "$body");
230     auto StubDecl = cloneFunctionDecl(*M, F);
231     StubDecl->setName(StubUnmangledName);
232     F.replaceAllUsesWith(StubDecl);
233     auto StubName = Mangle(StubUnmangledName);
234     auto BodyName = Mangle(F.getName());
235     if (auto CallbackAddr = CCMgr.getCompileCallback(
236             [BodyName, &TargetVSO, &ES]() -> JITTargetAddress {
237               if (auto Sym = lookup({&TargetVSO}, BodyName))
238                 return Sym->getAddress();
239               else {
240                 ES.reportError(Sym.takeError());
241                 return 0;
242               }
243             })) {
244       auto Flags = JITSymbolFlags::fromGlobalValue(F);
245       Flags &= ~JITSymbolFlags::Weak;
246       StubCallbacksAndLinkages[std::move(StubName)] =
247           std::make_pair(*CallbackAddr, Flags);
248     } else {
249       ES.reportError(CallbackAddr.takeError());
250       R.failMaterialization();
251       return;
252     }
253   }
254 
255   // Build the stub inits map.
256   IndirectStubsManager::StubInitsMap StubInits;
257   for (auto &KV : StubCallbacksAndLinkages)
258     StubInits[*KV.first] = KV.second;
259 
260   // Build the function-body-extracting materialization unit.
261   auto SR = GetSymbolResolver(K);
262   if (auto Err = R.getTargetVSO().define(
263           llvm::make_unique<ExtractingIRMaterializationUnit>(
264               ES, *this, std::move(M), SR))) {
265     ES.reportError(std::move(Err));
266     R.failMaterialization();
267     return;
268   }
269 
270   // Replace the fallback symbol resolver: We will re-use M's VModuleKey for
271   // the GlobalsModule.
272   SetSymbolResolver(K, SR);
273 
274   // Build the stubs.
275   // FIXME: Remove function bodies materialization unit if stub creation fails.
276   auto &StubsMgr = getStubsManager(TargetVSO);
277   if (auto Err = StubsMgr.createStubs(StubInits)) {
278     ES.reportError(std::move(Err));
279     R.failMaterialization();
280     return;
281   }
282 
283   // Resolve and finalize stubs.
284   SymbolMap ResolvedStubs;
285   for (auto &KV : StubCallbacksAndLinkages) {
286     if (auto Sym = StubsMgr.findStub(*KV.first, false))
287       ResolvedStubs[KV.first] = Sym;
288     else
289       llvm_unreachable("Stub went missing");
290   }
291 
292   R.resolve(ResolvedStubs);
293 
294   BaseLayer.emit(std::move(R), std::move(K), std::move(GlobalsModule));
295 }
296 
297 IndirectStubsManager &CompileOnDemandLayer2::getStubsManager(const VSO &V) {
298   std::lock_guard<std::mutex> Lock(CODLayerMutex);
299   StubManagersMap::iterator I = StubsMgrs.find(&V);
300   if (I == StubsMgrs.end())
301     I = StubsMgrs.insert(std::make_pair(&V, BuildIndirectStubsManager())).first;
302   return *I->second;
303 }
304 
305 std::unique_ptr<Module> CompileOnDemandLayer2::extractFunctions(
306     Module &M, const SymbolNameSet &SymbolNames,
307     const SymbolNameToDefinitionMap &SymbolToDefinition) {
308   assert(!SymbolNames.empty() && "Can not extract an empty function set");
309 
310   std::string ExtractedModName;
311   {
312     raw_string_ostream ExtractedModNameStream(ExtractedModName);
313     ExtractedModNameStream << M.getName();
314     for (auto &Name : SymbolNames)
315       ExtractedModNameStream << "." << *Name;
316   }
317 
318   auto ExtractedFunctionsModule =
319       llvm::make_unique<Module>(ExtractedModName, GetAvailableContext());
320   ExtractedFunctionsModule->setDataLayout(M.getDataLayout());
321 
322   ValueToValueMapTy VMap;
323 
324   auto Materializer = createLambdaValueMaterializer([&](Value *V) -> Value * {
325     if (auto *F = dyn_cast<Function>(V))
326       return cloneFunctionDecl(*ExtractedFunctionsModule, *F);
327     else if (auto *GV = dyn_cast<GlobalVariable>(V))
328       return cloneGlobalVariableDecl(*ExtractedFunctionsModule, *GV);
329     return nullptr;
330   });
331 
332   std::vector<std::pair<Function *, Function *>> OrigToNew;
333   for (auto &FunctionName : SymbolNames) {
334     assert(SymbolToDefinition.count(FunctionName) &&
335            "No definition for symbol");
336     auto *OrigF = cast<Function>(SymbolToDefinition.find(FunctionName)->second);
337     auto *NewF = cloneFunctionDecl(*ExtractedFunctionsModule, *OrigF, &VMap);
338     OrigToNew.push_back(std::make_pair(OrigF, NewF));
339   }
340 
341   for (auto &KV : OrigToNew)
342     moveFunctionBody(*KV.first, VMap, &Materializer, KV.second);
343 
344   return ExtractedFunctionsModule;
345 }
346 
347 void CompileOnDemandLayer2::emitExtractedFunctionsModule(
348     MaterializationResponsibility R, std::unique_ptr<Module> M,
349     std::shared_ptr<SymbolResolver> Resolver) {
350   auto &TargetVSO = R.getTargetVSO();
351   auto K = getExecutionSession().allocateVModule();
352 
353   auto ExtractedFunctionsResolver = createSymbolResolver(
354       [=](SymbolFlagsMap &Flags, const SymbolNameSet &Symbols) {
355         return Resolver->lookupFlags(Flags, Symbols);
356       },
357       [=, &TargetVSO](std::shared_ptr<AsynchronousSymbolQuery> Query,
358                       SymbolNameSet Symbols) {
359         auto RemainingSymbols = TargetVSO.lookup(Query, std::move(Symbols));
360         return Resolver->lookup(std::move(Query), std::move(RemainingSymbols));
361       });
362 
363   SetSymbolResolver(K, std::move(ExtractedFunctionsResolver));
364   BaseLayer.emit(std::move(R), std::move(K), std::move(M));
365 }
366 
367 } // end namespace orc
368 } // end namespace llvm
369