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