[ORC] Add partitioning support to CompileOnDemandLayer2.

CompileOnDemandLayer2 now supports user-supplied partition functions (the
original CompileOnDemandLayer already supported these).

Partition functions are called with the list of requested global values
(i.e. global values that currently have queries waiting on them) and have an
opportunity to select extra global values to materialize at the same time.

Also adds testing infrastructure for the new feature to lli.

llvm-svn: 343396

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
index 6c77bb5..a5de16b 100644 (file)
--- a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
@@ -16,6 +16,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
 
 #include "llvm/ADT/APInt.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
@@ -62,19 +63,37 @@ namespace orc {
 class ExtractingIRMaterializationUnit;
 
 class CompileOnDemandLayer2 : public IRLayer {
-  friend class ExtractingIRMaterializationUnit;
+  friend class PartitioningIRMaterializationUnit;
 
 public:
   /// Builder for IndirectStubsManagers.
   using IndirectStubsManagerBuilder =
       std::function<std::unique_ptr<IndirectStubsManager>()>;
 
+  using GlobalValueSet = std::set<const GlobalValue *>;
+
+  /// Partitioning function.
+  using PartitionFunction =
+      std::function<Optional<GlobalValueSet>(GlobalValueSet Requested)>;
+
+  /// Off-the-shelf partitioning which compiles all requested symbols (usually
+  /// a single function at a time).
+  static Optional<GlobalValueSet> compileRequested(GlobalValueSet Requested);
+
+  /// Off-the-shelf partitioning which compiles whole modules whenever any
+  /// symbol in them is requested.
+  static Optional<GlobalValueSet> compileWholeModule(GlobalValueSet Requested);
+
+  /// Construct a CompileOnDemandLayer2.
   CompileOnDemandLayer2(ExecutionSession &ES, IRLayer &BaseLayer,
                         LazyCallThroughManager &LCTMgr,
                         IndirectStubsManagerBuilder BuildIndirectStubsManager);
 
-  Error add(JITDylib &V, VModuleKey K, ThreadSafeModule TSM) override;
+  /// Sets the partition function.
+  void setPartitionFunction(PartitionFunction Partition);
 
+  /// Emits the given module. This should not be called by clients: it will be
+  /// called by the JIT when a definition added via the add method is requested.
   void emit(MaterializationResponsibility R, VModuleKey K,
             ThreadSafeModule TSM) override;
 
@@ -96,8 +115,12 @@ private:
 
   PerDylibResources &getPerDylibResources(JITDylib &TargetD);
 
-  void emitExtractedFunctionsModule(MaterializationResponsibility R,
-                                    ThreadSafeModule TSM);
+  void cleanUpModule(Module &M);
+
+  void expandPartition(GlobalValueSet &Partition);
+
+  void emitPartition(MaterializationResponsibility R, ThreadSafeModule TSM,
+                     IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
 
   mutable std::mutex CODLayerMutex;
 
@@ -105,6 +128,7 @@ private:
   LazyCallThroughManager &LCTMgr;
   IndirectStubsManagerBuilder BuildIndirectStubsManager;
   PerDylibResourcesMap DylibResources;
+  PartitionFunction Partition = compileRequested;
 };
 
 /// Compile-on-demand layer.

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
index ca66e0c..7e113b8 100644 (file)
--- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
@@ -146,6 +146,12 @@ public:
     TransformLayer.setTransform(std::move(Transform));
   }
 
+  /// Sets the partition function.
+  void
+  setPartitionFunction(CompileOnDemandLayer2::PartitionFunction Partition) {
+    CODLayer.setPartitionFunction(std::move(Partition));
+  }
+
   /// Add a module to be lazily compiled to JITDylib JD.
   Error addLazyIRModule(JITDylib &JD, ThreadSafeModule M);
 

diff --git a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
index 295a9bc..78ee0e7 100644 (file)
--- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
@@ -14,144 +14,75 @@
 using namespace llvm;
 using namespace llvm::orc;
 
-static void extractAliases(MaterializationResponsibility &R, Module &M,
-                           MangleAndInterner &Mangle) {
-  SymbolAliasMap Aliases;
+static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM,
+                                         StringRef Suffix,
+                                         GVPredicate ShouldExtract) {
 
-  std::vector<GlobalAlias *> ModAliases;
-  for (auto &A : M.aliases())
-    ModAliases.push_back(&A);
-
-  for (auto *A : ModAliases) {
-    Constant *Aliasee = A->getAliasee();
-    assert(A->hasName() && "Anonymous alias?");
-    assert(Aliasee->hasName() && "Anonymous aliasee");
-    std::string AliasName = A->getName();
-
-    Aliases[Mangle(AliasName)] = SymbolAliasMapEntry(
-        {Mangle(Aliasee->getName()), JITSymbolFlags::fromGlobalValue(*A)});
-
-    if (isa<Function>(Aliasee)) {
-      auto *F = cloneFunctionDecl(M, *cast<Function>(Aliasee));
-      A->replaceAllUsesWith(F);
-      A->eraseFromParent();
-      F->setName(AliasName);
-    } else if (isa<GlobalVariable>(Aliasee)) {
-      auto *G = cloneGlobalVariableDecl(M, *cast<GlobalVariable>(Aliasee));
-      A->replaceAllUsesWith(G);
-      A->eraseFromParent();
-      G->setName(AliasName);
-    }
-  }
-
-  R.replace(symbolAliases(std::move(Aliases)));
-}
-
-static ThreadSafeModule extractAndClone(ThreadSafeModule &TSM, StringRef Suffix,
-                                        GVPredicate ShouldCloneDefinition) {
+  auto DeleteExtractedDefs = [](GlobalValue &GV) {
+    // Bump the linkage: this global will be provided by the external module.
+    GV.setLinkage(GlobalValue::ExternalLinkage);
 
-  auto DeleteClonedDefsAndPromoteDeclLinkages = [](GlobalValue &GV) {
-    // Delete the definition and bump the linkage in the source module.
+    // Delete the definition in the source module.
     if (isa<Function>(GV)) {
       auto &F = cast<Function>(GV);
       F.deleteBody();
       F.setPersonalityFn(nullptr);
     } else if (isa<GlobalVariable>(GV)) {
       cast<GlobalVariable>(GV).setInitializer(nullptr);
+    } else if (isa<GlobalAlias>(GV)) {
+      // We need to turn deleted aliases into function or variable decls based
+      // on the type of their aliasee.
+      auto &A = cast<GlobalAlias>(GV);
+      Constant *Aliasee = A.getAliasee();
+      assert(A.hasName() && "Anonymous alias?");
+      assert(Aliasee->hasName() && "Anonymous aliasee");
+      std::string AliasName = A.getName();
+
+      if (isa<Function>(Aliasee)) {
+        auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee));
+        A.replaceAllUsesWith(F);
+        A.eraseFromParent();
+        F->setName(AliasName);
+      } else if (isa<GlobalVariable>(Aliasee)) {
+        auto *G = cloneGlobalVariableDecl(*A.getParent(),
+                                          *cast<GlobalVariable>(Aliasee));
+        A.replaceAllUsesWith(G);
+        A.eraseFromParent();
+        G->setName(AliasName);
+      } else
+        llvm_unreachable("Alias to unsupported type");
     } else
       llvm_unreachable("Unsupported global type");
-
-    GV.setLinkage(GlobalValue::ExternalLinkage);
   };
 
-  auto NewTSMod = cloneToNewContext(TSM, ShouldCloneDefinition,
-                                    DeleteClonedDefsAndPromoteDeclLinkages);
+  auto NewTSMod = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs);
   auto &M = *NewTSMod.getModule();
   M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
 
   return NewTSMod;
 }
 
