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
17using namespace llvm;
18using namespace llvm::orc;
19
20static 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
69namespace llvm {
70namespace orc {
71
72class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
73public:
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
88private:
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
105std::optional<CompileOnDemandLayer::GlobalValueSet>
106CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) {
107  return std::move(Requested);
108}
109
110std::optional<CompileOnDemandLayer::GlobalValueSet>
111CompileOnDemandLayer::compileWholeModule(GlobalValueSet Requested) {
112  return std::nullopt;
113}
114
115CompileOnDemandLayer::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
122void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) {
123  this->Partition = std::move(Partition);
124}
125
126void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) {
127  this->AliaseeImpls = Imp;
128}
129void 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
184CompileOnDemandLayer::PerDylibResources &
185CompileOnDemandLayer::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
214void 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
227void 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 (const 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 (const auto *GV : GVsToAdd)
256    Partition.insert(GV);
257}
258
259void 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 == std::nullopt) {
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 (const 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 (const 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