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