Home | History | Annotate | Download | only in Orc 
      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/Bitcode/BitcodeReader.h"
     12 #include "llvm/Bitcode/BitcodeWriter.h"
     13 #include "llvm/IR/Mangler.h"
     14 #include "llvm/IR/Module.h"
     15 #include "llvm/Support/raw_ostream.h"
     16 #include "llvm/Transforms/Utils/Cloning.h"
     17 
     18 using namespace llvm;
     19 using namespace llvm::orc;
     20 
     21 namespace {
     22 
     23 template <typename MaterializerFtor>
     24 class LambdaValueMaterializer final : public ValueMaterializer {
     25 public:
     26   LambdaValueMaterializer(MaterializerFtor M) : M(std::move(M)) {}
     27 
     28   Value *materialize(Value *V) final { return M(V); }
     29 
     30 private:
     31   MaterializerFtor M;
     32 };
     33 
     34 template <typename MaterializerFtor>
     35 LambdaValueMaterializer<MaterializerFtor>
     36 createLambdaValueMaterializer(MaterializerFtor M) {
     37   return LambdaValueMaterializer<MaterializerFtor>(std::move(M));
     38 }
     39 } // namespace
     40 
     41 static void extractAliases(MaterializationResponsibility &R, Module &M,
     42                            MangleAndInterner &Mangle) {
     43   SymbolAliasMap Aliases;
     44 
     45   std::vector<GlobalAlias *> ModAliases;
     46   for (auto &A : M.aliases())
     47     ModAliases.push_back(&A);
     48 
     49   for (auto *A : ModAliases) {
     50     Constant *Aliasee = A->getAliasee();
     51     assert(A->hasName() && "Anonymous alias?");
     52     assert(Aliasee->hasName() && "Anonymous aliasee");
     53     std::string AliasName = A->getName();
     54 
     55     Aliases[Mangle(AliasName)] = SymbolAliasMapEntry(
     56         {Mangle(Aliasee->getName()), JITSymbolFlags::fromGlobalValue(*A)});
     57 
     58     if (isa<Function>(Aliasee)) {
     59       auto *F = cloneFunctionDecl(M, *cast<Function>(Aliasee));
     60       A->replaceAllUsesWith(F);
     61       A->eraseFromParent();
     62       F->setName(AliasName);
     63     } else if (isa<GlobalValue>(Aliasee)) {
     64       auto *G = cloneGlobalVariableDecl(M, *cast<GlobalVariable>(Aliasee));
     65       A->replaceAllUsesWith(G);
     66       A->eraseFromParent();
     67       G->setName(AliasName);
     68     }
     69   }
     70 
     71   R.replace(symbolAliases(std::move(Aliases)));
     72 }
     73 
     74 static std::unique_ptr<Module>
     75 extractAndClone(Module &M, LLVMContext &NewContext, StringRef Suffix,
     76                 function_ref<bool(const GlobalValue *)> ShouldCloneDefinition) {
     77   SmallVector<char, 1> ClonedModuleBuffer;
     78 
     79   {
     80     std::set<GlobalValue *> ClonedDefsInSrc;
     81     ValueToValueMapTy VMap;
     82     auto Tmp = CloneModule(M, VMap, [&](const GlobalValue *GV) {
     83       if (ShouldCloneDefinition(GV)) {
     84         ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
     85         return true;
     86       }
     87       return false;
     88     });
     89 
     90     for (auto *GV : ClonedDefsInSrc) {
     91       // Delete the definition and bump the linkage in the source module.
     92       if (isa<Function>(GV)) {
     93         auto &F = *cast<Function>(GV);
     94         F.deleteBody();
     95         F.setPersonalityFn(nullptr);
     96       } else if (isa<GlobalVariable>(GV)) {
     97         cast<GlobalVariable>(GV)->setInitializer(nullptr);
     98       } else
     99         llvm_unreachable("Unsupported global type");
    100 
    101       GV->setLinkage(GlobalValue::ExternalLinkage);
    102     }
    103 
    104     BitcodeWriter BCWriter(ClonedModuleBuffer);
    105 
    106     BCWriter.writeModule(*Tmp);
    107     BCWriter.writeSymtab();
    108     BCWriter.writeStrtab();
    109   }
    110 
    111   MemoryBufferRef ClonedModuleBufferRef(
    112       StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
    113       "cloned module buffer");
    114 
    115   auto ClonedModule =
    116       cantFail(parseBitcodeFile(ClonedModuleBufferRef, NewContext));
    117   ClonedModule->setModuleIdentifier((M.getName() + Suffix).str());
    118   return ClonedModule;
    119 }
    120 
    121 static std::unique_ptr<Module> extractGlobals(Module &M,
    122                                               LLVMContext &NewContext) {
    123   return extractAndClone(M, NewContext, ".globals", [](const GlobalValue *GV) {
    124     return isa<GlobalVariable>(GV);
    125   });
    126 }
    127 
    128 namespace llvm {
    129 namespace orc {
    130 
    131 class ExtractingIRMaterializationUnit : public IRMaterializationUnit {
    132 public:
    133   ExtractingIRMaterializationUnit(ExecutionSession &ES,
    134                                   CompileOnDemandLayer2 &Parent,
    135                                   std::unique_ptr<Module> M)
    136       : IRMaterializationUnit(ES, std::move(M)), Parent(Parent) {}
    137 
    138   ExtractingIRMaterializationUnit(std::unique_ptr<Module> M,
    139                                   SymbolFlagsMap SymbolFlags,
    140                                   SymbolNameToDefinitionMap SymbolToDefinition,
    141                                   CompileOnDemandLayer2 &Parent)
    142       : IRMaterializationUnit(std::move(M), std::move(SymbolFlags),
    143                               std::move(SymbolToDefinition)),
    144         Parent(Parent) {}
    145 
    146 private:
    147   void materialize(MaterializationResponsibility R) override {
    148     // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
    149     //        extracted module key, extracted module, and source module key
    150     //        together. This could be used, for example, to provide a specific
    151     //        memory manager instance to the linking layer.
    152 
    153     auto RequestedSymbols = R.getRequestedSymbols();
    154 
    155     // Extract the requested functions into a new module.
    156     std::unique_ptr<Module> ExtractedFunctionsModule;
    157     if (!RequestedSymbols.empty()) {
    158       std::string Suffix;
    159       std::set<const GlobalValue *> FunctionsToClone;
    160       for (auto &Name : RequestedSymbols) {
    161         auto I = SymbolToDefinition.find(Name);
    162         assert(I != SymbolToDefinition.end() && I->second != nullptr &&
    163                "Should have a non-null definition");
    164         FunctionsToClone.insert(I->second);
    165         Suffix += ".";
    166         Suffix += *Name;
    167       }
    168 
    169       std::lock_guard<std::mutex> Lock(SourceModuleMutex);
    170       ExtractedFunctionsModule =
    171           extractAndClone(*M, Parent.GetAvailableContext(), Suffix,
    172                           [&](const GlobalValue *GV) -> bool {
    173                             return FunctionsToClone.count(GV);
    174                           });
    175     }
    176 
    177     // Build a new ExtractingIRMaterializationUnit to delegate the unrequested
    178     // symbols to.
    179     SymbolFlagsMap DelegatedSymbolFlags;
    180     IRMaterializationUnit::SymbolNameToDefinitionMap
    181         DelegatedSymbolToDefinition;
    182     for (auto &KV : SymbolToDefinition) {
    183       if (RequestedSymbols.count(KV.first))
    184         continue;
    185       DelegatedSymbolFlags[KV.first] =
    186           JITSymbolFlags::fromGlobalValue(*KV.second);
    187       DelegatedSymbolToDefinition[KV.first] = KV.second;
    188     }
    189 
    190     if (!DelegatedSymbolFlags.empty()) {
    191       assert(DelegatedSymbolFlags.size() ==
    192                  DelegatedSymbolToDefinition.size() &&
    193              "SymbolFlags and SymbolToDefinition should have the same number "
    194              "of entries");
    195       R.replace(llvm::make_unique<ExtractingIRMaterializationUnit>(
    196           std::move(M), std::move(DelegatedSymbolFlags),
    197           std::move(DelegatedSymbolToDefinition), Parent));
    198     }
    199 
    200     if (ExtractedFunctionsModule)
    201       Parent.emitExtractedFunctionsModule(std::move(R),
    202                                           std::move(ExtractedFunctionsModule));
    203   }
    204 
    205   void discard(const VSO &V, SymbolStringPtr Name) override {
    206     // All original symbols were materialized by the CODLayer and should be
    207     // final. The function bodies provided by M should never be overridden.
    208     llvm_unreachable("Discard should never be called on an "
    209                      "ExtractingIRMaterializationUnit");
    210   }
    211 
    212   mutable std::mutex SourceModuleMutex;
    213   CompileOnDemandLayer2 &Parent;
    214 };
    215 
    216 CompileOnDemandLayer2::CompileOnDemandLayer2(
    217     ExecutionSession &ES, IRLayer &BaseLayer, JITCompileCallbackManager &CCMgr,
    218     IndirectStubsManagerBuilder BuildIndirectStubsManager,
    219     GetAvailableContextFunction GetAvailableContext)
    220     : IRLayer(ES), BaseLayer(BaseLayer), CCMgr(CCMgr),
    221       BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)),
    222       GetAvailableContext(std::move(GetAvailableContext)) {}
    223 
    224 Error CompileOnDemandLayer2::add(VSO &V, VModuleKey K,
    225                                  std::unique_ptr<Module> M) {
    226   return IRLayer::add(V, K, std::move(M));
    227 }
    228 
    229 void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
    230                                  std::unique_ptr<Module> M) {
    231   auto &ES = getExecutionSession();
    232   assert(M && "M should not be null");
    233 
    234   for (auto &GV : M->global_values())
    235     if (GV.hasWeakLinkage())
    236       GV.setLinkage(GlobalValue::ExternalLinkage);
    237 
    238   MangleAndInterner Mangle(ES, M->getDataLayout());
    239 
    240   extractAliases(R, *M, Mangle);
    241 
    242   auto GlobalsModule = extractGlobals(*M, GetAvailableContext());
    243 
    244   // Delete the bodies of any available externally functions, rename the
    245   // rest, and build the compile callbacks.
    246   std::map<SymbolStringPtr, std::pair<JITTargetAddress, JITSymbolFlags>>
    247       StubCallbacksAndLinkages;
    248   auto &TargetVSO = R.getTargetVSO();
    249 
    250   for (auto &F : M->functions()) {
    251     if (F.isDeclaration())
    252       continue;
    253 
    254     if (F.hasAvailableExternallyLinkage()) {
    255       F.deleteBody();
    256       F.setPersonalityFn(nullptr);
    257       continue;
    258     }
    259 
    260     assert(F.hasName() && "Function should have a name");
    261     std::string StubUnmangledName = F.getName();
    262     F.setName(F.getName() + "$body");
    263     auto StubDecl = cloneFunctionDecl(*M, F);
    264     StubDecl->setName(StubUnmangledName);
    265     StubDecl->setPersonalityFn(nullptr);
    266     StubDecl->setLinkage(GlobalValue::ExternalLinkage);
    267     F.replaceAllUsesWith(StubDecl);
    268 
    269     auto StubName = Mangle(StubUnmangledName);
    270     auto BodyName = Mangle(F.getName());
    271     if (auto CallbackAddr = CCMgr.getCompileCallback(
    272             [BodyName, &TargetVSO, &ES]() -> JITTargetAddress {
    273               if (auto Sym = lookup({&TargetVSO}, BodyName))
    274                 return Sym->getAddress();
    275               else {
    276                 ES.reportError(Sym.takeError());
    277                 return 0;
    278               }
    279             })) {
    280       auto Flags = JITSymbolFlags::fromGlobalValue(F);
    281       Flags &= ~JITSymbolFlags::Weak;
    282       StubCallbacksAndLinkages[std::move(StubName)] =
    283           std::make_pair(*CallbackAddr, Flags);
    284     } else {
    285       ES.reportError(CallbackAddr.takeError());
    286       R.failMaterialization();
    287       return;
    288     }
    289   }
    290 
    291   // Build the stub inits map.
    292   IndirectStubsManager::StubInitsMap StubInits;
    293   for (auto &KV : StubCallbacksAndLinkages)
    294     StubInits[*KV.first] = KV.second;
    295 
    296   // Build the function-body-extracting materialization unit.
    297   if (auto Err = R.getTargetVSO().define(
    298           llvm::make_unique<ExtractingIRMaterializationUnit>(ES, *this,
    299                                                              std::move(M)))) {
    300     ES.reportError(std::move(Err));
    301     R.failMaterialization();
    302     return;
    303   }
    304 
    305   // Build the stubs.
    306   // FIXME: Remove function bodies materialization unit if stub creation fails.
    307   auto &StubsMgr = getStubsManager(TargetVSO);
    308   if (auto Err = StubsMgr.createStubs(StubInits)) {
    309     ES.reportError(std::move(Err));
    310     R.failMaterialization();
    311     return;
    312   }
    313 
    314   // Resolve and finalize stubs.
    315   SymbolMap ResolvedStubs;
    316   for (auto &KV : StubCallbacksAndLinkages) {
    317     if (auto Sym = StubsMgr.findStub(*KV.first, false))
    318       ResolvedStubs[KV.first] = Sym;
    319     else
    320       llvm_unreachable("Stub went missing");
    321   }
    322 
    323   R.resolve(ResolvedStubs);
    324 
    325   BaseLayer.emit(std::move(R), std::move(K), std::move(GlobalsModule));
    326 }
    327 
    328 IndirectStubsManager &CompileOnDemandLayer2::getStubsManager(const VSO &V) {
    329   std::lock_guard<std::mutex> Lock(CODLayerMutex);
    330   StubManagersMap::iterator I = StubsMgrs.find(&V);
    331   if (I == StubsMgrs.end())
    332     I = StubsMgrs.insert(std::make_pair(&V, BuildIndirectStubsManager())).first;
    333   return *I->second;
    334 }
    335 
    336 void CompileOnDemandLayer2::emitExtractedFunctionsModule(
    337     MaterializationResponsibility R, std::unique_ptr<Module> M) {
    338   auto K = getExecutionSession().allocateVModule();
    339   BaseLayer.emit(std::move(R), std::move(K), std::move(M));
    340 }
    341 
    342 } // end namespace orc
    343 } // end namespace llvm
    344