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