contrib/llvm-project/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp

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