[ORC] Add (partial) weak symbol support to the CompileOnDemand layer.

This adds partial support for weak functions to the CompileOnDemandLayer by
modifying the addLogicalModule method to check for existing stub definitions
before building a new stub for a weak function. This scheme is sufficient to
support ODR definitions, but fails for general weak definitions if strong
definition is encountered after the first weak definition. (A more extensive
refactor will be required to fully support weak symbols).

This patch does *not* add weak symbol support to RuntimeDyld: I hope to add
that in the near future.



git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@277896 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/ExecutionEngine/JITSymbol.h

@@ -59,6 +59,10 @@ public:
     return (Flags & Common) == Common;
   }
 
+  bool isStrongDefinition() const {
+    return !isWeak() && !isCommon();
+  }
+
   /// @brief Returns true is the Weak flag is set.
   bool isExported() const {
     return (Flags & Exported) == Exported;

include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

@@ -281,12 +281,20 @@ private:
     // Create stub functions.
     const DataLayout &DL = SrcM.getDataLayout();
     {
+      LMResources.StubsMgr = CreateIndirectStubsManager();
+
       typename IndirectStubsMgrT::StubInitsMap StubInits;
       for (auto &F : SrcM) {
         // Skip declarations.
         if (F.isDeclaration())
           continue;
 
+        // Skip weak functions for which we already have definitions.
+        auto MangledName = mangle(F.getName(), DL);
+        if (F.hasWeakLinkage() || F.hasLinkOnceLinkage())
+          if (auto Sym = LD.findSymbol(MangledName, false))
+            continue;
+
         // Record all functions defined by this module.
         if (CloneStubsIntoPartitions)
           LMResources.StubsToClone.insert(&F);
@@ -295,7 +303,7 @@ private:
         // and set the compile action to compile the partition containing the
         // function.
         auto CCInfo = CompileCallbackMgr.getCompileCallback();
-        StubInits[mangle(F.getName(), DL)] =
+        StubInits[MangledName] =
           std::make_pair(CCInfo.getAddress(),
                          JITSymbolFlags::fromGlobalValue(F));
         CCInfo.setCompileAction([this, &LD, LMH, &F]() {
@@ -303,7 +311,6 @@ private:
         });
       }
 
-      LMResources.StubsMgr = CreateIndirectStubsManager();
       auto EC = LMResources.StubsMgr->createStubs(StubInits);
       (void)EC;
       // FIXME: This should be propagated back to the user. Stub creation may
@@ -383,8 +390,7 @@ private:
     // Build a resolver for the globals module and add it to the base layer.
     auto GVsResolver = createLambdaResolver(
         [&LD, LMH](const std::string &Name) {
-          auto &LMResources = LD.getLogicalModuleResources(LMH);
-          if (auto Sym = LMResources.StubsMgr->findStub(Name, false))
+          if (auto Sym = LD.findSymbol(Name, false))
             return Sym;
           auto &LDResolver = LD.getDylibResources().ExternalSymbolResolver;
           return LDResolver->findSymbolInLogicalDylib(Name);

lib/ExecutionEngine/JITSymbol.cpp

@@ -19,7 +19,7 @@ using namespace llvm;
 
 JITSymbolFlags llvm::JITSymbolFlags::fromGlobalValue(const GlobalValue &GV) {
   JITSymbolFlags Flags = JITSymbolFlags::None;
-  if (GV.hasWeakLinkage())
+  if (GV.hasWeakLinkage() || GV.hasLinkOnceLinkage())
     Flags |= JITSymbolFlags::Weak;
   if (GV.hasCommonLinkage())
     Flags |= JITSymbolFlags::Common;

test/ExecutionEngine/OrcLazy/Inputs/weak-function-2.ll

@@ -0,0 +1,13 @@
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.12.0"
+
+; Function Attrs: nounwind ssp uwtable
+define linkonce_odr i32 @baz() #0 {
+entry:
+  ret i32 0
+}
+
+define i8* @bar() {
+entry:
+  ret i8* bitcast (i32 ()* @baz to i8*)
+}

test/ExecutionEngine/OrcLazy/weak-function.ll

@@ -0,0 +1,28 @@
+; RUN: lli -jit-kind=orc-lazy -extra-module %p/Inputs/weak-function-2.ll %s
+;
+; Check that functions in two different modules agree on the address of weak
+; function 'baz'
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.12.0"
+
+define linkonce_odr i32 @baz() {
+entry:
+  ret i32 0
+}
+
+define i8* @foo() {
+entry:
+  ret i8* bitcast (i32 ()* @baz to i8*)
+}
+
+declare i8* @bar()
+
+define i32 @main(i32 %argc, i8** %argv) {
+entry:
+  %call = tail call i8* @foo()
+  %call1 = tail call i8* @bar()
+  %cmp = icmp ne i8* %call, %call1
+  %conv = zext i1 %cmp to i32
+  ret i32 %conv
+}
+

tools/lli/OrcLazyJIT.cpp

@@ -144,8 +144,7 @@ int llvm::runOrcLazyJIT(std::vector<std::unique_ptr<Module>> Ms, int ArgC,
                OrcInlineStubs);
 
   // Add the module, look up main and run it.
-  for (auto &M : Ms)
-    J.addModule(std::move(M));
+  J.addModuleSet(std::move(Ms));
   auto MainSym = J.findSymbol("main");
 
   if (!MainSym) {

tools/lli/OrcLazyJIT.h

@@ -38,7 +38,7 @@ public:
   typedef orc::CompileOnDemandLayer<IRDumpLayerT, CompileCallbackMgr> CODLayerT;
   typedef CODLayerT::IndirectStubsManagerBuilderT
     IndirectStubsManagerBuilder;
-  typedef CODLayerT::ModuleSetHandleT ModuleHandleT;
+  typedef CODLayerT::ModuleSetHandleT ModuleSetHandleT;
 
   OrcLazyJIT(std::unique_ptr<TargetMachine> TM,
              std::unique_ptr<CompileCallbackMgr> CCMgr,
@@ -62,18 +62,21 @@ public:
       DtorRunner.runViaLayer(CODLayer);
   }
 
-  ModuleHandleT addModule(std::unique_ptr<Module> M) {
-    // Attach a data-layout if one isn't already present.
-    if (M->getDataLayout().isDefault())
-      M->setDataLayout(DL);
+  ModuleSetHandleT addModuleSet(std::vector<std::unique_ptr<Module>> Ms) {
+    // Attach a data-layouts if they aren't already present.
+    for (auto &M : Ms)
+      if (M->getDataLayout().isDefault())
+        M->setDataLayout(DL);
 
     // Record the static constructors and destructors. We have to do this before
     // we hand over ownership of the module to the JIT.
     std::vector<std::string> CtorNames, DtorNames;
-    for (auto Ctor : orc::getConstructors(*M))
-      CtorNames.push_back(mangle(Ctor.Func->getName()));
-    for (auto Dtor : orc::getDestructors(*M))
-      DtorNames.push_back(mangle(Dtor.Func->getName()));
+    for (auto &M : Ms) {
+      for (auto Ctor : orc::getConstructors(*M))
+        CtorNames.push_back(mangle(Ctor.Func->getName()));
+      for (auto Dtor : orc::getDestructors(*M))
+        DtorNames.push_back(mangle(Dtor.Func->getName()));
+    }
 
     // Symbol resolution order:
     //   1) Search the JIT symbols.
@@ -84,24 +87,18 @@ public:
         [this](const std::string &Name) -> JITSymbol {
           if (auto Sym = CODLayer.findSymbol(Name, true))
             return Sym;
-          if (auto Sym = CXXRuntimeOverrides.searchOverrides(Name))
-            return Sym;
-
+          return CXXRuntimeOverrides.searchOverrides(Name);
+        },
+        [](const std::string &Name) {
           if (auto Addr =
               RTDyldMemoryManager::getSymbolAddressInProcess(Name))
             return JITSymbol(Addr, JITSymbolFlags::Exported);
-
-          return JITSymbol(nullptr);
-        },
-        [](const std::string &Name) {
           return JITSymbol(nullptr);
         }
       );
 
     // Add the module to the JIT.
-    std::vector<std::unique_ptr<Module>> S;
-    S.push_back(std::move(M));
-    auto H = CODLayer.addModuleSet(std::move(S),
+    auto H = CODLayer.addModuleSet(std::move(Ms),
 				   llvm::make_unique<SectionMemoryManager>(),
 				   std::move(Resolver));
 
@@ -119,7 +116,7 @@ public:
     return CODLayer.findSymbol(mangle(Name), true);
   }
 
-  JITSymbol findSymbolIn(ModuleHandleT H, const std::string &Name) {
+  JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name) {
     return CODLayer.findSymbolIn(H, mangle(Name), true);
   }