llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.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-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
  10 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
  11 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.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 = 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       } else
  52         llvm_unreachable("Alias to unsupported type");
  53     } else
  54       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
  65 namespace llvm {
  66 namespace orc {
  67
  68 class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
  69 public:
  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
  83 private:
  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 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   IRPartitionLayer &Parent;
  97 };
  98
  99 } // namespace orc
 100 } // namespace llvm
 101
 102 IRPartitionLayer::IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer)
 103     : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer) {}
 104
 105 void IRPartitionLayer::setPartitionFunction(PartitionFunction Partition) {
 106   this->Partition = Partition;
 107 }
 108
 109 std::optional<IRPartitionLayer::GlobalValueSet>
 110 IRPartitionLayer::compileRequested(GlobalValueSet Requested) {
 111   return std::move(Requested);
 112 }
 113
 114 std::optional<IRPartitionLayer::GlobalValueSet>
 115 IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) {
 116   return std::nullopt;
 117 }
 118
 119 void 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
 138 void 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
 151 void 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 the
 155   //     partiton.
 156   // (3) If any global variable is in the partition then all global variables
 157   //     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
 183 void 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 the
 188   //        extracted module key, extracted module, and source module key
 189   //        together. This could be used, for example, to provide a specific
 190   //        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 partition
 207   /// 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 empty
 212   // partition, which means "materialize nothing"). Emit the whole module
 213   // 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 symbol
 236   // linkages/names, expand the partition to include any required symbols
 237   // (i.e. symbols that can't be separated from our partition), and
 238   // then extract the partition.
 239   //
 240   // FIXME: We apply this promotion once per partitioning. It's safe, but
 241   // 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       for (const auto *GV : *GVsToExtract)
 264         HashGVs.push_back(GV);
 265       llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) {
 266         return LHS->getName() < RHS->getName();
 267       });
 268       hash_code HC(0);
 269       for (const auto *GV : HashGVs) {
 270         assert(GV->hasName() && "All GVs to extract should be named by now");
 271         auto GVName = GV->getName();
 272         HC = hash_combine(HC, hash_combine_range(GVName.begin(), GVName.end()));
 273       }
 274       raw_string_ostream(SubModuleName)
 275           << ".submodule."
 276           << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}",
 277                      static_cast<size_t>(HC))
 278           << ".ll";
 279     }
 280
 281     // Extract the requested partiton (plus any necessary aliases) and
 282     // put the rest back into the impl dylib.
 283     auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
 284       return GVsToExtract->count(&GV);
 285     };
 286
 287     return extractSubModule(TSM, SubModuleName, ShouldExtract);
 288   });
 289
 290   if (!ExtractedTSM) {
 291     ES.reportError(ExtractedTSM.takeError());
 292     R->failMaterialization();
 293     return;
 294   }
 295
 296   if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
 297           ES, *getManglingOptions(), std::move(TSM), *this))) {
 298     ES.reportError(std::move(Err));
 299     R->failMaterialization();
 300     return;
 301   }
 302   BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));
 303 }