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