-static ThreadSafeModule extractGlobals(ThreadSafeModule &TSM) {
-  return extractAndClone(TSM, ".globals", [](const GlobalValue &GV) {
-    return isa<GlobalVariable>(GV);
-  });
-}
-
 namespace llvm {
 namespace orc {
 
-class ExtractingIRMaterializationUnit : public IRMaterializationUnit {
+class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
 public:
-  ExtractingIRMaterializationUnit(ExecutionSession &ES,
-                                  CompileOnDemandLayer2 &Parent,
-                                  ThreadSafeModule TSM)
+  PartitioningIRMaterializationUnit(ExecutionSession &ES, ThreadSafeModule TSM,
+                                    CompileOnDemandLayer2 &Parent)
       : IRMaterializationUnit(ES, std::move(TSM)), Parent(Parent) {}
 
-  ExtractingIRMaterializationUnit(ThreadSafeModule TSM,
-                                  SymbolFlagsMap SymbolFlags,
-                                  SymbolNameToDefinitionMap SymbolToDefinition,
-                                  CompileOnDemandLayer2 &Parent)
+  PartitioningIRMaterializationUnit(
+      ThreadSafeModule TSM, SymbolFlagsMap SymbolFlags,
+      SymbolNameToDefinitionMap SymbolToDefinition,
+      CompileOnDemandLayer2 &Parent)
       : IRMaterializationUnit(std::move(TSM), std::move(SymbolFlags),
                               std::move(SymbolToDefinition)),
         Parent(Parent) {}
 
 private:
   void materialize(MaterializationResponsibility R) override {
-    // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
-    //        extracted module key, extracted module, and source module key
-    //        together. This could be used, for example, to provide a specific
-    //        memory manager instance to the linking layer.
-
-    auto RequestedSymbols = R.getRequestedSymbols();
-
-    // Extract the requested functions into a new module.
-    ThreadSafeModule ExtractedFunctionsModule;
-    if (!RequestedSymbols.empty()) {
-      std::string Suffix;
-      std::set<const GlobalValue *> FunctionsToClone;
-      for (auto &Name : RequestedSymbols) {
-        auto I = SymbolToDefinition.find(Name);
-        assert(I != SymbolToDefinition.end() && I->second != nullptr &&
-               "Should have a non-null definition");
-        FunctionsToClone.insert(I->second);
-        Suffix += ".";
-        Suffix += *Name;
-      }
-
-      std::lock_guard<std::mutex> Lock(SourceModuleMutex);
-      ExtractedFunctionsModule =
-          extractAndClone(TSM, Suffix, [&](const GlobalValue &GV) -> bool {
-            return FunctionsToClone.count(&GV);
-          });
-    }
-
-    // Build a new ExtractingIRMaterializationUnit to delegate the unrequested
-    // symbols to.
-    SymbolFlagsMap DelegatedSymbolFlags;
-    IRMaterializationUnit::SymbolNameToDefinitionMap
-        DelegatedSymbolToDefinition;
-    for (auto &KV : SymbolToDefinition) {
-      if (RequestedSymbols.count(KV.first))
-        continue;
-      DelegatedSymbolFlags[KV.first] =
-          JITSymbolFlags::fromGlobalValue(*KV.second);
-      DelegatedSymbolToDefinition[KV.first] = KV.second;
-    }
-
-    if (!DelegatedSymbolFlags.empty()) {
-      assert(DelegatedSymbolFlags.size() ==
-                 DelegatedSymbolToDefinition.size() &&
-             "SymbolFlags and SymbolToDefinition should have the same number "
-             "of entries");
-      R.replace(llvm::make_unique<ExtractingIRMaterializationUnit>(
-          std::move(TSM), std::move(DelegatedSymbolFlags),
-          std::move(DelegatedSymbolToDefinition), Parent));
-    }
-
-    if (ExtractedFunctionsModule)
-      Parent.emitExtractedFunctionsModule(std::move(R),
-                                          std::move(ExtractedFunctionsModule));
+    Parent.emitPartition(std::move(R), std::move(TSM),
+                         std::move(SymbolToDefinition));
   }
 
   void discard(const JITDylib &V, SymbolStringPtr Name) override {
@@ -165,74 +96,71 @@ private:
   CompileOnDemandLayer2 &Parent;
 };
 
+Optional<CompileOnDemandLayer2::GlobalValueSet>
+CompileOnDemandLayer2::compileRequested(GlobalValueSet Requested) {
+  return std::move(Requested);
+}
+
+Optional<CompileOnDemandLayer2::GlobalValueSet>
+CompileOnDemandLayer2::compileWholeModule(GlobalValueSet Requested) {
+  return None;
+}
+
 CompileOnDemandLayer2::CompileOnDemandLayer2(
     ExecutionSession &ES, IRLayer &BaseLayer, LazyCallThroughManager &LCTMgr,
     IndirectStubsManagerBuilder BuildIndirectStubsManager)
     : IRLayer(ES), BaseLayer(BaseLayer), LCTMgr(LCTMgr),
       BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {}
 
-Error CompileOnDemandLayer2::add(JITDylib &V, VModuleKey K,
-                                 ThreadSafeModule TSM) {
-  return IRLayer::add(V, K, std::move(TSM));
+void CompileOnDemandLayer2::setPartitionFunction(PartitionFunction Partition) {
+  this->Partition = std::move(Partition);
 }
 
 void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
                                  ThreadSafeModule TSM) {
+  assert(TSM.getModule() && "Null module");
+
   auto &ES = getExecutionSession();
-  assert(TSM && "M should not be null");
   auto &M = *TSM.getModule();
 
-  for (auto &GV : M.global_values())
-    if (GV.hasWeakLinkage())
-      GV.setLinkage(GlobalValue::ExternalLinkage);
-
-  MangleAndInterner Mangle(ES, M.getDataLayout());
-
-  extractAliases(R, *TSM.getModule(), Mangle);
-
-  auto GlobalsModule = extractGlobals(TSM);
-
-  // Delete the bodies of any available externally functions and build the
-  // lazy reexports alias map.
-  std::map<SymbolStringPtr, std::pair<JITTargetAddress, JITSymbolFlags>>
-      StubCallbacksAndLinkages;
-  auto &TargetJD = R.getTargetJITDylib();
-  auto &Resources = getPerDylibResources(TargetJD);
-  auto &ImplD = Resources.getImplDylib();
+  // First, do some cleanup on the module:
+  cleanUpModule(M);
 
-  SymbolAliasMap LazyReexports;
-  for (auto &F : M.functions()) {
-    if (F.isDeclaration())
-      continue;
+  // Now sort the callables and non-callables, build re-exports and lodge the
+  // actual module with the implementation dylib.
+  auto &PDR = getPerDylibResources(R.getTargetJITDylib());
 
-    if (F.hasAvailableExternallyLinkage()) {
-      F.deleteBody();
-      F.setPersonalityFn(nullptr);
+  MangleAndInterner Mangle(ES, M.getDataLayout());
+  SymbolAliasMap NonCallables;
+  SymbolAliasMap Callables;
+  for (auto &GV : M.global_values()) {
+    assert(GV.hasName() && !GV.hasLocalLinkage() &&
+           "GlobalValues must have been promoted before adding to "
+           "CompileOnDemandLayer");
+    if (GV.isDeclaration() || GV.hasAppendingLinkage())
       continue;
-    }
-
-    auto Flags = JITSymbolFlags::fromGlobalValue(F);
-    assert(Flags.isCallable() && "Non-callable definition in functions module");
 
-    auto MangledName = Mangle(F.getName());
-    LazyReexports[MangledName] = SymbolAliasMapEntry(MangledName, Flags);
+    auto Name = Mangle(GV.getName());
+    auto Flags = JITSymbolFlags::fromGlobalValue(GV);
+    if (Flags.isCallable())
+      Callables[Name] = SymbolAliasMapEntry(Name, Flags);
+    else
+      NonCallables[Name] = SymbolAliasMapEntry(Name, Flags);
   }
 
-  // Add the functions module to the implementation dylib using an extracting
-  // materialization unit.
-  if (auto Err =
-          ImplD.define(llvm::make_unique<ExtractingIRMaterializationUnit>(
-              ES, *this, std::move(TSM)))) {
+  // Create a partitioning materialization unit and lodge it with the
+  // implementation dylib.
+  if (auto Err = PDR.getImplDylib().define(
+          llvm::make_unique<PartitioningIRMaterializationUnit>(
+              ES, std::move(TSM), *this))) {
     ES.reportError(std::move(Err));
     R.failMaterialization();
     return;
   }
 
-  // Handle responsibility for function symbols by returning lazy reexports.
-  auto &ISMgr = Resources.getISManager();
-  R.replace(lazyReexports(LCTMgr, ISMgr, ImplD, LazyReexports));
-
-  BaseLayer.emit(std::move(R), std::move(K), std::move(GlobalsModule));
+  R.replace(reexports(PDR.getImplDylib(), std::move(NonCallables)));
+  R.replace(lazyReexports(LCTMgr, PDR.getISManager(), PDR.getImplDylib(),
+                          std::move(Callables)));
 }
 
 CompileOnDemandLayer2::PerDylibResources &
@@ -251,10 +179,100 @@ CompileOnDemandLayer2::getPerDylibResources(JITDylib &TargetD) {
   return I->second;
 }
 
-void CompileOnDemandLayer2::emitExtractedFunctionsModule(
-    MaterializationResponsibility R, ThreadSafeModule TSM) {
-  auto K = getExecutionSession().allocateVModule();
-  BaseLayer.emit(std::move(R), std::move(K), std::move(TSM));
+void CompileOnDemandLayer2::cleanUpModule(Module &M) {
+  for (auto &F : M.functions()) {
+    if (F.isDeclaration())
+      continue;
+
+    if (F.hasAvailableExternallyLinkage()) {
+      F.deleteBody();
+      F.setPersonalityFn(nullptr);
+      continue;
+    }
+  }
+}
+
+void CompileOnDemandLayer2::expandPartition(GlobalValueSet &Partition) {
+  // Expands the partition to ensure the following rules hold:
+  // (1) If any alias is in the partition, its aliasee is also in the partition.
+  // (2) If any aliasee is in the partition, its aliases are also in the
+  //     partiton.
+  // (3) If any global variable is in the partition then all global variables
+  //     are in the partition.
+  assert(!Partition.empty() && "Unexpected empty partition");
+
+  const Module &M = *(*Partition.begin())->getParent();
+  bool ContainsGlobalVariables = false;
+  std::vector<const GlobalValue *> GVsToAdd;
+
+  for (auto *GV : Partition)
+    if (isa<GlobalAlias>(GV))
+      GVsToAdd.push_back(
+          cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));
+    else if (isa<GlobalVariable>(GV))
+      ContainsGlobalVariables = true;
+
+  for (auto &A : M.aliases())
+    if (Partition.count(cast<GlobalValue>(A.getAliasee())))
+      GVsToAdd.push_back(&A);
+
+  if (ContainsGlobalVariables)
+    for (auto &G : M.globals())
+      GVsToAdd.push_back(&G);
+
+  for (auto *GV : GVsToAdd)
+    Partition.insert(GV);
+}
+
+void CompileOnDemandLayer2::emitPartition(
+    MaterializationResponsibility R, ThreadSafeModule TSM,
+    IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {
+
+  // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
+  //        extracted module key, extracted module, and source module key
+  //        together. This could be used, for example, to provide a specific
+  //        memory manager instance to the linking layer.
+
+  auto &ES = getExecutionSession();
+
+  GlobalValueSet RequestedGVs;
+  for (auto &Name : R.getRequestedSymbols()) {
+    assert(Defs.count(Name) && "No definition for symbol");
+    RequestedGVs.insert(Defs[Name]);
+  }
+
+  auto GVsToExtract = Partition(RequestedGVs);
+
+  // Take a 'None' partition to mean the whole module (as opposed to an empty
+  // partition, which means "materialize nothing"). Emit the whole module
+  // unmodified to the base layer.
+  if (GVsToExtract == None) {
+    Defs.clear();
+    BaseLayer.emit(std::move(R), ES.allocateVModule(), std::move(TSM));
+    return;
+  }
+
+  // If the partition is empty, return the whole module to the symbol table.
+  if (GVsToExtract->empty()) {
+    R.replace(llvm::make_unique<PartitioningIRMaterializationUnit>(
+        std::move(TSM), R.getSymbols(), std::move(Defs), *this));
+    return;
+  }
+
+  expandPartition(*GVsToExtract);
+
+  // Extract the requested partiton (plus any necessary aliases) and
+  // put the rest back into the impl dylib.
+  auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
+    return GVsToExtract->count(&GV);
+  };
+
+  auto ExtractedTSM = extractSubModule(TSM, ".submodule", ShouldExtract);
+
+  R.replace(llvm::make_unique<PartitioningIRMaterializationUnit>(
+      ES, std::move(TSM), *this));
+
+  BaseLayer.emit(std::move(R), ES.allocateVModule(), std::move(ExtractedTSM));
 }
 
 } // end namespace orc

diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 7d38c32..193a5d8 100644 (file)
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -372,6 +372,8 @@ SymbolNameSet MaterializationResponsibility::getRequestedSymbols() const {
 }
 
 void MaterializationResponsibility::resolve(const SymbolMap &Symbols) {
+  LLVM_DEBUG(dbgs() << "In " << JD.getName() << " resolving " << Symbols
+                    << "\n");
 #ifndef NDEBUG
   for (auto &KV : Symbols) {
     auto I = SymbolFlags.find(KV.first);
@@ -435,8 +437,8 @@ void MaterializationResponsibility::replace(
     SymbolFlags.erase(KV.first);
 
   LLVM_DEBUG(JD.getExecutionSession().runSessionLocked([&]() {
-    dbgs() << "In " << JD.getName() << " replacing symbols with MU@" << MU.get()
-           << " (" << MU->getName() << ")\n";
+    dbgs() << "In " << JD.getName() << " replacing symbols with " << *MU
+           << "\n";
   }););
 
   JD.replace(std::move(MU));

diff --git a/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp b/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
index cf09174..9525b16 100644 (file)
--- a/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
@@ -29,11 +29,11 @@ ThreadSafeModule cloneToNewContext(ThreadSafeModule &TSM,
   SmallVector<char, 1> ClonedModuleBuffer;
 
   {
-    std::vector<GlobalValue *> ClonedDefsInSrc;
+    std::set<GlobalValue *> ClonedDefsInSrc;
     ValueToValueMapTy VMap;
     auto Tmp = CloneModule(*TSM.getModule(), VMap, [&](const GlobalValue *GV) {
       if (ShouldCloneDef(*GV)) {
-        ClonedDefsInSrc.push_back(const_cast<GlobalValue *>(GV));
+        ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
         return true;
       }
       return false;

diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/JITSymbol.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/JITSymbol.cpp
index b6ef8ad..0553c21 100644 (file)
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/JITSymbol.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/JITSymbol.cpp
@@ -13,6 +13,7 @@
 
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/Object/ObjectFile.h"
 
@@ -26,8 +27,13 @@ JITSymbolFlags llvm::JITSymbolFlags::fromGlobalValue(const GlobalValue &GV) {
     Flags |= JITSymbolFlags::Common;
   if (!GV.hasLocalLinkage() && !GV.hasHiddenVisibility())
     Flags |= JITSymbolFlags::Exported;
+
   if (isa<Function>(GV))
     Flags |= JITSymbolFlags::Callable;
+  else if (isa<GlobalAlias>(GV) &&
+           isa<Function>(cast<GlobalAlias>(GV).getAliasee()))
+    Flags |= JITSymbolFlags::Callable;
+
   return Flags;
 }
 

diff --git a/llvm/test/ExecutionEngine/OrcLazy/basic-whole-module-partitions.ll b/llvm/test/ExecutionEngine/OrcLazy/basic-whole-module-partitions.ll
new file mode 100644 (file)
index 0000000..3470666
--- /dev/null
+++ b/llvm/test/ExecutionEngine/OrcLazy/basic-whole-module-partitions.ll
@@ -0,0 +1,21 @@
+; RUN: lli -jit-kind=orc-lazy -orc-lazy-debug=funcs-to-stdout \
+; RUN:   %s | FileCheck --check-prefix=CHECK-PER-FUNCTION %s
+; RUN: lli -jit-kind=orc-lazy -per-module-lazy -orc-lazy-debug=funcs-to-stdout \
+; RUN:   %s | FileCheck --check-prefix=CHECK-WHOLE-MODULE %s
+;
+; CHECK-PER-FUNCTION-NOT: foo
+; CHECK-WHOLE-MODULE: foo
+;
+; Checks that the whole module is emitted when -per-module-lazy is enabled,
+; even though foo is not called.
+; Also checks that the foo function is not emitted when -per-module-lazy is off.
+
+define void @foo() {
+entry:
+  ret void
+}
+
+define i32 @main(i32 %argc, i8** nocapture readnone %argv) {
+entry:
+  ret i32 0
+}

diff --git a/llvm/test/ExecutionEngine/OrcLazy/module-flags.ll b/llvm/test/ExecutionEngine/OrcLazy/module-flags.ll
index c1240a8..5ed5bf1 100644 (file)
--- a/llvm/test/ExecutionEngine/OrcLazy/module-flags.ll
+++ b/llvm/test/ExecutionEngine/OrcLazy/module-flags.ll
@@ -1,6 +1,6 @@
 ; RUN: lli -jit-kind=orc-lazy -orc-lazy-debug=mods-to-stdout %s | FileCheck %s
 ;
-; CHECK: module-flags.ll.globals
+; CHECK: module-flags.ll.submodule
 ; CHECK-NOT: Module End
 ; CHECK: The Answer is {{.*}}42
 

diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp
index fab7b35..18cb661 100644 (file)
--- a/llvm/tools/lli/lli.cpp
+++ b/llvm/tools/lli/lli.cpp
@@ -108,6 +108,12 @@ namespace {
                     cl::desc("calls the given entry-point on a new thread "
                              "(jit-kind=orc-lazy only)"));
 
+  cl::opt<bool> PerModuleLazy(
+      "per-module-lazy",
+      cl::desc("Performs lazy compilation on whole module boundaries "
+               "rather than individual functions"),
+      cl::init(false));
+
   // The MCJIT supports building for a target address space separate from
   // the JIT compilation process. Use a forked process and a copying
   // memory manager with IPC to execute using this functionality.
@@ -349,6 +355,7 @@ static void reportError(SMDiagnostic Err, const char *ProgName) {
 }
 
 int runOrcLazyJIT(const char *ProgName);
+void disallowOrcOptions();
 
 //===----------------------------------------------------------------------===//
 // main Driver function
@@ -374,19 +381,8 @@ int main(int argc, char **argv, char * const *envp) {
 
   if (UseJITKind == JITKind::OrcLazy)
     return runOrcLazyJIT(argv[0]);
-  else {
-    // Make sure nobody used an orc-lazy specific option accidentally.
-
-    if (LazyJITCompileThreads != 0) {
-      errs() << "-compile-threads requires -jit-kind=orc-lazy\n";
-      exit(1);
-    }
-
-    if (!ThreadEntryPoints.empty()) {
-      errs() << "-thread-entry requires -jit-kind=orc-lazy\n";
-      exit(1);
-    }
-  }
+  else
+    disallowOrcOptions();
 
   LLVMContext Context;
 
@@ -794,6 +790,9 @@ int runOrcLazyJIT(const char *ProgName) {
   }
   auto J = ExitOnErr(orc::LLLazyJIT::Create(std::move(JTMB), DL, LazyJITCompileThreads));
 
+  if (PerModuleLazy)
+    J->setPartitionFunction(orc::CompileOnDemandLayer2::compileWholeModule);
+
   auto Dump = createDebugDumper();
 
   J->setLazyCompileTransform([&](orc::ThreadSafeModule TSM,
@@ -873,6 +872,25 @@ int runOrcLazyJIT(const char *ProgName) {
   return Result;
 }
 
+void disallowOrcOptions() {
+  // Make sure nobody used an orc-lazy specific option accidentally.
+
+  if (LazyJITCompileThreads != 0) {
+    errs() << "-compile-threads requires -jit-kind=orc-lazy\n";
+    exit(1);
+  }
+
+  if (!ThreadEntryPoints.empty()) {
+    errs() << "-thread-entry requires -jit-kind=orc-lazy\n";
+    exit(1);
+  }
+
+  if (PerModuleLazy) {
+    errs() << "-per-module-lazy requires -jit-kind=orc-lazy\n";
+    exit(1);
+  }
+}
+
 std::unique_ptr<FDRawChannel> launchRemote() {
 #ifndef LLVM_ON_UNIX
   llvm_unreachable("launchRemote not supported on non-Unix platforms");