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 303