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

   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), true));
 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 = getExecutionSession().createJITDylib(
 170         TargetD.getName() + ".impl", false);
 171     TargetD.withSearchOrderDo([&](const JITDylibSearchList &TargetSearchOrder) {
 172       auto NewSearchOrder = TargetSearchOrder;
 173       assert(!NewSearchOrder.empty() &&
 174              NewSearchOrder.front().first == &TargetD &&
 175              NewSearchOrder.front().second == true &&
 176              "TargetD must be at the front of its own search order and match "
 177              "non-exported symbol");
 178       NewSearchOrder.insert(std::next(NewSearchOrder.begin()), {&ImplD, true});
 179       ImplD.setSearchOrder(std::move(NewSearchOrder), false);
 180     });
 181     PerDylibResources PDR(ImplD, BuildIndirectStubsManager());
 182     I = DylibResources.insert(std::make_pair(&TargetD, std::move(PDR))).first;
 183   }
 184
 185   return I->second;
 186 }
 187
 188 void CompileOnDemandLayer::cleanUpModule(Module &M) {
 189   for (auto &F : M.functions()) {
 190     if (F.isDeclaration())
 191       continue;
 192
 193     if (F.hasAvailableExternallyLinkage()) {
 194       F.deleteBody();
 195       F.setPersonalityFn(nullptr);
 196       continue;
 197     }
 198   }
 199 }
 200
 201 void CompileOnDemandLayer::expandPartition(GlobalValueSet &Partition) {
 202   // Expands the partition to ensure the following rules hold:
 203   // (1) If any alias is in the partition, its aliasee is also in the partition.
 204   // (2) If any aliasee is in the partition, its aliases are also in the
 205   //     partiton.
 206   // (3) If any global variable is in the partition then all global variables
 207   //     are in the partition.
 208   assert(!Partition.empty() && "Unexpected empty partition");
 209
 210   const Module &M = *(*Partition.begin())->getParent();
 211   bool ContainsGlobalVariables = false;
 212   std::vector<const GlobalValue *> GVsToAdd;
 213
 214   for (auto *GV : Partition)
 215     if (isa<GlobalAlias>(GV))
 216       GVsToAdd.push_back(
 217           cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));
 218     else if (isa<GlobalVariable>(GV))
 219       ContainsGlobalVariables = true;
 220
 221   for (auto &A : M.aliases())
 222     if (Partition.count(cast<GlobalValue>(A.getAliasee())))
 223       GVsToAdd.push_back(&A);
 224
 225   if (ContainsGlobalVariables)
 226     for (auto &G : M.globals())
 227       GVsToAdd.push_back(&G);
 228
 229   for (auto *GV : GVsToAdd)
 230     Partition.insert(GV);
 231 }
 232
 233 void CompileOnDemandLayer::emitPartition(
 234     MaterializationResponsibility R, ThreadSafeModule TSM,
 235     IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {
 236
 237   // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
 238   //        extracted module key, extracted module, and source module key
 239   //        together. This could be used, for example, to provide a specific
 240   //        memory manager instance to the linking layer.
 241
 242   auto &ES = getExecutionSession();
 243
 244   GlobalValueSet RequestedGVs;
 245   for (auto &Name : R.getRequestedSymbols()) {
 246     assert(Defs.count(Name) && "No definition for symbol");
 247     RequestedGVs.insert(Defs[Name]);
 248   }
 249
 250   auto GVsToExtract = Partition(RequestedGVs);
 251
 252   // Take a 'None' partition to mean the whole module (as opposed to an empty
 253   // partition, which means "materialize nothing"). Emit the whole module
 254   // unmodified to the base layer.
 255   if (GVsToExtract == None) {
 256     Defs.clear();
 257     BaseLayer.emit(std::move(R), std::move(TSM));
 258     return;
 259   }
 260
 261   // If the partition is empty, return the whole module to the symbol table.
 262   if (GVsToExtract->empty()) {
 263     R.replace(llvm::make_unique<PartitioningIRMaterializationUnit>(
 264         std::move(TSM), R.getSymbols(), std::move(Defs), *this));
 265     return;
 266   }
 267
 268   // Ok -- we actually need to partition the symbols. Promote the symbol
 269   // linkages/names.
 270   // FIXME: We apply this once per partitioning. It's safe, but overkill.
 271   {
 272     auto PromotedGlobals = PromoteSymbols(*TSM.getModule());
 273     if (!PromotedGlobals.empty()) {
 274       MangleAndInterner Mangle(ES, TSM.getModule()->getDataLayout());
 275       SymbolFlagsMap SymbolFlags;
 276       for (auto &GV : PromotedGlobals)
 277         SymbolFlags[Mangle(GV->getName())] =
 278             JITSymbolFlags::fromGlobalValue(*GV);
 279       if (auto Err = R.defineMaterializing(SymbolFlags)) {
 280         ES.reportError(std::move(Err));
 281         R.failMaterialization();
 282         return;
 283       }
 284     }
 285   }
 286
 287   expandPartition(*GVsToExtract);
 288
 289   // Extract the requested partiton (plus any necessary aliases) and
 290   // put the rest back into the impl dylib.
 291   auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
 292     return GVsToExtract->count(&GV);
 293   };
 294
 295   auto ExtractedTSM = extractSubModule(TSM, ".submodule", ShouldExtract);
 296   R.replace(llvm::make_unique<PartitioningIRMaterializationUnit>(
 297       ES, std::move(TSM), R.getVModuleKey(), *this));
 298
 299   BaseLayer.emit(std::move(R), std::move(ExtractedTSM));
 300 }
 301
 302 } // end namespace orc
 303 } // end namespace llvm