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