CompileOnDemandLayer.cpp revision 360784 
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
13using namespace llvm;
14using namespace llvm::orc;
15
16static 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
65namespace llvm {
66namespace orc {
67
68class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
69public:
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
86private:
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
103Optional<CompileOnDemandLayer::GlobalValueSet>
104CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) {
105  return std::move(Requested);
106}
107
108Optional<CompileOnDemandLayer::GlobalValueSet>
109CompileOnDemandLayer::compileWholeModule(GlobalValueSet Requested) {
110  return None;
111}
112
113CompileOnDemandLayer::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
120void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) {
121  this->Partition = std::move(Partition);
122}
123
124void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) {
125  this->AliaseeImpls = Imp;
126}
127void 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
170CompileOnDemandLayer::PerDylibResources &
171CompileOnDemandLayer::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
197void 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
210void 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
242void 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