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