303 lines · cpp
1//===----- IRPartitionLayer.cpp - Partition IR module into submodules -----===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"10#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"11#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"12 13using namespace llvm;14using namespace llvm::orc;15 16static 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 based33 // 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 = std::string(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 } else52 llvm_unreachable("Alias to unsupported type");53 } else54 llvm_unreachable("Unsupported global type");55 };56 57 auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs);58 NewTSM.withModuleDo([&](Module &M) {59 M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());60 });61 62 return NewTSM;63}64 65namespace llvm {66namespace orc {67 68class PartitioningIRMaterializationUnit : public IRMaterializationUnit {69public:70 PartitioningIRMaterializationUnit(ExecutionSession &ES,71 const IRSymbolMapper::ManglingOptions &MO,72 ThreadSafeModule TSM,73 IRPartitionLayer &Parent)74 : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {}75 76 PartitioningIRMaterializationUnit(77 ThreadSafeModule TSM, Interface I,78 SymbolNameToDefinitionMap SymbolToDefinition, IRPartitionLayer &Parent)79 : IRMaterializationUnit(std::move(TSM), std::move(I),80 std::move(SymbolToDefinition)),81 Parent(Parent) {}82 83private:84 void materialize(std::unique_ptr<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 be91 // 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 IRPartitionLayer &Parent;97};98 99} // namespace orc100} // namespace llvm101 102IRPartitionLayer::IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer)103 : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer) {}104 105void IRPartitionLayer::setPartitionFunction(PartitionFunction Partition) {106 this->Partition = Partition;107}108 109std::optional<IRPartitionLayer::GlobalValueSet>110IRPartitionLayer::compileRequested(GlobalValueSet Requested) {111 return std::move(Requested);112}113 114std::optional<IRPartitionLayer::GlobalValueSet>115IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) {116 return std::nullopt;117}118 119void IRPartitionLayer::emit(std::unique_ptr<MaterializationResponsibility> R,120 ThreadSafeModule TSM) {121 assert(TSM && "Null module");122 123 auto &ES = getExecutionSession();124 TSM.withModuleDo([&](Module &M) {125 // First, do some cleanup on the module:126 cleanUpModule(M);127 });128 129 // Create a partitioning materialization unit and pass the responsibility.130 if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(131 ES, *getManglingOptions(), std::move(TSM), *this))) {132 ES.reportError(std::move(Err));133 R->failMaterialization();134 return;135 }136}137 138void IRPartitionLayer::cleanUpModule(Module &M) {139 for (auto &F : M.functions()) {140 if (F.isDeclaration())141 continue;142 143 if (F.hasAvailableExternallyLinkage()) {144 F.deleteBody();145 F.setPersonalityFn(nullptr);146 continue;147 }148 }149}150 151void IRPartitionLayer::expandPartition(GlobalValueSet &Partition) {152 // Expands the partition to ensure the following rules hold:153 // (1) If any alias is in the partition, its aliasee is also in the partition.154 // (2) If any aliasee is in the partition, its aliases are also in the155 // partiton.156 // (3) If any global variable is in the partition then all global variables157 // are in the partition.158 assert(!Partition.empty() && "Unexpected empty partition");159 160 const Module &M = *(*Partition.begin())->getParent();161 bool ContainsGlobalVariables = false;162 std::vector<const GlobalValue *> GVsToAdd;163 164 for (const auto *GV : Partition)165 if (isa<GlobalAlias>(GV))166 GVsToAdd.push_back(167 cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));168 else if (isa<GlobalVariable>(GV))169 ContainsGlobalVariables = true;170 171 for (auto &A : M.aliases())172 if (Partition.count(cast<GlobalValue>(A.getAliasee())))173 GVsToAdd.push_back(&A);174 175 if (ContainsGlobalVariables)176 for (auto &G : M.globals())177 GVsToAdd.push_back(&G);178 179 for (const auto *GV : GVsToAdd)180 Partition.insert(GV);181}182 183void IRPartitionLayer::emitPartition(184 std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM,185 IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {186 187 // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the188 // extracted module key, extracted module, and source module key189 // together. This could be used, for example, to provide a specific190 // memory manager instance to the linking layer.191 192 auto &ES = getExecutionSession();193 GlobalValueSet RequestedGVs;194 for (auto &Name : R->getRequestedSymbols()) {195 if (Name == R->getInitializerSymbol())196 TSM.withModuleDo([&](Module &M) {197 for (auto &GV : getStaticInitGVs(M))198 RequestedGVs.insert(&GV);199 });200 else {201 assert(Defs.count(Name) && "No definition for symbol");202 RequestedGVs.insert(Defs[Name]);203 }204 }205 206 /// Perform partitioning with the context lock held, since the partition207 /// function is allowed to access the globals to compute the partition.208 auto GVsToExtract =209 TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); });210 211 // Take a 'None' partition to mean the whole module (as opposed to an empty212 // partition, which means "materialize nothing"). Emit the whole module213 // unmodified to the base layer.214 if (GVsToExtract == std::nullopt) {215 Defs.clear();216 BaseLayer.emit(std::move(R), std::move(TSM));217 return;218 }219 220 // If the partition is empty, return the whole module to the symbol table.221 if (GVsToExtract->empty()) {222 if (auto Err =223 R->replace(std::make_unique<PartitioningIRMaterializationUnit>(224 std::move(TSM),225 MaterializationUnit::Interface(R->getSymbols(),226 R->getInitializerSymbol()),227 std::move(Defs), *this))) {228 getExecutionSession().reportError(std::move(Err));229 R->failMaterialization();230 return;231 }232 return;233 }234 235 // Ok -- we actually need to partition the symbols. Promote the symbol236 // linkages/names, expand the partition to include any required symbols237 // (i.e. symbols that can't be separated from our partition), and238 // then extract the partition.239 //240 // FIXME: We apply this promotion once per partitioning. It's safe, but241 // overkill.242 auto ExtractedTSM = TSM.withModuleDo([&](Module &M)243 -> Expected<ThreadSafeModule> {244 auto PromotedGlobals = PromoteSymbols(M);245 if (!PromotedGlobals.empty()) {246 247 MangleAndInterner Mangle(ES, M.getDataLayout());248 SymbolFlagsMap SymbolFlags;249 IRSymbolMapper::add(ES, *getManglingOptions(), PromotedGlobals,250 SymbolFlags);251 252 if (auto Err = R->defineMaterializing(SymbolFlags))253 return std::move(Err);254 }255 256 expandPartition(*GVsToExtract);257 258 // Submodule name is given by hashing the names of the globals.259 std::string SubModuleName;260 {261 std::vector<const GlobalValue *> HashGVs;262 HashGVs.reserve(GVsToExtract->size());263 llvm::append_range(HashGVs, *GVsToExtract);264 llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) {265 return LHS->getName() < RHS->getName();266 });267 hash_code HC(0);268 for (const auto *GV : HashGVs) {269 assert(GV->hasName() && "All GVs to extract should be named by now");270 auto GVName = GV->getName();271 HC = hash_combine(HC, hash_combine_range(GVName));272 }273 raw_string_ostream(SubModuleName)274 << ".submodule."275 << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}",276 static_cast<size_t>(HC))277 << ".ll";278 }279 280 // Extract the requested partiton (plus any necessary aliases) and281 // put the rest back into the impl dylib.282 auto ShouldExtract = [&](const GlobalValue &GV) -> bool {283 return GVsToExtract->count(&GV);284 };285 286 return extractSubModule(TSM, SubModuleName, ShouldExtract);287 });288 289 if (!ExtractedTSM) {290 ES.reportError(ExtractedTSM.takeError());291 R->failMaterialization();292 return;293 }294 295 if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(296 ES, *getManglingOptions(), std::move(TSM), *this))) {297 ES.reportError(std::move(Err));298 R->failMaterialization();299 return;300 }301 BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));302}303