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